<EOL>with tf.compat.v1.name_scope(name, values=[x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>rank = util.prefer_static_rank(x)<EOL>if axis < <NUM_LIT:0>:<EOL><INDENT>axis = rank + axis<EOL><DEDENT>shift = rank - <NUM_LIT:1> - axis<EOL>x_rotated = util.rotate_transpose(x, shift)<EOL>if center:<EOL><INDENT>x_rotated -= tf.reduce_mean(<EOL>input_tensor=x_rotated, axis=-<NUM_LIT:1>, keepdims=True)<EOL><DEDENT>x_len = util.prefer_static_shape(x_rotated)[-<NUM_LIT:1>]<EOL>x_len_float64 = tf.cast(x_len, np.float64)<EOL>target_length = tf.pow(<EOL>np.float64(<NUM_LIT>), tf.math.ceil(<EOL>tf.math.log(x_len_float64 * <NUM_LIT:2>) / np.log(<NUM_LIT>)))<EOL>pad_length = tf.cast(target_length - x_len_float64, np.int32)<EOL>x_rotated_pad = util.pad(x_rotated, axis=-<NUM_LIT:1>, back=True, count=pad_length)<EOL>dtype = x.dtype<EOL>if not dtype.is_complex:<EOL><INDENT>if not dtype.is_floating:<EOL><INDENT>raise TypeError('<STR_LIT>'<EOL>'<STR_LIT>'.format(dtype))<EOL><DEDENT>x_rotated_pad = tf.complex(x_rotated_pad,<EOL>dtype.real_dtype.as_numpy_dtype(<NUM_LIT:0.>))<EOL><DEDENT>fft_x_rotated_pad = tf.signal.fft(x_rotated_pad)<EOL>spectral_density = fft_x_rotated_pad * tf.math.conj(fft_x_rotated_pad)<EOL>shifted_product = tf.signal.ifft(spectral_density)<EOL>shifted_product = tf.cast(shifted_product, dtype)<EOL>know_static_shape = True<EOL>if not x_rotated.shape.is_fully_defined():<EOL><INDENT>know_static_shape = False<EOL><DEDENT>if max_lags is None:<EOL><INDENT>max_lags = x_len - <NUM_LIT:1><EOL><DEDENT>else:<EOL><INDENT>max_lags = tf.convert_to_tensor(value=max_lags, name='<STR_LIT>')<EOL>max_lags_ = tf.get_static_value(max_lags)<EOL>if max_lags_ is None or not know_static_shape:<EOL><INDENT>know_static_shape = False<EOL>max_lags = tf.minimum(x_len - <NUM_LIT:1>, max_lags)<EOL><DEDENT>else:<EOL><INDENT>max_lags = min(x_len - <NUM_LIT:1>, max_lags_)<EOL><DEDENT><DEDENT>shifted_product_chopped = shifted_product[..., :max_lags + <NUM_LIT:1>]<EOL>if know_static_shape:<EOL><INDENT>chopped_shape = x_rotated.shape.as_list()<EOL>chopped_shape[-<NUM_LIT:1>] = min(x_len, max_lags + <NUM_LIT:1>)<EOL>shifted_product_chopped.set_shape(chopped_shape)<EOL><DEDENT>x_len = tf.cast(x_len, dtype.real_dtype)<EOL>max_lags = tf.cast(max_lags, dtype.real_dtype)<EOL>denominator = x_len - tf.range(<NUM_LIT:0.>, max_lags + <NUM_LIT:1.>)<EOL>denominator = tf.cast(denominator, dtype)<EOL>shifted_product_rotated = shifted_product_chopped / denominator<EOL>if normalize:<EOL><INDENT>shifted_product_rotated /= shifted_product_rotated[..., :<NUM_LIT:1>]<EOL><DEDENT>return util.rotate_transpose(shifted_product_rotated, -shift)<EOL><DEDENT>

Auto correlation along one axis.

    Given a `1-D` wide sense stationary (WSS) sequence `X`, the auto correlation
    `RXX` may be defined as  (with `E` expectation and `Conj` complex conjugate)

    ```
    RXX[m] := E{ W[m] Conj(W[0]) } = E{ W[0] Conj(W[-m]) },
    W[n]   := (X[n] - MU) / S,
    MU     := E{ X[0] },
    S**2   := E{ (X[0] - MU) Conj(X[0] - MU) }.
    ```

    This function takes the viewpoint that `x` is (along one axis) a finite
    sub-sequence of a realization of (WSS) `X`, and then uses `x` to produce an
    estimate of `RXX[m]` as follows:

    After extending `x` from length `L` to `inf` by zero padding, the auto
    correlation estimate `rxx[m]` is computed for `m = 0, 1, ..., max_lags` as

    ```
    rxx[m] := (L - m)**-1 sum_n w[n + m] Conj(w[n]),
    w[n]   := (x[n] - mu) / s,
    mu     := L**-1 sum_n x[n],
    s**2   := L**-1 sum_n (x[n] - mu) Conj(x[n] - mu)
    ```

    The error in this estimate is proportional to `1 / sqrt(len(x) - m)`, so users
    often set `max_lags` small enough so that the entire output is meaningful.

    Note that since `mu` is an imperfect estimate of `E{ X[0] }`, and we divide by
    `len(x) - m` rather than `len(x) - m - 1`, our estimate of auto correlation
    contains a slight bias, which goes to zero as `len(x) - m --> infinity`.

    Args:
      x:  `float32` or `complex64` `Tensor`.
      axis:  Python `int`. The axis number along which to compute correlation.
        Other dimensions index different batch members.
      max_lags:  Positive `int` tensor.  The maximum value of `m` to consider (in
        equation above).  If `max_lags >= x.shape[axis]`, we effectively re-set
        `max_lags` to `x.shape[axis] - 1`.
      center:  Python `bool`.  If `False`, do not subtract the mean estimate `mu`
        from `x[n]` when forming `w[n]`.
      normalize:  Python `bool`.  If `False`, do not divide by the variance
        estimate `s**2` when forming `w[n]`.
      name:  `String` name to prepend to created ops.

    Returns:
      `rxx`: `Tensor` of same `dtype` as `x`.  `rxx.shape[i] = x.shape[i]` for
        `i != axis`, and `rxx.shape[axis] = max_lags + 1`.

    Raises:
      TypeError:  If `x` is not a supported type.


def cholesky_covariance(x, sample_axis=<NUM_LIT:0>, keepdims=False, name=None):

with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[x, sample_axis]):<EOL><INDENT>sample_axis = tf.convert_to_tensor(value=sample_axis, dtype=tf.int32)<EOL>cov = covariance(<EOL>x, sample_axis=sample_axis, event_axis=-<NUM_LIT:1>, keepdims=keepdims)<EOL>return tf.linalg.cholesky(cov)<EOL><DEDENT>

Cholesky factor of the covariance matrix of vector-variate random samples.

    This function can be use to fit a multivariate normal to data.

    ```python
    tf.enable_eager_execution()
    import tensorflow_probability as tfp
    tfd = tfp.distributions

    # Assume data.shape = (1000, 2).  1000 samples of a random variable in R^2.
    observed_data = read_data_samples(...)

    # The mean is easy
    mu = tf.reduce_mean(observed_data, axis=0)

    # Get the scale matrix
    L = tfp.stats.cholesky_covariance(observed_data)

    # Make the best fit multivariate normal (under maximum likelihood condition).
    mvn = tfd.MultivariateNormalTriL(loc=mu, scale_tril=L)

    # Plot contours of the pdf.
    xs, ys = tf.meshgrid(
        tf.linspace(-5., 5., 50), tf.linspace(-5., 5., 50), indexing='ij')
    xy = tf.stack((tf.reshape(xs, [-1]), tf.reshape(ys, [-1])), axis=-1)
    pdf = tf.reshape(mvn.prob(xy), (50, 50))
    CS = plt.contour(xs, ys, pdf, 10)
    plt.clabel(CS, inline=1, fontsize=10)
    ```

    Why does this work?
    Given vector-variate random variables `X = (X1, ..., Xd)`, one may obtain the
    sample covariance matrix in `R^{d x d}` (see `tfp.stats.covariance`).

    The [Cholesky factor](https://en.wikipedia.org/wiki/Cholesky_decomposition)
    of this matrix is analogous to standard deviation for scalar random variables:
    Suppose `X` has covariance matrix `C`, with Cholesky factorization `C = L L^T`
    Then multiplying a vector of iid random variables which have unit variance by
    `L` produces a vector with covariance `L L^T`, which is the same as `X`.

    ```python
    observed_data = read_data_samples(...)
    L = tfp.stats.cholesky_covariance(observed_data, sample_axis=0)

    # Make fake_data with the same covariance as observed_data.
    uncorrelated_normal = tf.random_normal(shape=(500, 10))
    fake_data = tf.linalg.matvec(L, uncorrelated_normal)
    ```

    Args:
      x:  Numeric `Tensor`.  The rightmost dimension of `x` indexes events. E.g.
        dimensions of a random vector.
      sample_axis: Scalar or vector `Tensor` designating axis holding samples.
        Default value: `0` (leftmost dimension). Cannot be the rightmost dimension
          (since this indexes events).
      keepdims:  Boolean.  Whether to keep the sample axis as singletons.
      name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., `'covariance'`).

    Returns:
      chol:  `Tensor` of same `dtype` as `x`.  The last two dimensions hold
        lower triangular matrices (the Cholesky factors).


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[x, y, event_axis, sample_axis]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>x -= tf.reduce_mean(input_tensor=x, axis=sample_axis, keepdims=True)<EOL>if y is None:<EOL><INDENT>y = x<EOL><DEDENT>else:<EOL><INDENT>y = tf.convert_to_tensor(value=y, name='<STR_LIT:y>', dtype=x.dtype)<EOL>x.shape.assert_is_compatible_with(y.shape)<EOL>y -= tf.reduce_mean(input_tensor=y, axis=sample_axis, keepdims=True)<EOL><DEDENT>if event_axis is None:<EOL><INDENT>return tf.reduce_mean(<EOL>input_tensor=x * tf.math.conj(y), axis=sample_axis, keepdims=keepdims)<EOL><DEDENT>if sample_axis is None:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(event_axis))<EOL><DEDENT>event_axis = _make_positive_axis(event_axis, tf.rank(x))<EOL>sample_axis = _make_positive_axis(sample_axis, tf.rank(x))<EOL>if _is_list_like(event_axis) and _is_list_like(sample_axis):<EOL><INDENT>if set(event_axis).intersection(sample_axis):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(<EOL>sample_axis, event_axis))<EOL><DEDENT>if (np.diff(sorted(event_axis)) > <NUM_LIT:1>).any():<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(event_axis))<EOL><DEDENT>batch_axis = list(<EOL>sorted(<EOL>set(range(x.shape.ndims)).difference(sample_axis + event_axis)))<EOL><DEDENT>else:<EOL><INDENT>batch_axis, _ = tf.compat.v1.setdiff1d(<EOL>tf.range(<NUM_LIT:0>, tf.rank(x)), tf.concat((sample_axis, event_axis), <NUM_LIT:0>))<EOL><DEDENT>event_axis = tf.convert_to_tensor(<EOL>value=event_axis, name='<STR_LIT>', dtype=tf.int32)<EOL>sample_axis = tf.convert_to_tensor(<EOL>value=sample_axis, name='<STR_LIT>', dtype=tf.int32)<EOL>batch_axis = tf.convert_to_tensor(<EOL>value=batch_axis, name='<STR_LIT>', dtype=tf.int32)<EOL>perm_for_xy = tf.concat((batch_axis, event_axis, sample_axis), <NUM_LIT:0>)<EOL>x_permed = tf.transpose(a=x, perm=perm_for_xy)<EOL>y_permed = tf.transpose(a=y, perm=perm_for_xy)<EOL>batch_ndims = tf.size(input=batch_axis)<EOL>batch_shape = tf.shape(input=x_permed)[:batch_ndims]<EOL>event_ndims = tf.size(input=event_axis)<EOL>event_shape = tf.shape(input=x_permed)[batch_ndims:batch_ndims +<EOL>event_ndims]<EOL>sample_shape = tf.shape(input=x_permed)[batch_ndims + event_ndims:]<EOL>sample_ndims = tf.size(input=sample_shape)<EOL>n_samples = tf.reduce_prod(input_tensor=sample_shape)<EOL>n_events = tf.reduce_prod(input_tensor=event_shape)<EOL>x_permed_flat = tf.reshape(<EOL>x_permed, tf.concat((batch_shape, event_shape, [n_samples]), <NUM_LIT:0>))<EOL>y_permed_flat = tf.reshape(<EOL>y_permed, tf.concat((batch_shape, event_shape, [n_samples]), <NUM_LIT:0>))<EOL>x_permed_flat = tf.reshape(<EOL>x_permed, tf.concat((batch_shape, [n_events], [n_samples]), <NUM_LIT:0>))<EOL>y_permed_flat = tf.reshape(<EOL>y_permed, tf.concat((batch_shape, [n_events], [n_samples]), <NUM_LIT:0>))<EOL>cov = tf.matmul(<EOL>x_permed_flat, y_permed_flat, adjoint_b=True) / tf.cast(<EOL>n_samples, x.dtype)<EOL>cov = tf.reshape(<EOL>cov,<EOL>tf.concat(<EOL>(<EOL>batch_shape,<EOL>event_shape**<NUM_LIT:2>,<EOL>tf.ones([sample_ndims], tf.int32)),<EOL><NUM_LIT:0>))<EOL>cov = tf.transpose(a=cov, perm=tf.math.invert_permutation(perm_for_xy))<EOL>e_start = event_axis[<NUM_LIT:0>]<EOL>e_len = <NUM_LIT:1> + event_axis[-<NUM_LIT:1>] - event_axis[<NUM_LIT:0>]<EOL>cov = tf.reshape(<EOL>cov,<EOL>tf.concat((tf.shape(input=cov)[:e_start], event_shape, event_shape,<EOL>tf.shape(input=cov)[e_start + e_len:]), <NUM_LIT:0>))<EOL>if not keepdims:<EOL><INDENT>squeeze_axis = tf.where(sample_axis < e_start, sample_axis,<EOL>sample_axis + e_len)<EOL>cov = _squeeze(cov, axis=squeeze_axis)<EOL><DEDENT>return cov<EOL><DEDENT>

Sample covariance between observations indexed by `event_axis`.

    Given `N` samples of scalar random variables `X` and `Y`, covariance may be
    estimated as

    ```none
    Cov[X, Y] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(Y_n - Ybar)}
    Xbar := N^{-1} sum_{n=1}^N X_n
    Ybar := N^{-1} sum_{n=1}^N Y_n
    ```

    For vector-variate random variables `X = (X1, ..., Xd)`, `Y = (Y1, ..., Yd)`,
    one is often interested in the covariance matrix, `C_{ij} := Cov[Xi, Yj]`.

    ```python
    x = tf.random_normal(shape=(100, 2, 3))
    y = tf.random_normal(shape=(100, 2, 3))

    # cov[i, j] is the sample covariance between x[:, i, j] and y[:, i, j].
    cov = tfp.stats.covariance(x, y, sample_axis=0, event_axis=None)

    # cov_matrix[i, m, n] is the sample covariance of x[:, i, m] and y[:, i, n]
    cov_matrix = tfp.stats.covariance(x, y, sample_axis=0, event_axis=-1)
    ```

    Notice we divide by `N` (the numpy default), which does not create `NaN`
    when `N = 1`, but is slightly biased.

    Args:
      x:  A numeric `Tensor` holding samples.
      y:  Optional `Tensor` with same `dtype` and `shape` as `x`.
        Default value: `None` (`y` is effectively set to `x`).
      sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
        `None` (meaning all axis hold samples).
        Default value: `0` (leftmost dimension).
      event_axis:  Scalar or vector `Tensor`, or `None` (scalar events).
        Axis indexing random events, whose covariance we are interested in.
        If a vector, entries must form a contiguous block of dims. `sample_axis`
        and `event_axis` should not intersect.
        Default value: `-1` (rightmost axis holds events).
      keepdims:  Boolean.  Whether to keep the sample axis as singletons.
      name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., `'covariance'`).

    Returns:
      cov: A `Tensor` of same `dtype` as the `x`, and rank equal to
        `rank(x) - len(sample_axis) + 2 * len(event_axis)`.

    Raises:
      AssertionError:  If `x` and `y` are found to have different shape.
      ValueError:  If `sample_axis` and `event_axis` are found to overlap.
      ValueError:  If `event_axis` is found to not be contiguous.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[x, y, event_axis, sample_axis]):<EOL><INDENT>x /= stddev(x, sample_axis=sample_axis, keepdims=True)<EOL>if y is not None:<EOL><INDENT>y /= stddev(y, sample_axis=sample_axis, keepdims=True)<EOL><DEDENT>return covariance(<EOL>x=x,<EOL>y=y,<EOL>event_axis=event_axis,<EOL>sample_axis=sample_axis,<EOL>keepdims=keepdims)<EOL><DEDENT>

Sample correlation (Pearson) between observations indexed by `event_axis`.

    Given `N` samples of scalar random variables `X` and `Y`, correlation may be
    estimated as

    ```none
    Corr[X, Y] := Cov[X, Y] / Sqrt(Cov[X, X] * Cov[Y, Y]),
    where
    Cov[X, Y] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(Y_n - Ybar)}
    Xbar := N^{-1} sum_{n=1}^N X_n
    Ybar := N^{-1} sum_{n=1}^N Y_n
    ```

    Correlation is always in the interval `[-1, 1]`, and `Corr[X, X] == 1`.

    For vector-variate random variables `X = (X1, ..., Xd)`, `Y = (Y1, ..., Yd)`,
    one is often interested in the correlation matrix, `C_{ij} := Corr[Xi, Yj]`.

    ```python
    x = tf.random_normal(shape=(100, 2, 3))
    y = tf.random_normal(shape=(100, 2, 3))

    # corr[i, j] is the sample correlation between x[:, i, j] and y[:, i, j].
    corr = tfp.stats.correlation(x, y, sample_axis=0, event_axis=None)

    # corr_matrix[i, m, n] is the sample correlation of x[:, i, m] and y[:, i, n]
    corr_matrix = tfp.stats.correlation(x, y, sample_axis=0, event_axis=-1)
    ```

    Notice we divide by `N` (the numpy default), which does not create `NaN`
    when `N = 1`, but is slightly biased.

    Args:
      x:  A numeric `Tensor` holding samples.
      y:  Optional `Tensor` with same `dtype` and `shape` as `x`.
        Default value: `None` (`y` is effectively set to `x`).
      sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
        `None` (meaning all axis hold samples).
        Default value: `0` (leftmost dimension).
      event_axis:  Scalar or vector `Tensor`, or `None` (scalar events).
        Axis indexing random events, whose correlation we are interested in.
        If a vector, entries must form a contiguous block of dims. `sample_axis`
        and `event_axis` should not intersect.
        Default value: `-1` (rightmost axis holds events).
      keepdims:  Boolean.  Whether to keep the sample axis as singletons.
      name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., `'correlation'`).

    Returns:
      corr: A `Tensor` of same `dtype` as the `x`, and rank equal to
        `rank(x) - len(sample_axis) + 2 * len(event_axis)`.

    Raises:
      AssertionError:  If `x` and `y` are found to have different shape.
      ValueError:  If `sample_axis` and `event_axis` are found to overlap.
      ValueError:  If `event_axis` is found to not be contiguous.


def stddev(x, sample_axis=<NUM_LIT:0>, keepdims=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', values=[x, sample_axis]):<EOL><INDENT>return tf.sqrt(variance(x, sample_axis=sample_axis, keepdims=keepdims))<EOL><DEDENT>

Estimate standard deviation using samples.

    Given `N` samples of scalar valued random variable `X`, standard deviation may
    be estimated as

    ```none
    Stddev[X] := Sqrt[Var[X]],
    Var[X] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(X_n - Xbar)},
    Xbar := N^{-1} sum_{n=1}^N X_n
    ```

    ```python
    x = tf.random_normal(shape=(100, 2, 3))

    # stddev[i, j] is the sample standard deviation of the (i, j) batch member.
    stddev = tfp.stats.stddev(x, sample_axis=0)
    ```

    Scaling a unit normal by a standard deviation produces normal samples
    with that standard deviation.

    ```python
    observed_data = read_data_samples(...)
    stddev = tfp.stats.stddev(observed_data)

    # Make fake_data with the same standard deviation as observed_data.
    fake_data = stddev * tf.random_normal(shape=(100,))
    ```

    Notice we divide by `N` (the numpy default), which does not create `NaN`
    when `N = 1`, but is slightly biased.

    Args:
      x:  A numeric `Tensor` holding samples.
      sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
        `None` (meaning all axis hold samples).
        Default value: `0` (leftmost dimension).
      keepdims:  Boolean.  Whether to keep the sample axis as singletons.
      name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., `'stddev'`).

    Returns:
      stddev: A `Tensor` of same `dtype` as the `x`, and rank equal to
        `rank(x) - len(sample_axis)`


def variance(x, sample_axis=<NUM_LIT:0>, keepdims=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', values=[x, sample_axis]):<EOL><INDENT>return covariance(<EOL>x, y=None, sample_axis=sample_axis, event_axis=None, keepdims=keepdims)<EOL><DEDENT>

Estimate variance using samples.

    Given `N` samples of scalar valued random variable `X`, variance may
    be estimated as

    ```none
    Var[X] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(X_n - Xbar)}
    Xbar := N^{-1} sum_{n=1}^N X_n
    ```

    ```python
    x = tf.random_normal(shape=(100, 2, 3))

    # var[i, j] is the sample variance of the (i, j) batch member of x.
    var = tfp.stats.variance(x, sample_axis=0)
    ```

    Notice we divide by `N` (the numpy default), which does not create `NaN`
    when `N = 1`, but is slightly biased.

    Args:
      x:  A numeric `Tensor` holding samples.
      sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
        `None` (meaning all axis hold samples).
        Default value: `0` (leftmost dimension).
      keepdims:  Boolean.  Whether to keep the sample axis as singletons.
      name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., `'variance'`).

    Returns:
      var: A `Tensor` of same `dtype` as the `x`, and rank equal to
        `rank(x) - len(sample_axis)`


return isinstance(x, (tuple, list))<EOL>

Helper which returns `True` if input is `list`-like.


values = tf.convert_to_tensor(value=values, name='<STR_LIT>')<EOL>values_ = tf.get_static_value(values)<EOL>if values_ is None:<EOL><INDENT>return values + tf.zeros([<NUM_LIT:1>], dtype=values.dtype)<EOL><DEDENT>if values_.ndim > <NUM_LIT:1>:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(values_.shape))<EOL><DEDENT>values_ = values_ + np.zeros([<NUM_LIT:1>], dtype=values_.dtype)<EOL>return list(values_)<EOL>

Return a list (preferred) or 1d Tensor from values, if values.ndims < 2.


axis = _make_list_or_1d_tensor(axis)<EOL>ndims = tf.convert_to_tensor(value=ndims, name='<STR_LIT>', dtype=tf.int32)<EOL>ndims_ = tf.get_static_value(ndims)<EOL>if _is_list_like(axis) and ndims_ is not None:<EOL><INDENT>positive_axis = []<EOL>for a in axis:<EOL><INDENT>if a < <NUM_LIT:0>:<EOL><INDENT>a = ndims_ + a<EOL><DEDENT>positive_axis.append(a)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>axis = tf.convert_to_tensor(value=axis, name='<STR_LIT>', dtype=tf.int32)<EOL>positive_axis = tf.where(axis >= <NUM_LIT:0>, axis, axis + ndims)<EOL><DEDENT>return positive_axis<EOL>

Rectify possibly negatively axis. Prefer return Python list.


x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if axis is None:<EOL><INDENT>return tf.squeeze(x, axis=None)<EOL><DEDENT>axis = tf.convert_to_tensor(value=axis, name='<STR_LIT>', dtype=tf.int32)<EOL>axis += tf.zeros([<NUM_LIT:1>], dtype=axis.dtype)  <EOL>keep_axis, _ = tf.compat.v1.setdiff1d(tf.range(<NUM_LIT:0>, tf.rank(x)), axis)<EOL>return tf.reshape(x, tf.gather(tf.shape(input=x), keep_axis))<EOL>

A version of squeeze that works with dynamic axis.


im_bytes = tf.io.read_file(filepath)<EOL>im = tf.image.decode_image(im_bytes, channels=CHANNELS)<EOL>im = tf.image.convert_image_dtype(im, tf.float32)<EOL>return im<EOL>


return tf.concat(tf.unstack(seq), <NUM_LIT:1>)<EOL>

Joins a sequence side-by-side into a single image.


filepath = os.path.join(FLAGS.data_dir, DATA_SPRITES_DIR)<EOL>if not tf.io.gfile.exists(filepath):<EOL><INDENT>if not tf.io.gfile.exists(FLAGS.data_dir):<EOL><INDENT>tf.io.gfile.makedirs(FLAGS.data_dir)<EOL><DEDENT>zip_name = "<STR_LIT>".format(filepath)<EOL>urllib.request.urlretrieve(DATA_SPRITES_URL, zip_name)<EOL>with zipfile.ZipFile(zip_name, "<STR_LIT:r>") as zip_file:<EOL><INDENT>zip_file.extractall(FLAGS.data_dir)<EOL><DEDENT>tf.io.gfile.remove(zip_name)<EOL><DEDENT>return filepath<EOL>

Downloads the sprites data and returns the saved filepath.


def create_character(skin, hair, top, pants):

dtype = skin.dtype<EOL>hair_mask = tf.cast(hair[..., -<NUM_LIT:1>:] <= <NUM_LIT:0>, dtype)<EOL>top_mask = tf.cast(top[..., -<NUM_LIT:1>:] <= <NUM_LIT:0>, dtype)<EOL>pants_mask = tf.cast(pants[..., -<NUM_LIT:1>:] <= <NUM_LIT:0>, dtype)<EOL>char = (skin * hair_mask) + hair<EOL>char = (char * top_mask) + top<EOL>char = (char * pants_mask) + pants<EOL>return char<EOL>

Creates a character sprite from a set of attribute sprites.


def create_seq(character, action_metadata, direction, length=<NUM_LIT:8>, start=<NUM_LIT:0>):

sprite_start = (action_metadata[<NUM_LIT:0>]+direction) * FRAME_SIZE<EOL>sprite_end = (action_metadata[<NUM_LIT:0>]+direction+<NUM_LIT:1>) * FRAME_SIZE<EOL>sprite_line = character[sprite_start:sprite_end, ...]<EOL>frames = tf.stack(tf.split(sprite_line, <NUM_LIT>, axis=<NUM_LIT:1>))  <EOL>frames = frames[<NUM_LIT:0>:action_metadata[<NUM_LIT:1>]]<EOL>frames = tf.roll(frames, shift=-start, axis=<NUM_LIT:0>)<EOL>frames = tf.tile(frames, [<NUM_LIT:2>, <NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:1>])  <EOL>frames = frames[:length]<EOL>frames = tf.cast(frames, dtype=tf.float32)<EOL>frames.set_shape([length, FRAME_SIZE, FRAME_SIZE, CHANNELS])<EOL>return frames<EOL>

Creates a sequence.

    Args:
      character: A character sprite tensor.
      action_metadata: An action metadata tuple.
      direction: An integer representing the direction, i.e., the row
        offset within each action group corresponding to a particular
        direction.
      length: Desired length of the sequence. If this is longer than
        the number of available frames, it will roll over to the
        beginning.
      start: Index of possible frames at which to start the sequence.

    Returns:
      A sequence tensor.


def create_random_seq(character, action_metadata, direction, length=<NUM_LIT:8>):

start = tf.random.uniform([], maxval=action_metadata[<NUM_LIT:1>], dtype=tf.int32)<EOL>return create_seq(character, action_metadata, direction, length, start)<EOL>


def create_sprites_dataset(characters, actions, directions, channels=<NUM_LIT:3>,<EOL>length=<NUM_LIT:8>, shuffle=False, fake_data=False):

if fake_data:<EOL><INDENT>dummy_image = tf.random.normal([HEIGHT, WIDTH, CHANNELS])<EOL><DEDENT>else:<EOL><INDENT>basedir = download_sprites()<EOL><DEDENT>action_names = [action.name for action in actions]<EOL>action_metadata = [(action.start_row, action.frames) for action in actions]<EOL>direction_rows = [direction.row_offset for direction in directions]<EOL>chars = tf.data.Dataset.from_tensor_slices(characters)<EOL>act_names = tf.data.Dataset.from_tensor_slices(action_names).repeat()<EOL>acts_metadata = tf.data.Dataset.from_tensor_slices(action_metadata).repeat()<EOL>dir_rows = tf.data.Dataset.from_tensor_slices(direction_rows).repeat()<EOL>if shuffle:<EOL><INDENT>chars = chars.shuffle(len(characters))<EOL><DEDENT>dataset = tf.data.Dataset.zip((chars, act_names, acts_metadata, dir_rows))<EOL>skin_table = tf.contrib.lookup.index_table_from_tensor(sorted(SKIN_COLORS))<EOL>hair_table = tf.contrib.lookup.index_table_from_tensor(sorted(HAIRSTYLES))<EOL>top_table = tf.contrib.lookup.index_table_from_tensor(sorted(TOPS))<EOL>pants_table = tf.contrib.lookup.index_table_from_tensor(sorted(PANTS))<EOL>action_table = tf.contrib.lookup.index_table_from_tensor(sorted(action_names))<EOL>def process_example(attrs, act_name, act_metadata, dir_row_offset):<EOL><INDENT>"""<STR_LIT>"""<EOL>skin_name = attrs[<NUM_LIT:0>]<EOL>hair_name = attrs[<NUM_LIT:1>]<EOL>top_name = attrs[<NUM_LIT:2>]<EOL>pants_name = attrs[<NUM_LIT:3>]<EOL>if fake_data:<EOL><INDENT>char = dummy_image<EOL><DEDENT>else:<EOL><INDENT>skin = read_image(basedir + os.sep + skin_name)<EOL>hair = read_image(basedir + os.sep + hair_name)<EOL>top = read_image(basedir + os.sep + top_name)<EOL>pants = read_image(basedir + os.sep + pants_name)<EOL>char = create_character(skin, hair, top, pants)<EOL><DEDENT>if shuffle:<EOL><INDENT>seq = create_random_seq(char, act_metadata, dir_row_offset, length)<EOL><DEDENT>else:<EOL><INDENT>seq = create_seq(char, act_metadata, dir_row_offset, length)<EOL><DEDENT>seq = seq[..., :channels]  <EOL>skin_idx = skin_table.lookup(skin_name)<EOL>hair_idx = hair_table.lookup(hair_name)<EOL>top_idx = top_table.lookup(top_name)<EOL>pants_idx = pants_table.lookup(pants_name)<EOL>act_idx = action_table.lookup(act_name)<EOL>return (seq, skin_idx, hair_idx, top_idx, pants_idx, act_idx,<EOL>skin_name, hair_name, top_name, pants_name, act_name)<EOL><DEDENT>dataset = dataset.map(process_example)<EOL>return dataset<EOL>

Creates a tf.data pipeline for the sprites dataset.

    Args:
      characters: A list of (skin, hair, top, pants) tuples containing
        relative paths to the sprite png image for each attribute.
      actions: A list of Actions.
      directions: A list of Directions.
      channels: Number of image channels to yield.
      length: Desired length of the sequences.
      shuffle: Whether or not to shuffle the characters and sequences
        start frame.
      fake_data: Boolean for whether or not to yield synthetic data.

    Returns:
      A tf.data.Dataset yielding (seq, skin label index, hair label index,
      top label index, pants label index, action label index, skin label
      name, hair label_name, top label name, pants label name, action
      label name) tuples.


def __init__(self, channels=<NUM_LIT:3>, shuffle_train=True, fake_data=False):

super(SpritesDataset, self).__init__()<EOL>self.frame_size = FRAME_SIZE<EOL>self.channels = channels<EOL>self.length = <NUM_LIT:8><EOL>num_train = <NUM_LIT:1000><EOL>num_test = <NUM_LIT><EOL>characters = [(skin, hair, top, pants)<EOL>for skin in sorted(SKIN_COLORS)<EOL>for hair in sorted(HAIRSTYLES)<EOL>for top in sorted(TOPS)<EOL>for pants in sorted(PANTS)]<EOL>random.seed(<NUM_LIT>)<EOL>random.shuffle(characters)<EOL>train_chars = characters[:num_train]<EOL>test_chars = characters[num_train:]<EOL>num_train_actual = len(set(train_chars))<EOL>num_test_actual = len(set(test_chars))<EOL>num_train_test_overlap = len(set(train_chars) & set(test_chars))<EOL>if num_train_actual != num_train:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>".format(<EOL>num_train_actual))<EOL><DEDENT>if num_test_actual != num_test:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>".format(<EOL>num_test_actual))<EOL><DEDENT>if num_train_test_overlap > <NUM_LIT:0>:  <EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>".format(<EOL>num_train_test_overlap))<EOL><DEDENT>self.train = create_sprites_dataset(<EOL>train_chars, ACTIONS, DIRECTIONS, self.channels, self.length,<EOL>shuffle=shuffle_train, fake_data=fake_data)<EOL>self.test = create_sprites_dataset(<EOL>test_chars, ACTIONS, DIRECTIONS, self.channels, self.length,<EOL>shuffle=False, fake_data=fake_data)<EOL>

Creates the SpritesDataset and stores train and test datasets.

        The datasets yield (seq, skin label index, hair label index, top
        label index, pants label index, action label index, skin label name,
        hair label_name, top label name, pants label name, action label
        name) tuples.

        Args:
          channels: Number of image channels to yield.
          shuffle_train: Boolean for whether or not to shuffle the training
            set.
          fake_data: Boolean for whether or not to yield synthetic data.

        Raises:
          ValueError: If the number of training or testing examples is
            incorrect, or if there is overlap betweem the two datasets.


def make_encoder(base_depth, activation, latent_size, code_size):

conv = functools.partial(<EOL>tf.keras.layers.Conv2D, padding="<STR_LIT>", activation=activation)<EOL>encoder_net = tf.keras.Sequential([<EOL>conv(base_depth, <NUM_LIT:5>, <NUM_LIT:1>),<EOL>conv(base_depth, <NUM_LIT:5>, <NUM_LIT:2>),<EOL>conv(<NUM_LIT:2> * base_depth, <NUM_LIT:5>, <NUM_LIT:1>),<EOL>conv(<NUM_LIT:2> * base_depth, <NUM_LIT:5>, <NUM_LIT:2>),<EOL>conv(<NUM_LIT:4> * latent_size, <NUM_LIT:7>, padding="<STR_LIT>"),<EOL>tf.keras.layers.Flatten(),<EOL>tf.keras.layers.Dense(latent_size * code_size, activation=None),<EOL>tf.keras.layers.Reshape([latent_size, code_size])<EOL>])<EOL>def encoder(images):<EOL><INDENT>"""<STR_LIT>"""<EOL>images = <NUM_LIT:2> * tf.cast(images, dtype=tf.float32) - <NUM_LIT:1><EOL>codes = encoder_net(images)<EOL>return codes<EOL><DEDENT>return encoder<EOL>

Creates the encoder function.

    Args:
      base_depth: Layer base depth in encoder net.
      activation: Activation function in hidden layers.
      latent_size: The number of latent variables in the code.
      code_size: The dimensionality of each latent variable.

    Returns:
      encoder: A `callable` mapping a `Tensor` of images to a `Tensor` of shape
        `[..., latent_size, code_size]`.


def make_decoder(base_depth, activation, input_size, output_shape):

deconv = functools.partial(<EOL>tf.keras.layers.Conv2DTranspose, padding="<STR_LIT>", activation=activation)<EOL>conv = functools.partial(<EOL>tf.keras.layers.Conv2D, padding="<STR_LIT>", activation=activation)<EOL>decoder_net = tf.keras.Sequential([<EOL>tf.keras.layers.Reshape((<NUM_LIT:1>, <NUM_LIT:1>, input_size)),<EOL>deconv(<NUM_LIT:2> * base_depth, <NUM_LIT:7>, padding="<STR_LIT>"),<EOL>deconv(<NUM_LIT:2> * base_depth, <NUM_LIT:5>),<EOL>deconv(<NUM_LIT:2> * base_depth, <NUM_LIT:5>, <NUM_LIT:2>),<EOL>deconv(base_depth, <NUM_LIT:5>),<EOL>deconv(base_depth, <NUM_LIT:5>, <NUM_LIT:2>),<EOL>deconv(base_depth, <NUM_LIT:5>),<EOL>conv(output_shape[-<NUM_LIT:1>], <NUM_LIT:5>, activation=None),<EOL>tf.keras.layers.Reshape(output_shape),<EOL>])<EOL>def decoder(codes):<EOL><INDENT>"""<STR_LIT>"""<EOL>logits = decoder_net(codes)<EOL>return tfd.Independent(tfd.Bernoulli(logits=logits),<EOL>reinterpreted_batch_ndims=len(output_shape),<EOL>name="<STR_LIT>")<EOL><DEDENT>return decoder<EOL>

Creates the decoder function.

    Args:
      base_depth: Layer base depth in decoder net.
      activation: Activation function in hidden layers.
      input_size: The flattened latent input shape as an int.
      output_shape: The output image shape as a list.

    Returns:
      decoder: A `callable` mapping a `Tensor` of encodings to a
        `tfd.Distribution` instance over images.


<EOL>updated_ema_count = moving_averages.assign_moving_average(<EOL>vector_quantizer.ema_count,<EOL>tf.reduce_sum(input_tensor=one_hot_assignments, axis=[<NUM_LIT:0>, <NUM_LIT:1>]),<EOL>decay,<EOL>zero_debias=False)<EOL>updated_ema_means = moving_averages.assign_moving_average(<EOL>vector_quantizer.ema_means,<EOL>tf.reduce_sum(<EOL>input_tensor=tf.expand_dims(codes, <NUM_LIT:2>) *<EOL>tf.expand_dims(one_hot_assignments, <NUM_LIT:3>),<EOL>axis=[<NUM_LIT:0>, <NUM_LIT:1>]),<EOL>decay,<EOL>zero_debias=False)<EOL>perturbed_ema_count = updated_ema_count + <NUM_LIT><EOL>with tf.control_dependencies([commitment_loss]):<EOL><INDENT>update_means = tf.compat.v1.assign(<EOL>vector_quantizer.codebook,<EOL>updated_ema_means / perturbed_ema_count[..., tf.newaxis])<EOL>with tf.control_dependencies([update_means]):<EOL><INDENT>return tf.identity(commitment_loss)<EOL><DEDENT><DEDENT>

Add control dependencies to the commmitment loss to update the codebook.

    Args:
      vector_quantizer: An instance of the VectorQuantizer class.
      one_hot_assignments: The one-hot vectors corresponding to the matched
        codebook entry for each code in the batch.
      codes: A `float`-like `Tensor` containing the latent vectors to be compared
        to the codebook.
      commitment_loss: The commitment loss from comparing the encoder outputs to
        their neighboring codebook entries.
      decay: Decay factor for exponential moving average.

    Returns:
      commitment_loss: Commitment loss with control dependencies.


n = x.shape[<NUM_LIT:0>]<EOL>fig = figure.Figure(figsize=(n, <NUM_LIT:1>), frameon=False)<EOL>canvas = backend_agg.FigureCanvasAgg(fig)<EOL>for i in range(n):<EOL><INDENT>ax = fig.add_subplot(<NUM_LIT:1>, n, i+<NUM_LIT:1>)<EOL>ax.imshow(x[i].squeeze(),<EOL>interpolation="<STR_LIT:none>",<EOL>cmap=cm.get_cmap("<STR_LIT>"))<EOL>ax.axis("<STR_LIT>")<EOL><DEDENT>canvas.print_figure(fname, format="<STR_LIT>")<EOL>print("<STR_LIT>" % fname)<EOL>

Helper method to save a grid of images to a PNG file.

    Args:
      x: A numpy array of shape [n_images, height, width].
      fname: The filename to write to (including extension).


Helper method to save images visualizing model reconstructions.

    Args:
      images_val: Numpy array containing a batch of input images.
      reconstructed_images_val: Numpy array giving the expected output
        (mean) of the decoder.
      random_images_val: Optionally, a Numpy array giving the expected output
        (mean) of decoding samples from the prior, or `None`.
      log_dir: The directory to write images (Python `str`).
      prefix: A specific label for the saved visualizations, which
        determines their filenames (Python `str`).
      viz_n: The number of images from each batch to visualize (Python `int`).


class Dummy(object):<EOL><INDENT>pass<EOL><DEDENT>num_examples = <NUM_LIT:10><EOL>mnist_data = Dummy()<EOL>mnist_data.train = Dummy()<EOL>mnist_data.train.images = np.float32(np.random.randn(<EOL>num_examples, np.prod(IMAGE_SHAPE)))<EOL>mnist_data.train.labels = np.int32(np.random.permutation(<EOL>np.arange(num_examples)))<EOL>mnist_data.train.num_examples = num_examples<EOL>mnist_data.validation = Dummy()<EOL>mnist_data.validation.images = np.float32(np.random.randn(<EOL>num_examples, np.prod(IMAGE_SHAPE)))<EOL>mnist_data.validation.labels = np.int32(np.random.permutation(<EOL>np.arange(num_examples)))<EOL>mnist_data.validation.num_examples = num_examples<EOL>return mnist_data<EOL>

Builds fake MNIST-style data for unit testing.


filepath = os.path.join(directory, filename)<EOL>if tf.io.gfile.exists(filepath):<EOL><INDENT>return filepath<EOL><DEDENT>if not tf.io.gfile.exists(directory):<EOL><INDENT>tf.io.gfile.makedirs(directory)<EOL><DEDENT>url = os.path.join(BERNOULLI_PATH, filename)<EOL>print("<STR_LIT>" % (url, filepath))<EOL>urllib.request.urlretrieve(url, filepath)<EOL>return filepath<EOL>


amat_file = download(directory, FILE_TEMPLATE.format(split=split_name))<EOL>dataset = tf.data.TextLineDataset(amat_file)<EOL>str_to_arr = lambda string: np.array([c == b"<STR_LIT:1>" for c in string.split()])<EOL>def _parser(s):<EOL><INDENT>booltensor = tf.compat.v1.py_func(str_to_arr, [s], tf.bool)<EOL>reshaped = tf.reshape(booltensor, [<NUM_LIT>, <NUM_LIT>, <NUM_LIT:1>])<EOL>return tf.cast(reshaped, dtype=tf.float32), tf.constant(<NUM_LIT:0>, tf.int32)<EOL><DEDENT>return dataset.map(_parser)<EOL>

Returns Hugo Larochelle's binary static MNIST tf.data.Dataset.


def build_input_pipeline(data_dir, batch_size, heldout_size, mnist_type):

<EOL>if mnist_type in [MnistType.FAKE_DATA, MnistType.THRESHOLD]:<EOL><INDENT>if mnist_type == MnistType.FAKE_DATA:<EOL><INDENT>mnist_data = build_fake_data()<EOL><DEDENT>else:<EOL><INDENT>mnist_data = mnist.read_data_sets(data_dir)<EOL><DEDENT>training_dataset = tf.data.Dataset.from_tensor_slices(<EOL>(mnist_data.train.images, np.int32(mnist_data.train.labels)))<EOL>heldout_dataset = tf.data.Dataset.from_tensor_slices(<EOL>(mnist_data.validation.images,<EOL>np.int32(mnist_data.validation.labels)))<EOL><DEDENT>elif mnist_type == MnistType.BERNOULLI:<EOL><INDENT>training_dataset = load_bernoulli_mnist_dataset(data_dir, "<STR_LIT:train>")<EOL>heldout_dataset = load_bernoulli_mnist_dataset(data_dir, "<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>training_batches = training_dataset.repeat().batch(batch_size)<EOL>training_iterator = tf.compat.v1.data.make_one_shot_iterator(training_batches)<EOL>heldout_frozen = (heldout_dataset.take(heldout_size).<EOL>repeat().batch(heldout_size))<EOL>heldout_iterator = tf.compat.v1.data.make_one_shot_iterator(heldout_frozen)<EOL>handle = tf.compat.v1.placeholder(tf.string, shape=[])<EOL>feedable_iterator = tf.compat.v1.data.Iterator.from_string_handle(<EOL>handle, training_batches.output_types, training_batches.output_shapes)<EOL>images, labels = feedable_iterator.get_next()<EOL>images = tf.reshape(images, shape=[-<NUM_LIT:1>] + IMAGE_SHAPE)<EOL>if mnist_type in [MnistType.FAKE_DATA, MnistType.THRESHOLD]:<EOL><INDENT>images = tf.cast(images > <NUM_LIT:0.5>, dtype=tf.int32)<EOL><DEDENT>return images, labels, handle, training_iterator, heldout_iterator<EOL>

Builds an Iterator switching between train and heldout data.


distances = tf.norm(<EOL>tensor=tf.expand_dims(codes, <NUM_LIT:2>) -<EOL>tf.reshape(self.codebook, [<NUM_LIT:1>, <NUM_LIT:1>, self.num_codes, self.code_size]),<EOL>axis=<NUM_LIT:3>)<EOL>assignments = tf.argmin(input=distances, axis=<NUM_LIT:2>)<EOL>one_hot_assignments = tf.one_hot(assignments, depth=self.num_codes)<EOL>nearest_codebook_entries = tf.reduce_sum(<EOL>input_tensor=tf.expand_dims(one_hot_assignments, -<NUM_LIT:1>) *<EOL>tf.reshape(self.codebook, [<NUM_LIT:1>, <NUM_LIT:1>, self.num_codes, self.code_size]),<EOL>axis=<NUM_LIT:2>)<EOL>return nearest_codebook_entries, one_hot_assignments<EOL>

Uses codebook to find nearest neighbor for each code.

        Args:
          codes: A `float`-like `Tensor` containing the latent
            vectors to be compared to the codebook. These are rank-3 with shape
            `[batch_size, latent_size, code_size]`.

        Returns:
          nearest_codebook_entries: The 1-nearest neighbor in Euclidean distance for
            each code in the batch.
          one_hot_assignments: The one-hot vectors corresponding to the matched
            codebook entry for each code in the batch.


def image_summary(seqs, name, num=None):

seqs = tf.clip_by_value(seqs, <NUM_LIT:0.>, <NUM_LIT:1.>)<EOL>seqs = tf.unstack(seqs[:num])<EOL>joined_seqs = [tf.concat(tf.unstack(seq), <NUM_LIT:1>) for seq in seqs]<EOL>joined_seqs = tf.expand_dims(tf.concat(joined_seqs, <NUM_LIT:0>), <NUM_LIT:0>)<EOL>tf.compat.v2.summary.image(<EOL>name,<EOL>joined_seqs,<EOL>max_outputs=<NUM_LIT:1>,<EOL>step=tf.compat.v1.train.get_or_create_global_step())<EOL>

Visualizes sequences as TensorBoard summaries.

    Args:
      seqs: A tensor of shape [n, t, h, w, c].
      name: String name of this summary.
      num: Integer for the number of examples to visualize. Defaults to
        all examples.


def visualize_reconstruction(inputs, reconstruct, num=<NUM_LIT:3>, name="<STR_LIT>"):

reconstruct = tf.clip_by_value(reconstruct, <NUM_LIT:0.>, <NUM_LIT:1.>)<EOL>inputs_and_reconstruct = tf.concat((inputs[:num], reconstruct[:num]), axis=<NUM_LIT:0>)<EOL>image_summary(inputs_and_reconstruct, name)<EOL>

Visualizes the reconstruction of inputs in TensorBoard.

    Args:
      inputs: A tensor of the original inputs, of shape [batch, timesteps,
        h, w, c].
      reconstruct: A tensor of a reconstruction of inputs, of shape
        [batch, timesteps, h, w, c].
      num: Integer for the number of examples to visualize.
      name: String name of this summary.


def visualize_qualitative_analysis(inputs, model, samples=<NUM_LIT:1>, batch_size=<NUM_LIT:3>,<EOL>length=<NUM_LIT:8>):

average = lambda dist: tf.reduce_mean(<EOL>input_tensor=dist.mean(), axis=<NUM_LIT:0>)  <EOL>with tf.compat.v1.name_scope("<STR_LIT>"):<EOL><INDENT>reconstruct = functools.partial(model.reconstruct, inputs=inputs,<EOL>samples=samples)<EOL>visualize_reconstruction(inputs, average(reconstruct()))<EOL>visualize_reconstruction(inputs, average(reconstruct(sample_static=True)),<EOL>name="<STR_LIT>")<EOL>visualize_reconstruction(inputs, average(reconstruct(sample_dynamic=True)),<EOL>name="<STR_LIT>")<EOL>visualize_reconstruction(inputs, average(reconstruct(swap_static=True)),<EOL>name="<STR_LIT>")<EOL>visualize_reconstruction(inputs, average(reconstruct(swap_dynamic=True)),<EOL>name="<STR_LIT>")<EOL><DEDENT>with tf.compat.v1.name_scope("<STR_LIT>"):<EOL><INDENT>generate = functools.partial(model.generate, batch_size=batch_size,<EOL>length=length, samples=samples)<EOL>image_summary(average(generate(fix_static=True)), "<STR_LIT>")<EOL>image_summary(average(generate(fix_dynamic=True)), "<STR_LIT>")<EOL><DEDENT>

Visualizes a qualitative analysis of a given model.

    Args:
      inputs: A tensor of the original inputs, of shape [batch, timesteps,
        h, w, c].
      model: A DisentangledSequentialVAE model.
      samples: Number of samples to draw from the latent distributions.
      batch_size: Number of sequences to generate.
      length: Number of timesteps to generate for each sequence.


def summarize_dist_params(dist, name, name_scope="<STR_LIT>"):

Summarize the parameters of a distribution.

    Args:
      dist: A Distribution object with mean and standard deviation
        parameters.
      name: The name of the distribution.
      name_scope: The name scope of this summary.


mean = tf.reduce_mean(input_tensor=inputs)<EOL>with tf.compat.v1.name_scope(nats_name_scope):<EOL><INDENT>tf.compat.v2.summary.scalar(<EOL>name,<EOL>mean,<EOL>step=tf.compat.v1.train.get_or_create_global_step())<EOL><DEDENT>with tf.compat.v1.name_scope(bits_name_scope):<EOL><INDENT>tf.compat.v2.summary.scalar(<EOL>name,<EOL>mean / units / tf.math.log(<NUM_LIT>),<EOL>step=tf.compat.v1.train.get_or_create_global_step())<EOL><DEDENT>

Summarize the mean of a tensor in nats and bits per unit.

    Args:
      inputs: A tensor of values measured in nats.
      units: The units of the tensor with which to compute the mean bits
        per unit.
      name: The name of the tensor.
      nats_name_scope: The name scope of the nats summary.
      bits_name_scope: The name scope of the bits summary.


super(LearnableMultivariateNormalDiag, self).__init__()<EOL>with tf.compat.v1.name_scope(self._name):<EOL><INDENT>self.dimensions = dimensions<EOL>self._mean = tf.compat.v2.Variable(<EOL>tf.random.normal([dimensions], stddev=<NUM_LIT:0.1>), name="<STR_LIT>")<EOL>self._untransformed_stddev = tf.compat.v2.Variable(<EOL>tf.random.normal([dimensions], mean=<NUM_LIT>, stddev=<NUM_LIT:0.1>),<EOL>name="<STR_LIT>")<EOL><DEDENT>

Constructs a learnable multivariate diagonal normal model.

        Args:
          dimensions: An integer corresponding to the dimensionality of the
            distribution.


del inputs  <EOL>with tf.compat.v1.name_scope(self._name):<EOL><INDENT>return tfd.MultivariateNormalDiag(self.loc, self.scale_diag)<EOL><DEDENT>

Runs the model to generate multivariate normal distribution.

        Args:
          inputs: Unused.

        Returns:
          A MultivariateNormalDiag distribution with event shape
          [dimensions], batch shape [], and sample shape [sample_shape,
          dimensions].


return self._mean<EOL>


return tf.nn.softplus(self._untransformed_stddev) + <NUM_LIT><EOL>

The diagonal standard deviation of the normal distribution.


def __init__(self, dimensions, hidden_size):

super(LearnableMultivariateNormalDiagCell, self).__init__()<EOL>self.dimensions = dimensions<EOL>self.hidden_size = hidden_size<EOL>self.lstm_cell = tf.keras.layers.LSTMCell(hidden_size)<EOL>self.output_layer = tf.keras.layers.Dense(<NUM_LIT:2>*dimensions)<EOL>

Constructs a learnable multivariate diagonal normal cell.

        Args:
          dimensions: An integer corresponding to the dimensionality of the
            distribution.
          hidden_size: Dimensionality of the LSTM function parameters.


h0 = tf.zeros([<NUM_LIT:1>, self.hidden_size])<EOL>c0 = tf.zeros([<NUM_LIT:1>, self.hidden_size])<EOL>combined_shape = tf.concat((tf.convert_to_tensor(<EOL>value=sample_batch_shape, dtype=tf.int32), [self.dimensions]),<EOL>axis=-<NUM_LIT:1>)<EOL>previous_output = tf.zeros(combined_shape)<EOL>return previous_output, (h0, c0)<EOL>

Returns an initial state for the LSTM cell.

        Args:
          sample_batch_shape: A 0D or 1D tensor of the combined sample and
            batch shape.

        Returns:
          A tuple of the initial previous output at timestep 0 of shape
          [sample_batch_shape, dimensions], and the cell state.


<EOL>original_shape = inputs.shape<EOL>if len(original_shape) < <NUM_LIT:2>:<EOL><INDENT>inputs = tf.reshape(inputs, [<NUM_LIT:1>, -<NUM_LIT:1>])<EOL><DEDENT>out, state = self.lstm_cell(inputs, state)<EOL>out = self.output_layer(out)<EOL>correct_shape = tf.concat((original_shape[:-<NUM_LIT:1>], tf.shape(input=out)[-<NUM_LIT:1>:]),<EOL><NUM_LIT:0>)<EOL>out = tf.reshape(out, correct_shape)<EOL>loc = out[..., :self.dimensions]<EOL>scale_diag = tf.nn.softplus(out[..., self.dimensions:]) + <NUM_LIT>  <EOL>return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag), state<EOL>

Runs the model to generate a distribution for a single timestep.

        This generates a batched MultivariateNormalDiag distribution using
        the output of the recurrent model at the current timestep to
        parameterize the distribution.

        Args:
          inputs: The sampled value of `z` at the previous timestep, i.e.,
            `z_{t-1}`, of shape [..., dimensions].
            `z_0` should be set to the empty matrix.
          state: A tuple containing the (hidden, cell) state.

        Returns:
          A tuple of a MultivariateNormalDiag distribution, and the state of
          the recurrent function at the end of the current timestep. The
          distribution will have event shape [dimensions], batch shape
          [...], and sample shape [sample_shape, ..., dimensions].


def __init__(self, hidden_size, channels=<NUM_LIT:3>):

super(Decoder, self).__init__()<EOL>self.hidden_size = hidden_size<EOL>activation = tf.nn.leaky_relu<EOL>self.dense = tf.keras.layers.Dense(hidden_size, activation=activation)<EOL>conv_transpose = functools.partial(<EOL>tf.keras.layers.Conv2DTranspose, padding="<STR_LIT>", activation=activation)<EOL>self.conv_transpose1 = conv_transpose(<NUM_LIT>, <NUM_LIT:8>, <NUM_LIT:1>, padding="<STR_LIT>")<EOL>self.conv_transpose2 = conv_transpose(<NUM_LIT>, <NUM_LIT:3>, <NUM_LIT:2>)<EOL>self.conv_transpose3 = conv_transpose(<NUM_LIT>, <NUM_LIT:3>, <NUM_LIT:2>)<EOL>self.conv_transpose4 = conv_transpose(channels, <NUM_LIT:3>, <NUM_LIT:2>, activation=None)<EOL>

Constructs a probabilistic decoder.

        For each timestep, this model takes as input a concatenation of the
        dynamic and static latent variables `z_t` and `f`, respectively,
        outputs an intermediate representation via an affine function (i.e.,
        a one hidden layer MLP), then transforms this with four transpose
        convolution layers and up-sampling to the spatial shape of `x_t`.

        Together with the priors, this allows us to specify a generative
        model

        ```none
        p(x_{1:T}, z_{1:T}, f) = p(f) prod_{t=1}^T p(z_t | z_{<t}) p(x_t | z_t, f).
        ```

        Args:
          hidden_size: Dimensionality of the intermediate representations.
          channels: The depth of the output tensor.


<EOL>dynamic, static = inputs<EOL>timesteps = tf.shape(input=dynamic)[-<NUM_LIT:2>]<EOL>static = static[..., tf.newaxis, :] + tf.zeros([timesteps, <NUM_LIT:1>])<EOL>latents = tf.concat([dynamic, static], axis=-<NUM_LIT:1>)  <EOL>out = self.dense(latents)<EOL>out = tf.reshape(out, (-<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:1>, self.hidden_size))<EOL>out = self.conv_transpose1(out)<EOL>out = self.conv_transpose2(out)<EOL>out = self.conv_transpose3(out)<EOL>out = self.conv_transpose4(out)  <EOL>expanded_shape = tf.concat(<EOL>(tf.shape(input=latents)[:-<NUM_LIT:1>], tf.shape(input=out)[<NUM_LIT:1>:]), axis=<NUM_LIT:0>)<EOL>out = tf.reshape(out, expanded_shape)  <EOL>return tfd.Independent(<EOL>distribution=tfd.Normal(loc=out, scale=<NUM_LIT:1.>),<EOL>reinterpreted_batch_ndims=<NUM_LIT:3>,  <EOL>name="<STR_LIT>")<EOL>

Runs the model to generate a distribution p(x_t | z_t, f).

        Args:
          inputs: A tuple of (z_{1:T}, f), where `z_{1:T}` is a tensor of
            shape [..., batch_size, timesteps, latent_size_dynamic], and `f`
            is of shape [..., batch_size, latent_size_static].

        Returns:
          A batched Independent distribution wrapping a set of Normal
          distributions over the pixels of x_t, where the Independent
          distribution has event shape [height, width, channels], batch
          shape [batch_size, timesteps], and sample shape [sample_shape,
          batch_size, timesteps, height, width, channels].


super(Compressor, self).__init__()<EOL>self.hidden_size = hidden_size<EOL>conv = functools.partial(<EOL>tf.keras.layers.Conv2D, padding="<STR_LIT>", activation=tf.nn.leaky_relu)<EOL>self.conv1 = conv(<NUM_LIT>, <NUM_LIT:3>, <NUM_LIT:2>)<EOL>self.conv2 = conv(<NUM_LIT>, <NUM_LIT:3>, <NUM_LIT:2>)<EOL>self.conv3 = conv(<NUM_LIT>, <NUM_LIT:3>, <NUM_LIT:2>)<EOL>self.conv4 = conv(hidden_size, <NUM_LIT:8>, padding="<STR_LIT>")<EOL>

Constructs a convolutional compressor.

        This model takes as input `x_{1:T}` and outputs an intermediate
        representation for use in downstream probabilistic encoders.

        Args:
          hidden_size: Dimensionality of the intermediate representations.


image_shape = tf.shape(input=inputs)[-<NUM_LIT:3>:]<EOL>collapsed_shape = tf.concat(([-<NUM_LIT:1>], image_shape), axis=<NUM_LIT:0>)<EOL>out = tf.reshape(inputs, collapsed_shape)  <EOL>out = self.conv1(out)<EOL>out = self.conv2(out)<EOL>out = self.conv3(out)<EOL>out = self.conv4(out)<EOL>expanded_shape = tf.concat((tf.shape(input=inputs)[:-<NUM_LIT:3>], [-<NUM_LIT:1>]), axis=<NUM_LIT:0>)<EOL>return tf.reshape(out, expanded_shape)<EOL>

Runs the model to generate an intermediate representation of x_t.

        Args:
          inputs: A batch of image sequences `x_{1:T}` of shape
            `[sample_shape, batch_size, timesteps, height, width,
            channels]`.

        Returns:
          A batch of intermediate representations of shape [sample_shape,
          batch_size, timesteps, hidden_size].


def __init__(self, latent_size, hidden_size):

super(EncoderStatic, self).__init__()<EOL>self.latent_size = latent_size<EOL>self.hidden_size = hidden_size<EOL>self.bilstm = tf.keras.layers.Bidirectional(<EOL>tf.keras.layers.LSTM(hidden_size),<EOL>merge_mode="<STR_LIT>")<EOL>self.output_layer = tf.keras.layers.Dense(<NUM_LIT:2>*latent_size)<EOL>

Constructs an encoder for `f`.

        Args:
          latent_size: An integer corresponding to the dimensionality of the
            distribution.
          hidden_size: Dimensionality of the LSTM, RNN, and affine function
            parameters.


<EOL>collapsed_shape = tf.concat(([-<NUM_LIT:1>], tf.shape(input=inputs)[-<NUM_LIT:2>:]), axis=<NUM_LIT:0>)<EOL>out = tf.reshape(inputs, collapsed_shape)  <EOL>out = self.bilstm(out)  <EOL>expanded_shape = tf.concat((tf.shape(input=inputs)[:-<NUM_LIT:2>], [-<NUM_LIT:1>]), axis=<NUM_LIT:0>)<EOL>out = tf.reshape(out, expanded_shape)  <EOL>out = self.output_layer(out)  <EOL>loc = out[..., :self.latent_size]<EOL>scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + <NUM_LIT>  <EOL>return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)<EOL>

Runs the model to generate a distribution `q(f | x_{1:T})`.

        This generates a list of batched MultivariateNormalDiag
        distributions using the output of the recurrent model at each
        timestep to parameterize each distribution.

        Args:
          inputs: A batch of intermediate representations of image frames
            across all timesteps, of shape [..., batch_size, timesteps,
            hidden_size].

        Returns:
          A batched MultivariateNormalDiag distribution with event shape
          [latent_size], batch shape [..., batch_size], and sample shape
          [sample_shape, ..., batch_size, latent_size].


super(EncoderDynamicFactorized, self).__init__()<EOL>self.latent_size = latent_size<EOL>self.hidden_size = hidden_size<EOL>self.dense = tf.keras.layers.Dense(hidden_size, activation=tf.nn.leaky_relu)<EOL>self.output_layer = tf.keras.layers.Dense(<NUM_LIT:2>*latent_size)<EOL>

Constructs a "factorized" encoder for `z_t`.

        Args:
          latent_size: An integer corresponding to the
            dimensionality of the distribution.
          hidden_size: Dimensionality of the affine function parameters.


out = self.dense(inputs)  <EOL>out = self.output_layer(out)  <EOL>loc = out[..., :self.latent_size]<EOL>scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + <NUM_LIT>  <EOL>return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)<EOL>

Runs the model to generate a distribution `q(z_{1:T} | x_{1:T})`.

        Args:
          inputs: A batch of intermediate representations of image frames
            across all timesteps, of shape [..., batch_size, timesteps,
            hidden_size].

        Returns:
          A batch of MultivariateNormalDiag distributions with event shape
          [latent_size], batch shape [..., batch_size, timesteps], and
          sample shape [sample_shape, ..., batch_size, timesteps,
          latent_size].


super(EncoderDynamicFull, self).__init__()<EOL>self.latent_size = latent_size<EOL>self.hidden_size = hidden_size<EOL>self.bilstm = tf.keras.layers.Bidirectional(<EOL>tf.keras.layers.LSTM(hidden_size, return_sequences=True),<EOL>merge_mode="<STR_LIT>")<EOL>self.rnn = tf.keras.layers.SimpleRNN(hidden_size, return_sequences=True)<EOL>self.output_layer = tf.keras.layers.Dense(<NUM_LIT:2>*latent_size)<EOL>

Constructs a "full" encoder for `z_t`.

        Args:
          latent_size: An integer corresponding to the
            dimensionality of the distribution.
          hidden_size: Dimensionality of the LSTM, RNN, and affine function
            parameters.


<EOL>features, static_sample = inputs<EOL>length = tf.shape(input=features)[-<NUM_LIT:2>]<EOL>static_sample = static_sample[..., tf.newaxis, :] + tf.zeros([length, <NUM_LIT:1>])<EOL>sample_shape_static = tf.shape(input=static_sample)[:-<NUM_LIT:3>]<EOL>sample_shape_inputs = tf.shape(input=features)[:-<NUM_LIT:3>]<EOL>broadcast_shape_inputs = tf.concat((sample_shape_static, [<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:1>]), <NUM_LIT:0>)<EOL>broadcast_shape_static = tf.concat((sample_shape_inputs, [<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:1>]), <NUM_LIT:0>)<EOL>features = features + tf.zeros(broadcast_shape_inputs)<EOL>static_sample = static_sample + tf.zeros(broadcast_shape_static)<EOL>combined = tf.concat((features, static_sample), axis=-<NUM_LIT:1>)<EOL>collapsed_shape = tf.concat(([-<NUM_LIT:1>], tf.shape(input=combined)[-<NUM_LIT:2>:]), axis=<NUM_LIT:0>)<EOL>out = tf.reshape(combined, collapsed_shape)<EOL>out = self.bilstm(out)  <EOL>out = self.rnn(out)  <EOL>expanded_shape = tf.concat(<EOL>(tf.shape(input=combined)[:-<NUM_LIT:2>], tf.shape(input=out)[<NUM_LIT:1>:]), axis=<NUM_LIT:0>)<EOL>out = tf.reshape(out, expanded_shape)  <EOL>out = self.output_layer(out)  <EOL>loc = out[..., :self.latent_size]<EOL>scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + <NUM_LIT>  <EOL>return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)<EOL>

Runs the model to generate a distribution `q(z_{1:T} | x_{1:T}, f)`.

        This generates a list of batched MultivariateNormalDiag
        distributions using the output of the recurrent model at each
        timestep to parameterize each distribution.

        Args:
          inputs: A tuple of a batch of intermediate representations of
            image frames across all timesteps of shape [..., batch_size,
            timesteps, dimensions], and a sample of the static latent
            variable `f` of shape [..., batch_size, latent_size].

        Returns:
          A batch of MultivariateNormalDiag distributions with event shape
          [latent_size], batch shape [broadcasted_shape, batch_size,
          timesteps], and sample shape [sample_shape, broadcasted_shape,
          batch_size, timesteps, latent_size], where `broadcasted_shape` is
          the broadcasted sampled shape between the inputs and static
          sample.


def __init__(self, latent_size_static, latent_size_dynamic,<EOL>hidden_size, channels, latent_posterior):

super(DisentangledSequentialVAE, self).__init__()<EOL>self.latent_size_static = latent_size_static<EOL>self.latent_size_dynamic = latent_size_dynamic<EOL>self.hidden_size = hidden_size<EOL>self.channels = channels<EOL>self.latent_posterior = latent_posterior<EOL>self.static_prior = LearnableMultivariateNormalDiag(latent_size_static)<EOL>self.dynamic_prior = LearnableMultivariateNormalDiagCell(<EOL>latent_size_dynamic, hidden_size)<EOL>self.decoder = Decoder(hidden_size, channels)<EOL>self.compressor = Compressor(hidden_size)<EOL>self.static_encoder = EncoderStatic(latent_size_static, hidden_size)<EOL>if latent_posterior == "<STR_LIT>":<EOL><INDENT>self.dynamic_encoder = EncoderDynamicFactorized(<EOL>latent_size_dynamic, hidden_size)<EOL><DEDENT>else:<EOL><INDENT>self.dynamic_encoder = EncoderDynamicFull(<EOL>latent_size_dynamic, hidden_size)<EOL><DEDENT>

Constructs a Disentangled Sequential Variational Autoencoder.

        Args:
          latent_size_static: Integer dimensionality of the static,
            time-invariant latent variable `f`.
          latent_size_dynamic: Integer dimensionality of each dynamic,
            time-variant latent variable `z_t`.
          hidden_size: Integer dimensionality of the model intermediates.
          channels: Integer depth of the output of the decoder.
          latent_posterior: Either "factorized" or "full" to indicate the
            formulation for the latent posterior `q`.


def generate(self, batch_size, length, samples=<NUM_LIT:1>, fix_static=False,<EOL>fix_dynamic=False):

static_sample, _ = self.sample_static_prior(samples, batch_size, fix_static)<EOL>dynamic_sample, _ = self.sample_dynamic_prior(samples, batch_size, length,<EOL>fix_dynamic)<EOL>likelihood = self.decoder((dynamic_sample, static_sample))<EOL>return likelihood<EOL>

Generate new sequences.

        Args:
          batch_size: Number of sequences to generate.
          length: Number of timesteps to generate for each sequence.
          samples: Number of samples to draw from the latent distributions.
          fix_static: Boolean for whether or not to share the same random
            sample of the static latent variable `f` from its prior across
            all examples.
          fix_dynamic: Boolean for whether or not to share the same random
            sample of the dynamic latent variable `z_{1:T}` from its prior
            across all examples.

        Returns:
          A batched Independent distribution wrapping a set of Normal
          distributions over the pixels of the generated sequences, where
          the Independent distribution has event shape [height, width,
          channels], batch shape [samples, batch_size, timesteps], and
          sample shape [sample_shape, samples, batch_size, timesteps,
          height, width, channels].


def reconstruct(self, inputs, samples=<NUM_LIT:1>, sample_static=False,<EOL>sample_dynamic=False, swap_static=False, swap_dynamic=False,<EOL>fix_static=False, fix_dynamic=False):

batch_size = tf.shape(input=inputs)[-<NUM_LIT:5>]<EOL>length = len(tf.unstack(inputs, axis=-<NUM_LIT:4>))  <EOL>features = self.compressor(inputs)  <EOL>if sample_static:<EOL><INDENT>static_sample, _ = self.sample_static_prior(<EOL>samples, batch_size, fix_static)<EOL><DEDENT>else:<EOL><INDENT>static_sample, _ = self.sample_static_posterior(features, samples)<EOL><DEDENT>if swap_static:<EOL><INDENT>static_sample = tf.reverse(static_sample, axis=[<NUM_LIT:1>])<EOL><DEDENT>if sample_dynamic:<EOL><INDENT>dynamic_sample, _ = self.sample_dynamic_prior(<EOL>samples, batch_size, length, fix_dynamic)<EOL><DEDENT>else:<EOL><INDENT>dynamic_sample, _ = self.sample_dynamic_posterior(<EOL>features, samples, static_sample)<EOL><DEDENT>if swap_dynamic:<EOL><INDENT>dynamic_sample = tf.reverse(dynamic_sample, axis=[<NUM_LIT:1>])<EOL><DEDENT>likelihood = self.decoder((dynamic_sample, static_sample))<EOL>return likelihood<EOL>

Reconstruct the given input sequences.

        Args:
          inputs: A batch of image sequences `x_{1:T}` of shape
            `[batch_size, timesteps, height, width, channels]`.
          samples: Number of samples to draw from the latent distributions.
          sample_static: Boolean for whether or not to randomly sample the
            static latent variable `f` from its prior distribution.
          sample_dynamic: Boolean for whether or not to randomly sample the
            dynamic latent variable `z_{1:T}` from its prior distribution.
          swap_static: Boolean for whether or not to swap the encodings for
            the static latent variable `f` between the examples.
          swap_dynamic: Boolean for whether or not to swap the encodings for
            the dynamic latent variable `z_{1:T}` between the examples.
          fix_static: Boolean for whether or not to share the same random
            sample of the static latent variable `f` from its prior across
            all examples.
          fix_dynamic: Boolean for whether or not to share the same random
            sample of the dynamic latent variable `z_{1:T}` from its prior
            across all examples.

        Returns:
          A batched Independent distribution wrapping a set of Normal
          distributions over the pixels of the reconstruction of the input,
          where the Independent distribution has event shape [height, width,
          channels], batch shape [samples, batch_size, timesteps], and
          sample shape [sample_shape, samples, batch_size, timesteps,
          height, width, channels].


def sample_static_prior(self, samples, batch_size, fixed=False):

dist = self.static_prior()<EOL>if fixed:  <EOL><INDENT>sample = dist.sample((samples, <NUM_LIT:1>)) + tf.zeros([batch_size, <NUM_LIT:1>])<EOL><DEDENT>else:<EOL><INDENT>sample = dist.sample((samples, batch_size))<EOL><DEDENT>return sample, dist<EOL>

Sample the static latent prior.

        Args:
          samples: Number of samples to draw from the latent distribution.
          batch_size: Number of sequences to sample.
          fixed: Boolean for whether or not to share the same random
            sample across all sequences.

        Returns:
          A tuple of a sample tensor of shape [samples, batch_size,
          latent_size], and a MultivariateNormalDiag distribution from which
          the tensor was sampled, with event shape [latent_size], and batch
          shape [].


def sample_static_posterior(self, inputs, samples):

dist = self.static_encoder(inputs)<EOL>sample = dist.sample(samples)<EOL>return sample, dist<EOL>

Sample the static latent posterior.

        Args:
          inputs: A batch of intermediate representations of image frames
            across all timesteps, of shape [..., batch_size, timesteps,
            hidden_size].
          samples: Number of samples to draw from the latent distribution.

        Returns:
          A tuple of a sample tensor of shape [samples, batch_size,
          latent_size], and a MultivariateNormalDiag distribution from which
          the tensor was sampled, with event shape [latent_size], and batch
          shape [..., batch_size].


def sample_dynamic_prior(self, samples, batch_size, length, fixed=False):

if fixed:<EOL><INDENT>sample_batch_size = <NUM_LIT:1><EOL><DEDENT>else:<EOL><INDENT>sample_batch_size = batch_size<EOL><DEDENT>sample, state = self.dynamic_prior.zero_state([samples, sample_batch_size])<EOL>locs = []<EOL>scale_diags = []<EOL>sample_list = []<EOL>for _ in range(length):<EOL><INDENT>dist, state = self.dynamic_prior(sample, state)<EOL>sample = dist.sample()<EOL>locs.append(dist.parameters["<STR_LIT>"])<EOL>scale_diags.append(dist.parameters["<STR_LIT>"])<EOL>sample_list.append(sample)<EOL><DEDENT>sample = tf.stack(sample_list, axis=<NUM_LIT:2>)<EOL>loc = tf.stack(locs, axis=<NUM_LIT:2>)<EOL>scale_diag = tf.stack(scale_diags, axis=<NUM_LIT:2>)<EOL>if fixed:  <EOL><INDENT>sample = sample + tf.zeros([batch_size, <NUM_LIT:1>, <NUM_LIT:1>])<EOL><DEDENT>return sample, tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)<EOL>

Sample the dynamic latent prior.

        Args:
          samples: Number of samples to draw from the latent distribution.
          batch_size: Number of sequences to sample.
          length: Number of timesteps to sample for each sequence.
          fixed: Boolean for whether or not to share the same random
            sample across all sequences.

        Returns:
          A tuple of a sample tensor of shape [samples, batch_size, length
          latent_size], and a MultivariateNormalDiag distribution from which
          the tensor was sampled, with event shape [latent_size], and batch
          shape [samples, 1, length] if fixed or [samples, batch_size,
          length] otherwise.


def sample_dynamic_posterior(self, inputs, samples, static_sample=None):

if self.latent_posterior == "<STR_LIT>":<EOL><INDENT>dist = self.dynamic_encoder(inputs)<EOL>samples = dist.sample(samples)  <EOL><DEDENT>else:  <EOL><INDENT>if static_sample is None:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>dist = self.dynamic_encoder((inputs, static_sample))<EOL>samples = dist.sample()  <EOL><DEDENT>return samples, dist<EOL>

Sample the static latent posterior.

        Args:
          inputs: A batch of intermediate representations of image frames
            across all timesteps, of shape [..., batch_size, timesteps,
            hidden_size].
          samples: Number of samples to draw from the latent distribution.
          static_sample: A tensor sample of the static latent variable `f`
            of shape [..., batch_size, latent_size]. Only used
            for the full dynamic posterior formulation.

        Returns:
          A tuple of a sample tensor of shape [samples, batch_size, length
          latent_size], and a MultivariateNormalDiag distribution from which
          the tensor was sampled, with event shape [latent_size], and batch
          shape [broadcasted_shape, batch_size, length], where
          `broadcasted_shape` is the broadcasted sampled shape between the
          inputs and static sample.

        Raises:
          ValueError: If the "full" latent posterior formulation is being
            used, yet a static latent sample was not provided.


filters = [<NUM_LIT:64>, <NUM_LIT>, <NUM_LIT>, <NUM_LIT>, <NUM_LIT>]<EOL>kernels = [<NUM_LIT:3>, <NUM_LIT:3>, <NUM_LIT:3>, <NUM_LIT:3>, <NUM_LIT:3>]<EOL>strides = [<NUM_LIT:2>, <NUM_LIT:2>, <NUM_LIT:2>, <NUM_LIT:2>, <NUM_LIT:2>]<EOL>def _untransformed_scale_constraint(t):<EOL><INDENT>return tf.clip_by_value(t, -<NUM_LIT:1000>,<EOL>tf.math.log(kernel_posterior_scale_constraint))<EOL><DEDENT>kernel_posterior_fn = tfp.layers.default_mean_field_normal_fn(<EOL>untransformed_scale_initializer=tf.compat.v1.initializers.random_normal(<EOL>mean=kernel_posterior_scale_mean,<EOL>stddev=kernel_posterior_scale_stddev),<EOL>untransformed_scale_constraint=_untransformed_scale_constraint)<EOL>image = tf.keras.layers.Input(shape=input_shape, dtype='<STR_LIT>')<EOL>x = image<EOL>for i in range(len(kernels)):<EOL><INDENT>x = _vggconv_block(<EOL>x,<EOL>filters[i],<EOL>kernels[i],<EOL>strides[i],<EOL>kernel_posterior_fn)<EOL><DEDENT>x = tf.keras.layers.Flatten()(x)<EOL>x = tfp.layers.DenseFlipout(<EOL>num_classes,<EOL>kernel_posterior_fn=kernel_posterior_fn)(x)<EOL>model = tf.keras.Model(inputs=image, outputs=x, name='<STR_LIT>')<EOL>return model<EOL>

Constructs a VGG16 model.

    Args:
      input_shape: A `tuple` indicating the Tensor shape.
      num_classes: `int` representing the number of class labels.
      kernel_posterior_scale_mean: Python `int` number for the kernel
        posterior's scale (log variance) mean. The smaller the mean the closer
        is the initialization to a deterministic network.
      kernel_posterior_scale_stddev: Python `float` number for the initial kernel
        posterior's scale stddev.
        ```
        q(W|x) ~ N(mu, var),
        log_var ~ N(kernel_posterior_scale_mean, kernel_posterior_scale_stddev)
        ````
      kernel_posterior_scale_constraint: Python `float` number for the log value
        to constrain the log variance throughout training.
        i.e. log_var <= log(kernel_posterior_scale_constraint).

    Returns:
      tf.keras.Model.


def _vggconv_block(x, filters, kernel, stride, kernel_posterior_fn):

out = tfp.layers.Convolution2DFlipout(<EOL>filters,<EOL>kernel,<EOL>padding='<STR_LIT>',<EOL>kernel_posterior_fn=kernel_posterior_fn)(x)<EOL>out = tf.keras.layers.BatchNormalization()(out)<EOL>out = tf.keras.layers.Activation('<STR_LIT:relu>')(out)<EOL>out = tfp.layers.Convolution2DFlipout(<EOL>filters,<EOL>kernel,<EOL>padding='<STR_LIT>',<EOL>kernel_posterior_fn=kernel_posterior_fn)(out)<EOL>out = tf.keras.layers.BatchNormalization()(out)<EOL>out = tf.keras.layers.Activation('<STR_LIT:relu>')(out)<EOL>out = tf.keras.layers.MaxPooling2D(<EOL>pool_size=(<NUM_LIT:2>, <NUM_LIT:2>), strides=stride)(out)<EOL>return out<EOL>


filters = [<NUM_LIT:64>, <NUM_LIT>, <NUM_LIT>, <NUM_LIT>]<EOL>kernels = [<NUM_LIT:3>, <NUM_LIT:3>, <NUM_LIT:3>, <NUM_LIT:3>]<EOL>strides = [<NUM_LIT:1>, <NUM_LIT:2>, <NUM_LIT:2>, <NUM_LIT:2>]<EOL>def _untransformed_scale_constraint(t):<EOL><INDENT>return tf.clip_by_value(t, -<NUM_LIT:1000>,<EOL>tf.math.log(kernel_posterior_scale_constraint))<EOL><DEDENT>kernel_posterior_fn = tfp.layers.default_mean_field_normal_fn(<EOL>untransformed_scale_initializer=tf.compat.v1.initializers.random_normal(<EOL>mean=kernel_posterior_scale_mean,<EOL>stddev=kernel_posterior_scale_stddev),<EOL>untransformed_scale_constraint=_untransformed_scale_constraint)<EOL>image = tf.keras.layers.Input(shape=input_shape, dtype='<STR_LIT>')<EOL>x = tfp.layers.Convolution2DFlipout(<EOL><NUM_LIT:64>,<EOL><NUM_LIT:3>,<EOL>strides=<NUM_LIT:1>,<EOL>padding='<STR_LIT>',<EOL>kernel_posterior_fn=kernel_posterior_fn)(image)<EOL>for i in range(len(kernels)):<EOL><INDENT>x = _resnet_block(<EOL>x,<EOL>filters[i],<EOL>kernels[i],<EOL>strides[i],<EOL>kernel_posterior_fn)<EOL><DEDENT>x = tf.keras.layers.BatchNormalization()(x)<EOL>x = tf.keras.layers.Activation('<STR_LIT:relu>')(x)<EOL>x = tf.keras.layers.AveragePooling2D(<NUM_LIT:4>, <NUM_LIT:1>)(x)<EOL>x = tf.keras.layers.Flatten()(x)<EOL>x = tfp.layers.DenseFlipout(<EOL>num_classes,<EOL>kernel_posterior_fn=kernel_posterior_fn)(x)<EOL>model = tf.keras.Model(inputs=image, outputs=x, name='<STR_LIT>')<EOL>return model<EOL>

Constructs a ResNet18 model.

    Args:
      input_shape: A `tuple` indicating the Tensor shape.
      num_classes: `int` representing the number of class labels.
      kernel_posterior_scale_mean: Python `int` number for the kernel
        posterior's scale (log variance) mean. The smaller the mean the closer
        is the initialization to a deterministic network.
      kernel_posterior_scale_stddev: Python `float` number for the initial kernel
        posterior's scale stddev.
        ```
        q(W|x) ~ N(mu, var),
        log_var ~ N(kernel_posterior_scale_mean, kernel_posterior_scale_stddev)
        ````
      kernel_posterior_scale_constraint: Python `float` number for the log value
        to constrain the log variance throughout training.
        i.e. log_var <= log(kernel_posterior_scale_constraint).

    Returns:
      tf.keras.Model.


def _resnet_block(x, filters, kernel, stride, kernel_posterior_fn):

x = tf.keras.layers.BatchNormalization()(x)<EOL>x = tf.keras.layers.Activation('<STR_LIT:relu>')(x)<EOL>if stride != <NUM_LIT:1> or filters != x.shape[<NUM_LIT:1>]:<EOL><INDENT>shortcut = _projection_shortcut(x, filters, stride, kernel_posterior_fn)<EOL><DEDENT>else:<EOL><INDENT>shortcut = x<EOL><DEDENT>x = tfp.layers.Convolution2DFlipout(<EOL>filters,<EOL>kernel,<EOL>strides=stride,<EOL>padding='<STR_LIT>',<EOL>kernel_posterior_fn=kernel_posterior_fn)(x)<EOL>x = tf.keras.layers.BatchNormalization()(x)<EOL>x = tf.keras.layers.Activation('<STR_LIT:relu>')(x)<EOL>x = tfp.layers.Convolution2DFlipout(<EOL>filters,<EOL>kernel,<EOL>strides=<NUM_LIT:1>,<EOL>padding='<STR_LIT>',<EOL>kernel_posterior_fn=kernel_posterior_fn)(x)<EOL>x = tf.keras.layers.add([x, shortcut])<EOL>return x<EOL>


return tf.clip_by_value(x, <NUM_LIT>, <NUM_LIT>)<EOL>

Clips the Dirichlet parameters to the numerically stable KL region.


def make_encoder(activation, num_topics, layer_sizes):

Create the encoder function.

    Args:
      activation: Activation function to use.
      num_topics: The number of topics.
      layer_sizes: The number of hidden units per layer in the encoder.

    Returns:
      encoder: A `callable` mapping a bag-of-words `Tensor` to a
        `tfd.Distribution` instance over topics.


topics_words_logits = tf.compat.v1.get_variable(<EOL>"<STR_LIT>",<EOL>shape=[num_topics, num_words],<EOL>initializer=tf.compat.v1.glorot_normal_initializer())<EOL>topics_words = tf.nn.softmax(topics_words_logits, axis=-<NUM_LIT:1>)<EOL>def decoder(topics):<EOL><INDENT>word_probs = tf.matmul(topics, topics_words)<EOL>return tfd.OneHotCategorical(probs=word_probs,<EOL>name="<STR_LIT>")<EOL><DEDENT>return decoder, topics_words<EOL>

Create the decoder function.

    Args:
      num_topics: The number of topics.
      num_words: The number of words.

    Returns:
      decoder: A `callable` mapping a `Tensor` of encodings to a
        `tfd.Distribution` instance over words.


def _softplus_inverse(x):<EOL><INDENT>return np.log(np.expm1(x))<EOL><DEDENT>logit_concentration = tf.compat.v1.get_variable(<EOL>"<STR_LIT>",<EOL>shape=[<NUM_LIT:1>, num_topics],<EOL>initializer=tf.compat.v1.initializers.constant(<EOL>_softplus_inverse(initial_value)))<EOL>concentration = _clip_dirichlet_parameters(<EOL>tf.nn.softplus(logit_concentration))<EOL>def prior():<EOL><INDENT>return tfd.Dirichlet(concentration=concentration,<EOL>name="<STR_LIT>")<EOL><DEDENT>prior_variables = [logit_concentration]<EOL>return prior, prior_variables<EOL>

Create the prior distribution.

    Args:
      num_topics: Number of topics.
      initial_value: The starting value for the prior parameters.

    Returns:
      prior: A `callable` that returns a `tf.distribution.Distribution`
          instance, the prior distribution.
      prior_variables: A `list` of `Variable` objects, the trainable parameters
          of the prior.


def model_fn(features, labels, mode, params, config):

del labels, config<EOL>encoder = make_encoder(params["<STR_LIT>"],<EOL>params["<STR_LIT>"],<EOL>params["<STR_LIT>"])<EOL>decoder, topics_words = make_decoder(params["<STR_LIT>"],<EOL>features.shape[<NUM_LIT:1>])<EOL>prior, prior_variables = make_prior(params["<STR_LIT>"],<EOL>params["<STR_LIT>"])<EOL>topics_prior = prior()<EOL>alpha = topics_prior.concentration<EOL>topics_posterior = encoder(features)<EOL>topics = topics_posterior.sample()<EOL>random_reconstruction = decoder(topics)<EOL>reconstruction = random_reconstruction.log_prob(features)<EOL>tf.compat.v1.summary.scalar("<STR_LIT>",<EOL>tf.reduce_mean(input_tensor=reconstruction))<EOL>kl = tfd.kl_divergence(topics_posterior, topics_prior)<EOL>tf.compat.v1.summary.scalar("<STR_LIT>", tf.reduce_mean(input_tensor=kl))<EOL>with tf.control_dependencies(<EOL>[tf.compat.v1.assert_greater(kl, -<NUM_LIT>, message="<STR_LIT>")]):<EOL><INDENT>kl = tf.identity(kl)<EOL><DEDENT>elbo = reconstruction - kl<EOL>avg_elbo = tf.reduce_mean(input_tensor=elbo)<EOL>tf.compat.v1.summary.scalar("<STR_LIT>", avg_elbo)<EOL>loss = -avg_elbo<EOL>global_step = tf.compat.v1.train.get_or_create_global_step()<EOL>optimizer = tf.compat.v1.train.AdamOptimizer(params["<STR_LIT>"])<EOL>grads_and_vars = optimizer.compute_gradients(loss)<EOL>grads_and_vars_except_prior = [<EOL>x for x in grads_and_vars if x[<NUM_LIT:1>] not in prior_variables]<EOL>def train_op_except_prior():<EOL><INDENT>return optimizer.apply_gradients(<EOL>grads_and_vars_except_prior,<EOL>global_step=global_step)<EOL><DEDENT>def train_op_all():<EOL><INDENT>return optimizer.apply_gradients(<EOL>grads_and_vars,<EOL>global_step=global_step)<EOL><DEDENT>train_op = tf.cond(<EOL>pred=global_step < params["<STR_LIT>"],<EOL>true_fn=train_op_except_prior,<EOL>false_fn=train_op_all)<EOL>words_per_document = tf.reduce_sum(input_tensor=features, axis=<NUM_LIT:1>)<EOL>log_perplexity = -elbo / words_per_document<EOL>tf.compat.v1.summary.scalar(<EOL>"<STR_LIT>", tf.exp(tf.reduce_mean(input_tensor=log_perplexity)))<EOL>(log_perplexity_tensor,<EOL>log_perplexity_update) = tf.compat.v1.metrics.mean(log_perplexity)<EOL>perplexity_tensor = tf.exp(log_perplexity_tensor)<EOL>topics = tf.compat.v1.py_func(<EOL>functools.partial(get_topics_strings, vocabulary=params["<STR_LIT>"]),<EOL>[topics_words, alpha],<EOL>tf.string,<EOL>stateful=False)<EOL>tf.compat.v1.summary.text("<STR_LIT>", topics)<EOL>return tf.estimator.EstimatorSpec(<EOL>mode=mode,<EOL>loss=loss,<EOL>train_op=train_op,<EOL>eval_metric_ops={<EOL>"<STR_LIT>": tf.compat.v1.metrics.mean(elbo),<EOL>"<STR_LIT>": tf.compat.v1.metrics.mean(reconstruction),<EOL>"<STR_LIT>": tf.compat.v1.metrics.mean(kl),<EOL>"<STR_LIT>": (perplexity_tensor, log_perplexity_update),<EOL>"<STR_LIT>": (topics, tf.no_op()),<EOL>},<EOL>)<EOL>

Build the model function for use in an estimator.

    Arguments:
      features: The input features for the estimator.
      labels: The labels, unused here.
      mode: Signifies whether it is train or test or predict.
      params: Some hyperparameters as a dictionary.
      config: The RunConfig, unused here.
    Returns:
      EstimatorSpec: A tf.estimator.EstimatorSpec instance.


def get_topics_strings(topics_words, alpha, vocabulary,<EOL>topics_to_print=<NUM_LIT:10>, words_per_topic=<NUM_LIT:10>):

alpha = np.squeeze(alpha, axis=<NUM_LIT:0>)<EOL>highest_weight_topics = np.argsort(-alpha, kind="<STR_LIT>")<EOL>top_words = np.argsort(-topics_words, axis=<NUM_LIT:1>)<EOL>res = []<EOL>for topic_idx in highest_weight_topics[:topics_to_print]:<EOL><INDENT>l = ["<STR_LIT>".format(topic_idx, alpha[topic_idx])]<EOL>l += [vocabulary[word] for word in top_words[topic_idx, :words_per_topic]]<EOL>res.append("<STR_LIT:U+0020>".join(l))<EOL><DEDENT>return np.array(res)<EOL>

Returns the summary of the learned topics.

    Arguments:
      topics_words: KxV tensor with topics as rows and words as columns.
      alpha: 1xK tensor of prior Dirichlet concentrations for the
          topics.
      vocabulary: A mapping of word's integer index to the corresponding string.
      topics_to_print: The number of topics with highest prior weight to
          summarize.
      words_per_topic: Number of wodrs per topic to return.
    Returns:
      summary: A np.array with strings.


filepath = os.path.join(directory, filename)<EOL>if tf.io.gfile.exists(filepath):<EOL><INDENT>return filepath<EOL><DEDENT>if not tf.io.gfile.exists(directory):<EOL><INDENT>tf.io.gfile.makedirs(directory)<EOL><DEDENT>url = os.path.join(ROOT_PATH, filename)<EOL>print("<STR_LIT>" % (url, filepath))<EOL>urllib.request.urlretrieve(url, filepath)<EOL>return filepath<EOL>


def newsgroups_dataset(directory, split_name, num_words, shuffle_and_repeat):

data = np.load(download(directory, FILE_TEMPLATE.format(split=split_name)))<EOL>data = data[:-<NUM_LIT:1>]<EOL>num_documents = data.shape[<NUM_LIT:0>]<EOL>indices = np.array([(row_idx, column_idx)<EOL>for row_idx, row in enumerate(data)<EOL>for column_idx in row])<EOL>sparse_matrix = scipy.sparse.coo_matrix(<EOL>(np.ones(indices.shape[<NUM_LIT:0>]), (indices[:, <NUM_LIT:0>], indices[:, <NUM_LIT:1>])),<EOL>shape=(num_documents, num_words),<EOL>dtype=np.float32)<EOL>sparse_matrix = sparse_matrix.tocsr()<EOL>dataset = tf.data.Dataset.range(num_documents)<EOL>if shuffle_and_repeat:<EOL><INDENT>dataset = dataset.shuffle(num_documents).repeat()<EOL><DEDENT>def get_row_py_func(idx):<EOL><INDENT>def get_row_python(idx_py):<EOL><INDENT>return np.squeeze(np.array(sparse_matrix[idx_py].todense()), axis=<NUM_LIT:0>)<EOL><DEDENT>py_func = tf.compat.v1.py_func(<EOL>get_row_python, [idx], tf.float32, stateful=False)<EOL>py_func.set_shape((num_words,))<EOL>return py_func<EOL><DEDENT>dataset = dataset.map(get_row_py_func)<EOL>return dataset<EOL>


num_words = <NUM_LIT:1000><EOL>vocabulary = [str(i) for i in range(num_words)]<EOL>random_sample = np.random.randint(<EOL><NUM_LIT:10>, size=(batch_size, num_words)).astype(np.float32)<EOL>def train_input_fn():<EOL><INDENT>dataset = tf.data.Dataset.from_tensor_slices(random_sample)<EOL>dataset = dataset.batch(batch_size)<EOL>return tf.compat.v1.data.make_one_shot_iterator(dataset.repeat()).get_next()<EOL><DEDENT>def eval_input_fn():<EOL><INDENT>dataset = tf.data.Dataset.from_tensor_slices(random_sample)<EOL>dataset = dataset.batch(batch_size)<EOL>return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()<EOL><DEDENT>return train_input_fn, eval_input_fn, vocabulary<EOL>


with open(download(data_dir, "<STR_LIT>"), "<STR_LIT:r>") as f:<EOL><INDENT>words_to_idx = pickle.load(f)<EOL><DEDENT>num_words = len(words_to_idx)<EOL>vocabulary = [None] * num_words<EOL>for word, idx in words_to_idx.items():<EOL><INDENT>vocabulary[idx] = word<EOL><DEDENT>def train_input_fn():<EOL><INDENT>dataset = newsgroups_dataset(<EOL>data_dir, "<STR_LIT:train>", num_words, shuffle_and_repeat=True)<EOL>dataset = dataset.batch(batch_size).prefetch(<NUM_LIT:32>)<EOL>return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()<EOL><DEDENT>def eval_input_fn():<EOL><INDENT>dataset = newsgroups_dataset(<EOL>data_dir, "<STR_LIT:test>", num_words, shuffle_and_repeat=False)<EOL>dataset = dataset.batch(batch_size)<EOL>return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()<EOL><DEDENT>return train_input_fn, eval_input_fn, vocabulary<EOL>

Builds iterators for train and evaluation data.

    Each object is represented as a bag-of-words vector.

    Arguments:
      data_dir: Folder in which to store the data.
      batch_size: Batch size for both train and evaluation.
    Returns:
      train_input_fn: A function that returns an iterator over the training data.
      eval_input_fn: A function that returns an iterator over the evaluation data.
      vocabulary: A mapping of word's integer index to the corresponding string.


return np.log(np.expm1(x))<EOL>


topics = ed.Dirichlet(concentration=concentration, name="<STR_LIT>")<EOL>word_probs = tf.matmul(topics, topics_words)<EOL>bag_of_words = ed.OneHotCategorical(probs=word_probs, name="<STR_LIT>")<EOL>return bag_of_words<EOL>

Latent Dirichlet Allocation in terms of its generative process.

    The model posits a distribution over bags of words and is parameterized by
    a concentration and the topic-word probabilities. It collapses per-word
    topic assignments.

    Args:
      concentration: A Tensor of shape [1, num_topics], which parameterizes the
        Dirichlet prior over topics.
      topics_words: A Tensor of shape [num_topics, num_words], where each row
        (topic) denotes the probability of each word being in that topic.

    Returns:
      bag_of_words: A random variable capturing a sample from the model, of shape
        [1, num_words]. It represents one generated document as a bag of words.


def make_lda_variational(activation, num_topics, layer_sizes):

encoder_net = tf.keras.Sequential()<EOL>for num_hidden_units in layer_sizes:<EOL><INDENT>encoder_net.add(<EOL>tf.keras.layers.Dense(<EOL>num_hidden_units,<EOL>activation=activation,<EOL>kernel_initializer=tf.compat.v1.glorot_normal_initializer()))<EOL><DEDENT>encoder_net.add(<EOL>tf.keras.layers.Dense(<EOL>num_topics,<EOL>activation=tf.nn.softplus,<EOL>kernel_initializer=tf.compat.v1.glorot_normal_initializer()))<EOL>def lda_variational(bag_of_words):<EOL><INDENT>concentration = _clip_dirichlet_parameters(encoder_net(bag_of_words))<EOL>return ed.Dirichlet(concentration=concentration, name="<STR_LIT>")<EOL><DEDENT>return lda_variational<EOL>

Creates the variational distribution for LDA.

    Args:
      activation: Activation function to use.
      num_topics: The number of topics.
      layer_sizes: The number of hidden units per layer in the encoder.

    Returns:
      lda_variational: A function that takes a bag-of-words Tensor as
        input and returns a distribution over topics.


def set_values(f, *args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>name = kwargs.get("<STR_LIT:name>")<EOL>if name in model_kwargs:<EOL><INDENT>kwargs["<STR_LIT:value>"] = model_kwargs[name]<EOL><DEDENT>return ed.interceptable(f)(*args, **kwargs)<EOL><DEDENT>return set_values<EOL>

Creates a value-setting interceptor.

    Args:
      **model_kwargs: dict of str to Tensor. Keys are the names of random variable
        in the model to which this interceptor is being applied. Values are
        Tensors to set their value to.

    Returns:
      set_values: Function which sets the value of intercepted ops.


del labels, config<EOL>logit_concentration = tf.compat.v1.get_variable(<EOL>"<STR_LIT>",<EOL>shape=[<NUM_LIT:1>, params["<STR_LIT>"]],<EOL>initializer=tf.compat.v1.initializers.constant(<EOL>_softplus_inverse(params["<STR_LIT>"])))<EOL>concentration = _clip_dirichlet_parameters(<EOL>tf.nn.softplus(logit_concentration))<EOL>num_words = features.shape[<NUM_LIT:1>]<EOL>topics_words_logits = tf.compat.v1.get_variable(<EOL>"<STR_LIT>",<EOL>shape=[params["<STR_LIT>"], num_words],<EOL>initializer=tf.compat.v1.glorot_normal_initializer())<EOL>topics_words = tf.nn.softmax(topics_words_logits, axis=-<NUM_LIT:1>)<EOL>lda_variational = make_lda_variational(<EOL>params["<STR_LIT>"],<EOL>params["<STR_LIT>"],<EOL>params["<STR_LIT>"])<EOL>with ed.tape() as variational_tape:<EOL><INDENT>_ = lda_variational(features)<EOL><DEDENT>with ed.tape() as model_tape:<EOL><INDENT>with ed.interception(<EOL>make_value_setter(topics=variational_tape["<STR_LIT>"])):<EOL><INDENT>posterior_predictive = latent_dirichlet_allocation(concentration,<EOL>topics_words)<EOL><DEDENT><DEDENT>log_likelihood = posterior_predictive.distribution.log_prob(features)<EOL>tf.compat.v1.summary.scalar("<STR_LIT>",<EOL>tf.reduce_mean(input_tensor=log_likelihood))<EOL>kl = variational_tape["<STR_LIT>"].distribution.kl_divergence(<EOL>model_tape["<STR_LIT>"].distribution)<EOL>tf.compat.v1.summary.scalar("<STR_LIT>", tf.reduce_mean(input_tensor=kl))<EOL>with tf.control_dependencies(<EOL>[tf.compat.v1.assert_greater(kl, -<NUM_LIT>, message="<STR_LIT>")]):<EOL><INDENT>kl = tf.identity(kl)<EOL><DEDENT>elbo = log_likelihood - kl<EOL>avg_elbo = tf.reduce_mean(input_tensor=elbo)<EOL>tf.compat.v1.summary.scalar("<STR_LIT>", avg_elbo)<EOL>loss = -avg_elbo<EOL>global_step = tf.compat.v1.train.get_or_create_global_step()<EOL>optimizer = tf.compat.v1.train.AdamOptimizer(params["<STR_LIT>"])<EOL>grads_and_vars = optimizer.compute_gradients(loss)<EOL>grads_and_vars_except_prior = [<EOL>x for x in grads_and_vars if x[<NUM_LIT:1>] != logit_concentration]<EOL>def train_op_except_prior():<EOL><INDENT>return optimizer.apply_gradients(<EOL>grads_and_vars_except_prior,<EOL>global_step=global_step)<EOL><DEDENT>def train_op_all():<EOL><INDENT>return optimizer.apply_gradients(<EOL>grads_and_vars,<EOL>global_step=global_step)<EOL><DEDENT>train_op = tf.cond(<EOL>pred=global_step < params["<STR_LIT>"],<EOL>true_fn=train_op_except_prior,<EOL>false_fn=train_op_all)<EOL>words_per_document = tf.reduce_sum(input_tensor=features, axis=<NUM_LIT:1>)<EOL>log_perplexity = -elbo / words_per_document<EOL>tf.compat.v1.summary.scalar(<EOL>"<STR_LIT>", tf.exp(tf.reduce_mean(input_tensor=log_perplexity)))<EOL>(log_perplexity_tensor,<EOL>log_perplexity_update) = tf.compat.v1.metrics.mean(log_perplexity)<EOL>perplexity_tensor = tf.exp(log_perplexity_tensor)<EOL>topics = tf.compat.v1.py_func(<EOL>functools.partial(get_topics_strings, vocabulary=params["<STR_LIT>"]),<EOL>[topics_words, concentration],<EOL>tf.string,<EOL>stateful=False)<EOL>tf.compat.v1.summary.text("<STR_LIT>", topics)<EOL>return tf.estimator.EstimatorSpec(<EOL>mode=mode,<EOL>loss=loss,<EOL>train_op=train_op,<EOL>eval_metric_ops={<EOL>"<STR_LIT>": tf.compat.v1.metrics.mean(elbo),<EOL>"<STR_LIT>": tf.compat.v1.metrics.mean(log_likelihood),<EOL>"<STR_LIT>": tf.compat.v1.metrics.mean(kl),<EOL>"<STR_LIT>": (perplexity_tensor, log_perplexity_update),<EOL>"<STR_LIT>": (topics, tf.no_op()),<EOL>},<EOL>)<EOL>

Builds the model function for use in an Estimator.

    Arguments:
      features: The input features for the Estimator.
      labels: The labels, unused here.
      mode: Signifies whether it is train or test or predict.
      params: Some hyperparameters as a dictionary.
      config: The RunConfig, unused here.

    Returns:
      EstimatorSpec: A tf.estimator.EstimatorSpec instance.


Returns the summary of the learned topics.

    Arguments:
      topics_words: KxV tensor with topics as rows and words as columns.
      alpha: 1xK tensor of prior Dirichlet concentrations for the
          topics.
      vocabulary: A mapping of word's integer index to the corresponding string.
      topics_to_print: The number of topics with highest prior weight to
          summarize.
      words_per_topic: Number of wodrs per topic to return.

    Returns:
      summary: A np.array with strings.


num_words = <NUM_LIT:1000><EOL>vocabulary = [str(i) for i in range(num_words)]<EOL>random_sample = np.random.randint(<EOL><NUM_LIT:10>, size=(batch_size, num_words)).astype(np.float32)<EOL>def train_input_fn():<EOL><INDENT>dataset = tf.data.Dataset.from_tensor_slices(random_sample)<EOL>dataset = dataset.batch(batch_size).repeat()<EOL>return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()<EOL><DEDENT>def eval_input_fn():<EOL><INDENT>dataset = tf.data.Dataset.from_tensor_slices(random_sample)<EOL>dataset = dataset.batch(batch_size)<EOL>return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()<EOL><DEDENT>return train_input_fn, eval_input_fn, vocabulary<EOL>


Builds iterators for train and evaluation data.

    Each object is represented as a bag-of-words vector.

    Arguments:
      data_dir: Folder in which to store the data.
      batch_size: Batch size for both train and evaluation.

    Returns:
      train_input_fn: A function that returns an iterator over the training data.
      eval_input_fn: A function that returns an iterator over the evaluation data.
      vocabulary: A mapping of word's integer index to the corresponding string.


def toy_logistic_data(num_examples, input_size=<NUM_LIT:2>, weights_prior_stddev=<NUM_LIT>):

random_weights = weights_prior_stddev * np.random.randn(input_size)<EOL>random_bias = np.random.randn()<EOL>design_matrix = np.random.rand(num_examples, input_size) * <NUM_LIT:2> - <NUM_LIT:1><EOL>logits = np.reshape(<EOL>np.dot(design_matrix, random_weights) + random_bias,<EOL>(-<NUM_LIT:1>, <NUM_LIT:1>))<EOL>p_labels = <NUM_LIT:1.> / (<NUM_LIT:1> + np.exp(-logits))<EOL>labels = np.int32(p_labels > np.random.rand(num_examples, <NUM_LIT:1>))<EOL>return random_weights, random_bias, np.float32(design_matrix), labels<EOL>

Generates synthetic data for binary classification.

    Args:
      num_examples: The number of samples to generate (scalar Python `int`).
      input_size: The input space dimension (scalar Python `int`).
      weights_prior_stddev: The prior standard deviation of the weight
        vector. (scalar Python `float`).

    Returns:
      random_weights: Sampled weights as a Numpy `array` of shape
        `[input_size]`.
      random_bias: Sampled bias as a scalar Python `float`.
      design_matrix: Points sampled uniformly from the cube `[-1,
         1]^{input_size}`, as a Numpy `array` of shape `(num_examples,
         input_size)`.
      labels: Labels sampled from the logistic model `p(label=1) =
        logistic(dot(features, random_weights) + random_bias)`, as a Numpy
        `int32` `array` of shape `(num_examples, 1)`.


def visualize_decision(features, labels, true_w_b, candidate_w_bs, fname):

fig = figure.Figure(figsize=(<NUM_LIT:6>, <NUM_LIT:6>))<EOL>canvas = backend_agg.FigureCanvasAgg(fig)<EOL>ax = fig.add_subplot(<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:1>)<EOL>ax.scatter(features[:, <NUM_LIT:0>], features[:, <NUM_LIT:1>],<EOL>c=np.float32(labels[:, <NUM_LIT:0>]),<EOL>cmap=cm.get_cmap("<STR_LIT>"),<EOL>edgecolors="<STR_LIT:k>")<EOL>def plot_weights(w, b, **kwargs):<EOL><INDENT>w1, w2 = w<EOL>x1s = np.linspace(-<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:100>)<EOL>x2s = -(w1  * x1s + b) / w2<EOL>ax.plot(x1s, x2s, **kwargs)<EOL><DEDENT>for w, b in candidate_w_bs:<EOL><INDENT>plot_weights(w, b,<EOL>alpha=<NUM_LIT:1.>/np.sqrt(len(candidate_w_bs)),<EOL>lw=<NUM_LIT:1>, color="<STR_LIT>")<EOL><DEDENT>if true_w_b is not None:<EOL><INDENT>plot_weights(*true_w_b, lw=<NUM_LIT:4>,<EOL>color="<STR_LIT>", label="<STR_LIT>")<EOL><DEDENT>ax.set_xlim([-<NUM_LIT>, <NUM_LIT>])<EOL>ax.set_ylim([-<NUM_LIT>, <NUM_LIT>])<EOL>ax.legend()<EOL>canvas.print_figure(fname, format="<STR_LIT>")<EOL>print("<STR_LIT>".format(fname))<EOL>

Utility method to visualize decision boundaries in R^2.

    Args:
      features: Input points, as a Numpy `array` of shape `[num_examples, 2]`.
      labels: Numpy `float`-like array of shape `[num_examples, 1]` giving a
        label for each point.
      true_w_b: A `tuple` `(w, b)` where `w` is a Numpy array of
         shape `[2]` and `b` is a scalar `float`, interpreted as a
         decision rule of the form `dot(features, w) + b > 0`.
      candidate_w_bs: Python `iterable` containing tuples of the same form as
         true_w_b.
      fname: The filename to save the plot as a PNG image (Python `str`).


def build_input_pipeline(x, y, batch_size):

training_dataset = tf.data.Dataset.from_tensor_slices((x, y))<EOL>training_batches = training_dataset.repeat().batch(batch_size)<EOL>training_iterator = tf.compat.v1.data.make_one_shot_iterator(training_batches)<EOL>batch_features, batch_labels = training_iterator.get_next()<EOL>return batch_features, batch_labels<EOL>

Build a Dataset iterator for supervised classification.

    Args:
      x: Numpy `array` of features, indexed by the first dimension.
      y: Numpy `array` of labels, with the same first dimension as `x`.
      batch_size: Number of elements in each training batch.

    Returns:
      batch_features: `Tensor` feed  features, of shape
        `[batch_size] + x.shape[1:]`.
      batch_labels: `Tensor` feed of labels, of shape
        `[batch_size] + y.shape[1:]`.


num_examples = <NUM_LIT:10><EOL>x_train = np.random.rand(num_examples, *IMAGE_SHAPE).astype(np.float32)<EOL>y_train = np.random.permutation(np.arange(num_examples)).astype(np.int32)<EOL>x_test = np.random.rand(num_examples, *IMAGE_SHAPE).astype(np.float32)<EOL>y_test = np.random.permutation(np.arange(num_examples)).astype(np.int32)<EOL>return (x_train, y_train), (x_test, y_test)<EOL>

Build fake CIFAR10-style data for unit testing.


symbols = []<EOL>for production in tf.unstack(productions, axis=<NUM_LIT:1>):<EOL><INDENT>lhs, rhs = self.production_rules[tf.argmax(input=production, axis=-<NUM_LIT:1>)]<EOL>if not symbols:  <EOL><INDENT>if lhs != self.start_symbol:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>symbols = rhs<EOL><DEDENT>else:<EOL><INDENT>index = symbols.index(lhs)<EOL>symbols = symbols[:index] + rhs + symbols[index + <NUM_LIT:1>:]<EOL><DEDENT><DEDENT>string = "<STR_LIT>".join(symbols)<EOL>return string<EOL>

Converts a sequence of productions into a string of terminal symbols.

        Args:
          productions: Tensor of shape [1, num_productions, num_production_rules].
            Slices along the `num_productions` dimension represent one-hot vectors.

        Returns:
          str that concatenates all terminal symbols from `productions`.

        Raises:
          ValueError: If the first production rule does not begin with
            `self.start_symbol`.


def mask(self, symbol, on_value, off_value):

mask_values = [on_value if lhs == symbol else off_value<EOL>for lhs, _ in self.production_rules]<EOL>mask_values = tf.reshape(mask_values, [<NUM_LIT:1>, len(self.production_rules)])<EOL>return mask_values<EOL>

Produces a masking tensor for (in)valid production rules.

        Args:
          symbol: str, a symbol in the grammar.
          on_value: Value to use for a valid production rule.
          off_value: Value to use for an invalid production rule.

        Returns:
          Tensor of shape [1, num_production_rules]. An element is `on_value`
          if its corresponding production rule has `symbol` on its left-hand-side;
          the element is `off_value` otherwise.


def __init__(self, grammar, latent_size, num_units):

super(ProbabilisticGrammar, self).__init__()<EOL>self.grammar = grammar<EOL>self.latent_size = latent_size<EOL>self.lstm = tf.compat.v1.nn.rnn_cell.LSTMCell(num_units)<EOL>self.output_layer = tf.keras.layers.Dense(len(grammar.production_rules))<EOL>

Constructs a probabilistic grammar.

        Args:
          grammar: An object representing a grammar. It has members
            `nonterminal_symbols`, `alphabet`, `production_rules`, and
            `start_symbol`, and a method `mask` determining (in)valid
            production rules given a symbol.
          latent_size: Number of dimensions in the latent code.
          num_units: Number of units in the LSTM cell.