@kullback_leibler.RegisterKL(<EOL>mvn_linear_operator.MultivariateNormalLinearOperator, GaussianProcess)<EOL>def _kl_mvn_gp(mvn, gp, name=None):

return _kl_compatible_gp(mvn, gp, name or '<STR_LIT>')<EOL>

Calculate the batched KL divergence KL(mvn || gp).

    Args:
      mvn: instance of a multivariate Normal distribution object (any subclass of
        MultivariateNormalLinearOperator)
      gp: instance of a GaussianProcess distribution object.
      name: (optional) Name to use for created operations.
        default is 'kl_mvn_gp'.

    Returns:
      Batchwise KL(mvn || gp)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[index_points, observation_noise_variance, jitter], tf.float32)<EOL>if index_points is not None:<EOL><INDENT>index_points = tf.convert_to_tensor(<EOL>value=index_points, dtype=dtype, name='<STR_LIT>')<EOL><DEDENT>jitter = tf.convert_to_tensor(value=jitter, dtype=dtype, name='<STR_LIT>')<EOL>observation_noise_variance = tf.convert_to_tensor(<EOL>value=observation_noise_variance,<EOL>dtype=dtype,<EOL>name='<STR_LIT>')<EOL>self._kernel = kernel<EOL>self._index_points = index_points<EOL>if mean_fn is None:<EOL><INDENT>mean_fn = lambda x: tf.zeros([<NUM_LIT:1>], dtype=dtype)<EOL><DEDENT>else:<EOL><INDENT>if not callable(mean_fn):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT><DEDENT>self._mean_fn = mean_fn<EOL>self._observation_noise_variance = observation_noise_variance<EOL>self._jitter = jitter<EOL>graph_parents = [observation_noise_variance, jitter]<EOL>if index_points is not None: graph_parents.append(index_points)<EOL>with tf.name_scope('<STR_LIT>'):<EOL><INDENT>super(GaussianProcess, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=graph_parents,<EOL>name=name)<EOL><DEDENT><DEDENT>

Instantiate a GaussianProcess Distribution.

        Args:
          kernel: `PositiveSemidefiniteKernel`-like instance representing the
            GP's covariance function.
          index_points: `float` `Tensor` representing finite (batch of) vector(s) of
            points in the index set over which the GP is defined. Shape has the form
            `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
            dimensions and must equal `kernel.feature_ndims` and `e` is the number
            (size) of index points in each batch. Ultimately this distribution
            corresponds to a `e`-dimensional multivariate normal. The batch shape
            must be broadcastable with `kernel.batch_shape` and any batch dims
            yielded by `mean_fn`.
          mean_fn: Python `callable` that acts on `index_points` to produce a (batch
            of) vector(s) of mean values at `index_points`. Takes a `Tensor` of
            shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
            broadcastable with `[b1, ..., bB]`. Default value: `None` implies
            constant zero function.
          observation_noise_variance: `float` `Tensor` representing the variance
            of the noise in the Normal likelihood distribution of the model. May be
            batched, in which case the batch shape must be broadcastable with the
            shapes of all other batched parameters (`kernel.batch_shape`,
            `index_points`, etc.).
            Default value: `0.`
          jitter: `float` scalar `Tensor` added to the diagonal of the covariance
            matrix to ensure positive definiteness of the covariance matrix.
            Default value: `1e-6`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          allow_nan_stats: Python `bool`, default `True`. When `True`,
            statistics (e.g., mean, mode, variance) use the value "`NaN`" to
            indicate the result is undefined. When `False`, an exception is raised
            if one or more of the statistic's batch members are undefined.
            Default value: `False`.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: "GaussianProcess".

        Raises:
          ValueError: if `mean_fn` is not `None` and is not callable.


num_index_points = tf.compat.dimension_value(<EOL>index_points.shape[-(self.kernel.feature_ndims + <NUM_LIT:1>)])<EOL>if num_index_points is None:<EOL><INDENT>warnings.warn(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>return num_index_points == <NUM_LIT:1><EOL>

True if the given index_points would yield a univariate marginal.

        Args:
          index_points: the set of index set locations at which to compute the
          marginal Gaussian distribution. If this set is of size 1, the marginal is
          univariate.

        Returns:
          is_univariate: Boolean indicating whether the marginal is univariate or
          multivariate. In the case of dynamic shape in the number of index points,
          defaults to "multivariate" since that's the best we can do.


with self._name_scope('<STR_LIT>'):<EOL><INDENT>index_points = self._get_index_points(index_points)<EOL>covariance = self._compute_covariance(index_points)<EOL>loc = self._mean_fn(index_points)<EOL>if self._is_univariate_marginal(index_points):<EOL><INDENT>scale = tf.sqrt(covariance)<EOL>loc = tf.squeeze(loc, axis=-<NUM_LIT:1>)<EOL>return normal.Normal(<EOL>loc=loc,<EOL>scale=scale,<EOL>validate_args=self._validate_args,<EOL>allow_nan_stats=self._allow_nan_stats,<EOL>name='<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>scale = tf.linalg.LinearOperatorLowerTriangular(<EOL>tf.linalg.cholesky(_add_diagonal_shift(covariance, self.jitter)),<EOL>is_non_singular=True,<EOL>name='<STR_LIT>')<EOL>return mvn_linear_operator.MultivariateNormalLinearOperator(<EOL>loc=loc,<EOL>scale=scale,<EOL>validate_args=self._validate_args,<EOL>allow_nan_stats=self._allow_nan_stats,<EOL>name='<STR_LIT>')<EOL><DEDENT><DEDENT>

Compute the marginal of this GP over function values at `index_points`.

        Args:
          index_points: `float` `Tensor` representing finite (batch of) vector(s) of
            points in the index set over which the GP is defined. Shape has the form
            `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
            dimensions and must equal `kernel.feature_ndims` and `e` is the number
            (size) of index points in each batch. Ultimately this distribution
            corresponds to a `e`-dimensional multivariate normal. The batch shape
            must be broadcastable with `kernel.batch_shape` and any batch dims
            yielded by `mean_fn`.

        Returns:
          marginal: a `Normal` or `MultivariateNormalLinearOperator` distribution,
            according to whether `index_points` consists of one or many index
            points, respectively.


if self._index_points is None and index_points is None:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>return index_points if index_points is not None else self._index_points<EOL>

Return `index_points` if not None, else `self._index_points`.

        Args:
          index_points: if given, this is what is returned; else,
          `self._index_points`

        Returns:
          index_points: the given arg, if not None, else the class member
          `self._index_points`.

        Rases:
          ValueError: if `index_points` and `self._index_points` are both `None`.


parameters = dict(locals())<EOL>def _convert_to_tensor(x, name, dtype):<EOL><INDENT>return None if x is None else tf.convert_to_tensor(<EOL>value=x, name=name, dtype=dtype)<EOL><DEDENT>if loc is None and scale_tril is None:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale_tril], tf.float32)<EOL>loc = _convert_to_tensor(loc, name="<STR_LIT>", dtype=dtype)<EOL>scale_tril = _convert_to_tensor(<EOL>scale_tril, name="<STR_LIT>", dtype=dtype)<EOL>if scale_tril is None:<EOL><INDENT>scale = tf.linalg.LinearOperatorIdentity(<EOL>num_rows=distribution_util.dimension_size(loc, -<NUM_LIT:1>),<EOL>dtype=loc.dtype,<EOL>is_self_adjoint=True,<EOL>is_positive_definite=True,<EOL>assert_proper_shapes=validate_args)<EOL><DEDENT>else:<EOL><INDENT>scale = tf.linalg.LinearOperatorLowerTriangular(<EOL>scale_tril,<EOL>is_non_singular=True,<EOL>is_self_adjoint=False,<EOL>is_positive_definite=False)<EOL><DEDENT><DEDENT><DEDENT>super(MultivariateNormalTriL, self).__init__(<EOL>loc=loc,<EOL>scale=scale,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL>self._parameters = parameters<EOL>

Construct Multivariate Normal distribution on `R^k`.

        The `batch_shape` is the broadcast shape between `loc` and `scale`
        arguments.

        The `event_shape` is given by last dimension of the matrix implied by
        `scale`. The last dimension of `loc` (if provided) must broadcast with this.

        Recall that `covariance = scale @ scale.T`. A (non-batch) `scale` matrix is:

        ```none
        scale = scale_tril
        ```

        where `scale_tril` is lower-triangular `k x k` matrix with non-zero
        diagonal, i.e., `tf.diag_part(scale_tril) != 0`.

        Additional leading dimensions (if any) will index batches.

        Args:
          loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
            implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
            `b >= 0` and `k` is the event size.
          scale_tril: Floating-point, lower-triangular `Tensor` with non-zero
            diagonal elements. `scale_tril` has shape `[B1, ..., Bb, k, k]` where
            `b >= 0` and `k` is the event size.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`,
            statistics (e.g., mean, mode, variance) use the value "`NaN`" to
            indicate the result is undefined. When `False`, an exception is raised
            if one or more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          ValueError: if neither `loc` nor `scale_tril` are specified.


def quadrature_scheme_softmaxnormal_gauss_hermite(<EOL>normal_loc, normal_scale, quadrature_size,<EOL>validate_args=False, name=None):

with tf.name_scope(<EOL>name or "<STR_LIT>"):<EOL><INDENT>normal_loc = tf.convert_to_tensor(value=normal_loc, name="<STR_LIT>")<EOL>npdt = dtype_util.as_numpy_dtype(normal_loc.dtype)<EOL>normal_scale = tf.convert_to_tensor(<EOL>value=normal_scale, dtype=npdt, name="<STR_LIT>")<EOL>normal_scale = maybe_check_quadrature_param(<EOL>normal_scale, "<STR_LIT>", validate_args)<EOL>grid, probs = np.polynomial.hermite.hermgauss(deg=quadrature_size)<EOL>grid = grid.astype(npdt)<EOL>probs = probs.astype(npdt)<EOL>probs /= np.linalg.norm(probs, ord=<NUM_LIT:1>, keepdims=True)<EOL>probs = tf.convert_to_tensor(value=probs, name="<STR_LIT>", dtype=npdt)<EOL>grid = softmax(<EOL>-distribution_util.pad(<EOL>(normal_loc[..., tf.newaxis] +<EOL>np.sqrt(<NUM_LIT>) * normal_scale[..., tf.newaxis] * grid),<EOL>axis=-<NUM_LIT:2>,<EOL>front=True),<EOL>axis=-<NUM_LIT:2>)  <EOL>return grid, probs<EOL><DEDENT>

Use Gauss-Hermite quadrature to form quadrature on `K - 1` simplex.

    A `SoftmaxNormal` random variable `Y` may be generated via

    ```
    Y = SoftmaxCentered(X),
    X = Normal(normal_loc, normal_scale)
    ```

    Note: for a given `quadrature_size`, this method is generally less accurate
    than `quadrature_scheme_softmaxnormal_quantiles`.

    Args:
      normal_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
        The location parameter of the Normal used to construct the SoftmaxNormal.
      normal_scale: `float`-like `Tensor`. Broadcastable with `normal_loc`.
        The scale parameter of the Normal used to construct the SoftmaxNormal.
      quadrature_size: Python `int` scalar representing the number of quadrature
        points.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      name: Python `str` name prefixed to Ops created by this class.

    Returns:
      grid: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        convex combination of affine parameters for `K` components.
        `grid[..., :, n]` is the `n`-th grid point, living in the `K - 1` simplex.
      probs:  Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        associated with each grid point.


def quadrature_scheme_softmaxnormal_quantiles(<EOL>normal_loc, normal_scale, quadrature_size,<EOL>validate_args=False, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>normal_loc = tf.convert_to_tensor(value=normal_loc, name="<STR_LIT>")<EOL>dt = dtype_util.base_dtype(normal_loc.dtype)<EOL>normal_scale = tf.convert_to_tensor(<EOL>value=normal_scale, dtype=dt, name="<STR_LIT>")<EOL>normal_scale = maybe_check_quadrature_param(<EOL>normal_scale, "<STR_LIT>", validate_args)<EOL>dist = normal.Normal(loc=normal_loc, scale=normal_scale)<EOL>def _get_batch_ndims():<EOL><INDENT>"""<STR_LIT>"""<EOL>ndims = tensorshape_util.rank(dist.batch_shape)<EOL>if ndims is None:<EOL><INDENT>ndims = tf.shape(input=dist.batch_shape_tensor())[<NUM_LIT:0>]<EOL><DEDENT>return ndims<EOL><DEDENT>batch_ndims = _get_batch_ndims()<EOL>def _get_final_shape(qs):<EOL><INDENT>"""<STR_LIT>"""<EOL>bs = tensorshape_util.with_rank_at_least(dist.batch_shape, <NUM_LIT:1>)<EOL>num_components = tf.compat.dimension_value(bs[-<NUM_LIT:1>])<EOL>if num_components is not None:<EOL><INDENT>num_components += <NUM_LIT:1><EOL><DEDENT>tail = tf.TensorShape([num_components, qs])<EOL>return bs[:-<NUM_LIT:1>].concatenate(tail)<EOL><DEDENT>def _compute_quantiles():<EOL><INDENT>"""<STR_LIT>"""<EOL>zero = tf.zeros([], dtype=dist.dtype)<EOL>edges = tf.linspace(zero, <NUM_LIT:1.>, quadrature_size + <NUM_LIT:3>)[<NUM_LIT:1>:-<NUM_LIT:1>]<EOL>edges = tf.reshape(<EOL>edges,<EOL>shape=tf.concat(<EOL>[[-<NUM_LIT:1>], tf.ones([batch_ndims], dtype=tf.int32)], axis=<NUM_LIT:0>))<EOL>quantiles = dist.quantile(edges)<EOL>quantiles = softmax_centered_bijector.SoftmaxCentered().forward(quantiles)<EOL>perm = tf.concat([tf.range(<NUM_LIT:1>, <NUM_LIT:1> + batch_ndims), [<NUM_LIT:0>]], axis=<NUM_LIT:0>)<EOL>quantiles = tf.transpose(a=quantiles, perm=perm)<EOL>tensorshape_util.set_shape(<EOL>quantiles, _get_final_shape(quadrature_size + <NUM_LIT:1>))<EOL>return quantiles<EOL><DEDENT>quantiles = _compute_quantiles()<EOL>grid = (quantiles[..., :-<NUM_LIT:1>] + quantiles[..., <NUM_LIT:1>:]) / <NUM_LIT><EOL>tensorshape_util.set_shape(grid, _get_final_shape(quadrature_size))<EOL>probs = tf.fill(<EOL>dims=[quadrature_size], value=<NUM_LIT:1.> / tf.cast(quadrature_size, dist.dtype))<EOL>return grid, probs<EOL><DEDENT>

Use SoftmaxNormal quantiles to form quadrature on `K - 1` simplex.

    A `SoftmaxNormal` random variable `Y` may be generated via

    ```
    Y = SoftmaxCentered(X),
    X = Normal(normal_loc, normal_scale)
    ```

    Args:
      normal_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
        The location parameter of the Normal used to construct the SoftmaxNormal.
      normal_scale: `float`-like `Tensor`. Broadcastable with `normal_loc`.
        The scale parameter of the Normal used to construct the SoftmaxNormal.
      quadrature_size: Python `int` scalar representing the number of quadrature
        points.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      name: Python `str` name prefixed to Ops created by this class.

    Returns:
      grid: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        convex combination of affine parameters for `K` components.
        `grid[..., :, n]` is the `n`-th grid point, living in the `K - 1` simplex.
      probs:  Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        associated with each grid point.


def maybe_check_quadrature_param(param, name, validate_args):

with tf.name_scope("<STR_LIT>" + name):<EOL><INDENT>assertions = []<EOL>if tensorshape_util.rank(param.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(param.shape) == <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>name, tensorshape_util.rank(param.shape)))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>assert_util.assert_rank_at_least(<EOL>param,<EOL><NUM_LIT:1>,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>".format(name))))<EOL><DEDENT>if tensorshape_util.with_rank_at_least(param.shape, <NUM_LIT:1>)[-<NUM_LIT:1>] is not None:<EOL><INDENT>if tf.compat.dimension_value(param.shape[-<NUM_LIT:1>]) != <NUM_LIT:1>:<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>name,<EOL>tf.compat.dimension_value(<EOL>param.shape[-<NUM_LIT:1>])))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>assert_util.assert_equal(<EOL>tf.shape(input=param)[-<NUM_LIT:1>],<EOL><NUM_LIT:1>,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>".format(name))))<EOL><DEDENT>if assertions:<EOL><INDENT>return distribution_util.with_dependencies(assertions, param)<EOL><DEDENT>return param<EOL><DEDENT>

Helper which checks validity of `loc` and `scale` init args.


with tf.name_scope("<STR_LIT>"):<EOL><INDENT>batch_shape = grid.shape[:-<NUM_LIT:2>]<EOL>batch_shape_tensor = tf.shape(input=grid)[:-<NUM_LIT:2>]<EOL>event_shape = None<EOL>event_shape_tensor = None<EOL>def _set_event_shape(shape, shape_tensor):<EOL><INDENT>if event_shape is None:<EOL><INDENT>return shape, shape_tensor<EOL><DEDENT>return (tf.broadcast_static_shape(event_shape, shape),<EOL>tf.broadcast_dynamic_shape(event_shape_tensor, shape_tensor))<EOL><DEDENT>for aff in endpoint_affine:<EOL><INDENT>if aff.shift is not None:<EOL><INDENT>batch_shape = tf.broadcast_static_shape(batch_shape,<EOL>aff.shift.shape[:-<NUM_LIT:1>])<EOL>batch_shape_tensor = tf.broadcast_dynamic_shape(<EOL>batch_shape_tensor,<EOL>tf.shape(input=aff.shift)[:-<NUM_LIT:1>])<EOL>event_shape, event_shape_tensor = _set_event_shape(<EOL>aff.shift.shape[-<NUM_LIT:1>:],<EOL>tf.shape(input=aff.shift)[-<NUM_LIT:1>:])<EOL><DEDENT>if aff.scale is not None:<EOL><INDENT>batch_shape = tf.broadcast_static_shape(batch_shape,<EOL>aff.scale.batch_shape)<EOL>batch_shape_tensor = tf.broadcast_dynamic_shape(<EOL>batch_shape_tensor, aff.scale.batch_shape_tensor())<EOL>event_shape, event_shape_tensor = _set_event_shape(<EOL>tf.TensorShape([aff.scale.range_dimension]),<EOL>aff.scale.range_dimension_tensor()[tf.newaxis])<EOL><DEDENT><DEDENT>return batch_shape, batch_shape_tensor, event_shape, event_shape_tensor<EOL><DEDENT>

Helper to infer batch_shape and event_shape.


if len(loc) != <NUM_LIT:2>:<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>".format(len(loc)))<EOL><DEDENT>deg = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(grid.shape, <NUM_LIT:1>)[-<NUM_LIT:1>])<EOL>if deg is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>if loc is None or loc[<NUM_LIT:0>] is None and loc[<NUM_LIT:1>] is None:<EOL><INDENT>return [None]*deg<EOL><DEDENT>w = grid[..., tf.newaxis, :, :]<EOL>loc = [<EOL>x[..., tf.newaxis]  <EOL>if x is not None else None for x in loc<EOL>]<EOL>if loc[<NUM_LIT:0>] is None:<EOL><INDENT>x = w[..., <NUM_LIT:1>, :] * loc[<NUM_LIT:1>]                        <EOL><DEDENT>elif loc[<NUM_LIT:1>] is None:<EOL><INDENT>x = w[..., <NUM_LIT:0>, :] * loc[<NUM_LIT:0>]                        <EOL><DEDENT>else:<EOL><INDENT>delta = loc[<NUM_LIT:0>] - loc[<NUM_LIT:1>]<EOL>x = w[..., <NUM_LIT:0>, :] * delta + loc[<NUM_LIT:1>]                <EOL><DEDENT>return [x[..., k] for k in range(deg)]             <EOL><DEDENT>

Helper which interpolates between two locs.


if len(scale) != <NUM_LIT:2>:<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>".format(len(scale)))<EOL><DEDENT>deg = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(grid.shape, <NUM_LIT:1>)[-<NUM_LIT:1>])<EOL>if deg is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>return [linop_add_lib.add_operators([<EOL>linop_scale(grid[..., k, q], s)<EOL>for k, s in enumerate(scale)<EOL>])[<NUM_LIT:0>] for q in range(deg)]<EOL><DEDENT>

Helper which interpolates between two scales.


<EOL>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>def scaled_identity(w):<EOL><INDENT>return tf.linalg.LinearOperatorScaledIdentity(<EOL>num_rows=op.range_dimension_tensor(),<EOL>multiplier=w,<EOL>is_non_singular=op.is_non_singular,<EOL>is_self_adjoint=op.is_self_adjoint,<EOL>is_positive_definite=op.is_positive_definite)<EOL><DEDENT>if isinstance(op, tf.linalg.LinearOperatorIdentity):<EOL><INDENT>return scaled_identity(w)<EOL><DEDENT>if isinstance(op, tf.linalg.LinearOperatorScaledIdentity):<EOL><INDENT>return scaled_identity(w * op.multiplier)<EOL><DEDENT>if isinstance(op, tf.linalg.LinearOperatorDiag):<EOL><INDENT>return tf.linalg.LinearOperatorDiag(<EOL>diag=w[..., tf.newaxis] * op.diag_part(),<EOL>is_non_singular=op.is_non_singular,<EOL>is_self_adjoint=op.is_self_adjoint,<EOL>is_positive_definite=op.is_positive_definite)<EOL><DEDENT>if isinstance(op, tf.linalg.LinearOperatorLowerTriangular):<EOL><INDENT>return tf.linalg.LinearOperatorLowerTriangular(<EOL>tril=w[..., tf.newaxis, tf.newaxis] * op.to_dense(),<EOL>is_non_singular=op.is_non_singular,<EOL>is_self_adjoint=op.is_self_adjoint,<EOL>is_positive_definite=op.is_positive_definite)<EOL><DEDENT>raise NotImplementedError(<EOL>"<STR_LIT>".format(type(op).__name__))<EOL><DEDENT>

Creates weighted `LinOp` from existing `LinOp`.


args_ = [tf.get_static_value(x) for x in args]<EOL>if any(vec is None for vec in args_):<EOL><INDENT>return tf.concat(args, axis=<NUM_LIT:0>)<EOL><DEDENT>return [val for vec in args_ for val in vec]<EOL>

Concatenates input vectors, statically if possible.


if x is None:<EOL><INDENT>return y<EOL><DEDENT>if y is None:<EOL><INDENT>return x<EOL><DEDENT>return x + y<EOL>


return x[..., :, tf.newaxis] * x[..., tf.newaxis, :]<EOL>

Computes the outer-product of a (batch of) vector, i.e., x.T x.


with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>ndims = (<EOL>tensorshape_util.rank(x.shape)<EOL>if tensorshape_util.rank(x.shape) is not None else tf.rank(<EOL>x, name="<STR_LIT>"))<EOL>axis = tf.convert_to_tensor(value=axis, dtype=tf.int32, name="<STR_LIT>")<EOL>axis_ = tf.get_static_value(axis)<EOL>if axis_ is not None:<EOL><INDENT>axis = np.int(ndims + axis_ if axis_ < <NUM_LIT:0> else axis_)<EOL><DEDENT>else:<EOL><INDENT>axis = tf.where(axis < <NUM_LIT:0>, ndims + axis, axis)<EOL><DEDENT><DEDENT>return tf.nn.softmax(x, axis=axis)<EOL>

Equivalent to tf.nn.softmax but works around b/70297725.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>if not scale or len(scale) < <NUM_LIT:2>:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>if loc is None:<EOL><INDENT>loc = [None]*len(scale)<EOL><DEDENT>if len(loc) != len(scale):<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>dtype = dtype_util.base_dtype(scale[<NUM_LIT:0>].dtype)<EOL>loc = [<EOL>tf.convert_to_tensor(value=loc_, dtype=dtype, name="<STR_LIT>".format(k))<EOL>if loc_ is not None else None for k, loc_ in enumerate(loc)<EOL>]<EOL>for k, scale_ in enumerate(scale):<EOL><INDENT>if validate_args and not scale_.is_positive_definite:<EOL><INDENT>raise ValueError("<STR_LIT>".format(<EOL>k, scale_.is_positive_definite))<EOL><DEDENT>if dtype_util.base_dtype(scale_.dtype) != dtype:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>".format(<EOL>k, dtype_util.name(scale_.dtype), dtype_util.name(dtype)))<EOL><DEDENT><DEDENT>self._endpoint_affine = [<EOL>affine_linear_operator_bijector.AffineLinearOperator(<EOL>shift=loc_, scale=scale_,<EOL>validate_args=validate_args,<EOL>name="<STR_LIT>".format(k))<EOL>for k, (loc_, scale_) in enumerate(zip(loc, scale))]<EOL>if len(scale) != <NUM_LIT:2>:<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>".format(len(scale)))<EOL><DEDENT>mix_loc = tf.convert_to_tensor(value=mix_loc, dtype=dtype, name="<STR_LIT>")<EOL>temperature = tf.convert_to_tensor(<EOL>value=temperature, dtype=dtype, name="<STR_LIT>")<EOL>self._grid, probs = tuple(quadrature_fn(<EOL>mix_loc / temperature,<EOL><NUM_LIT:1.> / temperature,<EOL>quadrature_size,<EOL>validate_args))<EOL>self._mixture_distribution = categorical.Categorical(<EOL>logits=tf.math.log(probs),<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats)<EOL>asserts = distribution_util.maybe_check_scalar_distribution(<EOL>distribution, dtype, validate_args)<EOL>if asserts:<EOL><INDENT>self._grid = distribution_util.with_dependencies(<EOL>asserts, self._grid)<EOL><DEDENT>self._distribution = distribution<EOL>self._interpolated_affine = [<EOL>affine_linear_operator_bijector.AffineLinearOperator(<EOL>shift=loc_, scale=scale_,<EOL>validate_args=validate_args,<EOL>name="<STR_LIT>".format(k))<EOL>for k, (loc_, scale_) in enumerate(zip(<EOL>interpolate_loc(self._grid, loc),<EOL>interpolate_scale(self._grid, scale)))]<EOL>[<EOL>self._batch_shape_,<EOL>self._batch_shape_tensor_,<EOL>self._event_shape_,<EOL>self._event_shape_tensor_,<EOL>] = determine_batch_event_shapes(self._grid,<EOL>self._endpoint_affine)<EOL>super(VectorDiffeomixture, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=(<EOL>distribution._graph_parents  <EOL>+ [loc_ for loc_ in loc if loc_ is not None] +<EOL>[p for scale_ in scale for p in scale_.graph_parents]),<EOL>name=name)<EOL><DEDENT>

Constructs the VectorDiffeomixture on `R^d`.

        The vector diffeomixture (VDM) approximates the compound distribution

        ```none
        p(x) = int p(x | z) p(z) dz,
        where z is in the K-simplex, and
        p(x | z) := p(x | loc=sum_k z[k] loc[k], scale=sum_k z[k] scale[k])
        ```

        Args:
          mix_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`.
            In terms of samples, larger `mix_loc[..., k]` ==>
            `Z` is more likely to put more weight on its `kth` component.
          temperature: `float`-like `Tensor`. Broadcastable with `mix_loc`.
            In terms of samples, smaller `temperature` means one component is more
            likely to dominate.  I.e., smaller `temperature` makes the VDM look more
            like a standard mixture of `K` components.
          distribution: `tfp.distributions.Distribution`-like instance. Distribution
            from which `d` iid samples are used as input to the selected affine
            transformation. Must be a scalar-batch, scalar-event distribution.
            Typically `distribution.reparameterization_type = FULLY_REPARAMETERIZED`
            or it is a function of non-trainable parameters. WARNING: If you
            backprop through a VectorDiffeomixture sample and the `distribution`
            is not `FULLY_REPARAMETERIZED` yet is a function of trainable variables,
            then the gradient will be incorrect!
          loc: Length-`K` list of `float`-type `Tensor`s. The `k`-th element
            represents the `shift` used for the `k`-th affine transformation.  If
            the `k`-th item is `None`, `loc` is implicitly `0`.  When specified,
            must have shape `[B1, ..., Bb, d]` where `b >= 0` and `d` is the event
            size.
          scale: Length-`K` list of `LinearOperator`s. Each should be
            positive-definite and operate on a `d`-dimensional vector space. The
            `k`-th element represents the `scale` used for the `k`-th affine
            transformation. `LinearOperator`s must have shape `[B1, ..., Bb, d, d]`,
            `b >= 0`, i.e., characterizes `b`-batches of `d x d` matrices
          quadrature_size: Python `int` scalar representing number of
            quadrature points.  Larger `quadrature_size` means `q_N(x)` better
            approximates `p(x)`.
          quadrature_fn: Python callable taking `normal_loc`, `normal_scale`,
            `quadrature_size`, `validate_args` and returning `tuple(grid, probs)`
            representing the SoftmaxNormal grid and corresponding normalized weight.
            normalized) weight.
            Default value: `quadrature_scheme_softmaxnormal_quantiles`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`,
            statistics (e.g., mean, mode, variance) use the value "`NaN`" to
            indicate the result is undefined. When `False`, an exception is raised
            if one or more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          ValueError: if `not scale or len(scale) < 2`.
          ValueError: if `len(loc) != len(scale)`
          ValueError: if `quadrature_grid_and_probs is not None` and
            `len(quadrature_grid_and_probs[0]) != len(quadrature_grid_and_probs[1])`
          ValueError: if `validate_args` and any not scale.is_positive_definite.
          TypeError: if any scale.dtype != scale[0].dtype.
          TypeError: if any loc.dtype != scale[0].dtype.
          NotImplementedError: if `len(scale) != 2`.
          ValueError: if `not distribution.is_scalar_batch`.
          ValueError: if `not distribution.is_scalar_event`.


return self._mixture_distribution<EOL>

Distribution used to select a convex combination of affine transforms.


return self._distribution<EOL>

Base scalar-event, scalar-batch distribution.


return self._grid<EOL>

Grid of mixing probabilities, one for each grid point.


return self._endpoint_affine<EOL>

Affine transformation for each of `K` components.


return self._interpolated_affine<EOL>

Affine transformation for each convex combination of `K` components.


single_draw_shape = concat_vectors(self.batch_shape_tensor(),<EOL>self.event_shape_tensor())<EOL>m = tf.reshape(<EOL>self.distribution.mean(),  <EOL>shape=tf.ones_like(single_draw_shape, dtype=tf.int32))<EOL>m = tf.tile(m, multiples=single_draw_shape)<EOL>tensorshape_util.set_shape(<EOL>m, tensorshape_util.concatenate(self.batch_shape, self.event_shape))<EOL>return m<EOL>

Ensures `self.distribution.mean()` has `[batch, event]` shape.


parameters = dict(locals())<EOL>if not dtype_util.is_floating(scale.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if validate_args and not scale.is_positive_definite:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([df, loc, scale],<EOL>preferred_dtype=tf.float32)<EOL>with tf.control_dependencies([<EOL>assert_util.assert_positive(df, message="<STR_LIT>")<EOL>] if validate_args else []):<EOL><INDENT>self._df = tf.identity(<EOL>tf.convert_to_tensor(value=df, dtype=dtype), name="<STR_LIT>")<EOL><DEDENT>self._loc = tf.convert_to_tensor(value=loc, name="<STR_LIT>", dtype=dtype)<EOL>self._scale = scale<EOL><DEDENT>super(MultivariateStudentTLinearOperator, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>parameters=parameters,<EOL>graph_parents=[self._df, self._loc] + self._scale.graph_parents,<EOL>name=name,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats)<EOL>self._parameters = parameters<EOL>

Construct Multivariate Student's t-distribution on `R^k`.

        The `batch_shape` is the broadcast shape between `df`, `loc` and `scale`
        arguments.

        The `event_shape` is given by last dimension of the matrix implied by
        `scale`. The last dimension of `loc` must broadcast with this.

        Additional leading dimensions (if any) will index batches.

        Args:
          df: A positive floating-point `Tensor`. Has shape `[B1, ..., Bb]` where `b
            >= 0`.
          loc: Floating-point `Tensor`. Has shape `[B1, ..., Bb, k]` where `k` is
            the event size.
          scale: Instance of `LinearOperator` with a floating `dtype` and shape
            `[B1, ..., Bb, k, k]`.
          validate_args: Python `bool`, default `False`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are invalid,
            correct behavior is not guaranteed.
          allow_nan_stats: Python `bool`, default `True`. If `False`, raise an
            exception if a statistic (e.g. mean/variance/etc...) is undefined for
            any batch member If `True`, batch members with valid parameters leading
            to undefined statistics will return NaN for this statistic.
          name: The name to give Ops created by the initializer.

        Raises:
          TypeError: if not `scale.dtype.is_floating`.
          ValueError: if not `scale.is_positive_definite`.


return self._loc<EOL>

The location parameter of the distribution.

        `loc` applies an elementwise shift to the distribution.

        ```none
        X ~ MultivariateT(loc=0, scale=1)   # Identity scale, zero shift.
        Y = scale @ X + loc
        ```

        Returns:
          The `loc` `Tensor`.


return self._scale<EOL>

The scale parameter of the distribution.

        `scale` applies an affine scale to the distribution.

        ```none
        X ~ MultivariateT(loc=0, scale=1)   # Identity scale, zero shift.
        Y = scale @ X + loc
        ```

        Returns:
          The `scale` `LinearOperator`.


return self._df<EOL>

The degrees of freedom of the distribution.

        This controls the degrees of freedom of the distribution. The tails of the
        distribution get more heavier the smaller `df` is. As `df` goes to
        infinitiy, the distribution approaches the Multivariate Normal with the same
        `loc` and `scale`.

        Returns:
          The `df` `Tensor`.


def _std_var_helper(self, statistic, statistic_name, statistic_ndims,<EOL>df_factor_fn):

df = tf.reshape(<EOL>self.df,<EOL>tf.concat([<EOL>tf.shape(input=self.df),<EOL>tf.ones([statistic_ndims], dtype=tf.int32)<EOL>], -<NUM_LIT:1>))<EOL>df = _broadcast_to_shape(df, tf.shape(input=statistic))<EOL>denom = tf.where(df > <NUM_LIT>, df - <NUM_LIT>, tf.ones_like(df))<EOL>statistic = statistic * df_factor_fn(df / denom)<EOL>inf = dtype_util.as_numpy_dtype(self.dtype)(np.inf)<EOL>result_where_defined = tf.where(<EOL>df > <NUM_LIT>, statistic, tf.fill(tf.shape(input=statistic), inf, name="<STR_LIT>"))<EOL>if self.allow_nan_stats:<EOL><INDENT>nan = dtype_util.as_numpy_dtype(self.dtype)(np.nan)<EOL>return tf.where(df > <NUM_LIT:1.>, result_where_defined,<EOL>tf.fill(tf.shape(input=statistic), nan, name="<STR_LIT>"))<EOL><DEDENT>else:<EOL><INDENT>with tf.control_dependencies([<EOL>assert_util.assert_less(<EOL>tf.cast(<NUM_LIT:1.>, self.dtype),<EOL>df,<EOL>message=statistic_name +<EOL>"<STR_LIT>"),<EOL>]):<EOL><INDENT>return tf.identity(result_where_defined)<EOL><DEDENT><DEDENT>

Helper to compute stddev, covariance and variance.


def _maybe_validate_args(outcomes, logits, probs, validate_args):

assertions = []<EOL>def validate_equal_last_dim(tensor_a, tensor_b, message):<EOL><INDENT>if tensor_a.shape.is_fully_defined() and tensor_b.shape.is_fully_defined():<EOL><INDENT>if tensor_a.shape[-<NUM_LIT:1>] != tensor_b.shape[-<NUM_LIT:1>]:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_equal(<EOL>tf.shape(input=tensor_a)[-<NUM_LIT:1>],<EOL>tf.shape(input=tensor_b)[-<NUM_LIT:1>],<EOL>message=message))<EOL><DEDENT><DEDENT>if logits is not None:<EOL><INDENT>validate_equal_last_dim(<EOL>outcomes,<EOL>logits,<EOL>message='<STR_LIT>')<EOL><DEDENT>if probs is not None:<EOL><INDENT>validate_equal_last_dim(<EOL>outcomes,<EOL>probs,<EOL>message='<STR_LIT>')<EOL><DEDENT>message = '<STR_LIT>'<EOL>if outcomes.shape.ndims is not None:<EOL><INDENT>if outcomes.shape.ndims != <NUM_LIT:1>:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(tf.compat.v1.assert_rank(outcomes, <NUM_LIT:1>, message=message))<EOL><DEDENT>message = '<STR_LIT>'<EOL>if outcomes.shape.num_elements() is not None:<EOL><INDENT>if outcomes.shape.num_elements() == <NUM_LIT:0>:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_greater(<EOL>tf.size(input=outcomes), <NUM_LIT:0>, message=message))<EOL><DEDENT>if validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_equal(<EOL>tf.math.is_strictly_increasing(outcomes),<EOL>True,<EOL>message='<STR_LIT>'))<EOL><DEDENT>return assertions<EOL>

Validate `outcomes`, `logits` and `probs`'s shapes.


parameters = dict(locals())<EOL>with tf.compat.v1.name_scope(<EOL>name, values=[outcomes, logits, probs]) as name:<EOL><INDENT>self._outcomes = tf.convert_to_tensor(value=outcomes, name='<STR_LIT>')<EOL>if validate_args:<EOL><INDENT>assertions = _maybe_validate_args(self._outcomes, logits, probs,<EOL>validate_args)<EOL>with tf.control_dependencies(assertions):<EOL><INDENT>self._outcomes = tf.identity(self._outcomes)<EOL><DEDENT><DEDENT>if dtype_util.is_floating(self._outcomes.dtype):<EOL><INDENT>eps = np.finfo(dtype_util.as_numpy_dtype(self._outcomes.dtype)).eps<EOL>self._rtol = <NUM_LIT:10> * eps if rtol is None else rtol<EOL>self._atol = <NUM_LIT:10> * eps if atol is None else atol<EOL><DEDENT>else:<EOL><INDENT>self._rtol = None<EOL>self._atol = None<EOL><DEDENT>self._categorical = categorical.Categorical(<EOL>logits=logits,<EOL>probs=probs,<EOL>dtype=tf.int32,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats)<EOL><DEDENT>super(FiniteDiscrete, self).__init__(<EOL>dtype=self._outcomes.dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[<EOL>self._outcomes, self._categorical.logits, self._categorical.probs<EOL>],<EOL>name=name)<EOL>

Construct a finite discrete contribution.

        Args:
          outcomes: A 1-D floating or integer `Tensor`, representing a list of
            possible outcomes in strictly ascending order.
          logits: A floating N-D `Tensor`, `N >= 1`, representing the log
            probabilities of a set of FiniteDiscrete distributions. The first `N -
            1` dimensions index into a batch of independent distributions and the
            last dimension represents a vector of logits for each discrete value.
            Only one of `logits` or `probs` should be passed in.
          probs: A floating  N-D `Tensor`, `N >= 1`, representing the probabilities
            of a set of FiniteDiscrete distributions. The first `N - 1` dimensions
            index into a batch of independent distributions and the last dimension
            represents a vector of probabilities for each discrete value. Only one
            of `logits` or `probs` should be passed in.
          rtol: `Tensor` with same `dtype` as `outcomes`. The relative tolerance for
            floating number comparison. Only effective when `outcomes` is a floating
            `Tensor`. Default is `10 * eps`.
          atol: `Tensor` with same `dtype` as `outcomes`. The absolute tolerance for
            floating number comparison. Only effective when `outcomes` is a floating
            `Tensor`. Default is `10 * eps`.
          validate_args:  Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may render incorrect outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value '`NaN`' to indicate the
            result is undefined. When `False`, an exception is raised if one or more
            of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


def squared_frobenius_norm(x):<EOL><INDENT>"""<STR_LIT>"""<EOL>return tf.reduce_sum(input_tensor=tf.square(x), axis=[-<NUM_LIT:2>, -<NUM_LIT:1>])<EOL><DEDENT>def is_diagonal(x):<EOL><INDENT>"""<STR_LIT>"""<EOL>return (isinstance(x, tf.linalg.LinearOperatorIdentity) or<EOL>isinstance(x, tf.linalg.LinearOperatorScaledIdentity) or<EOL>isinstance(x, tf.linalg.LinearOperatorDiag))<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>if is_diagonal(a.scale) and is_diagonal(b.scale):<EOL><INDENT>b_inv_a = (a.stddev() / b.stddev())[..., tf.newaxis]<EOL><DEDENT>else:<EOL><INDENT>b_inv_a = b.scale.solve(a.scale.to_dense())<EOL><DEDENT>kl_div = (<EOL>b.scale.log_abs_determinant() - a.scale.log_abs_determinant() +<EOL><NUM_LIT:0.5> * (-tf.cast(a.scale.domain_dimension_tensor(), a.dtype) +<EOL>squared_frobenius_norm(b_inv_a) + squared_frobenius_norm(<EOL>b.scale.solve((b.mean() - a.mean())[..., tf.newaxis]))))<EOL>tensorshape_util.set_shape(<EOL>kl_div, tf.broadcast_static_shape(a.batch_shape, b.batch_shape))<EOL>return kl_div<EOL><DEDENT>

Batched KL divergence `KL(a || b)` for multivariate Normals.

    With `X`, `Y` both multivariate Normals in `R^k` with means `mu_a`, `mu_b` and
    covariance `C_a`, `C_b` respectively,

    ```
    KL(a || b) = 0.5 * ( L - k + T + Q ),
    L := Log[Det(C_b)] - Log[Det(C_a)]
    T := trace(C_b^{-1} C_a),
    Q := (mu_b - mu_a)^T C_b^{-1} (mu_b - mu_a),
    ```

    This `Op` computes the trace by solving `C_b^{-1} C_a`. Although efficient
    methods for solving systems with `C_b` may be available, a dense version of
    (the square root of) `C_a` is used, so performance is `O(B s k**2)` where `B`
    is the batch size, and `s` is the cost of solving `C_b x = y` for vectors `x`
    and `y`.

    Args:
      a: Instance of `MultivariateNormalLinearOperator`.
      b: Instance of `MultivariateNormalLinearOperator`.
      name: (optional) name to use for created ops. Default "kl_mvn".

    Returns:
      Batchwise `KL(a || b)`.


parameters = dict(locals())<EOL>if scale is None:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if not dtype_util.is_floating(scale.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>with tf.name_scope(name) as name:<EOL><INDENT>loc = loc if loc is None else tf.convert_to_tensor(<EOL>value=loc, name="<STR_LIT>", dtype=scale.dtype)<EOL>batch_shape, event_shape = distribution_util.shapes_from_loc_and_scale(<EOL>loc, scale)<EOL><DEDENT>super(MultivariateNormalLinearOperator, self).__init__(<EOL>distribution=normal.Normal(<EOL>loc=tf.zeros([], dtype=scale.dtype),<EOL>scale=tf.ones([], dtype=scale.dtype)),<EOL>bijector=affine_linear_operator_bijector.AffineLinearOperator(<EOL>shift=loc, scale=scale, validate_args=validate_args),<EOL>batch_shape=batch_shape,<EOL>event_shape=event_shape,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>self._parameters = parameters<EOL>

Construct Multivariate Normal distribution on `R^k`.

        The `batch_shape` is the broadcast shape between `loc` and `scale`
        arguments.

        The `event_shape` is given by last dimension of the matrix implied by
        `scale`. The last dimension of `loc` (if provided) must broadcast with this.

        Recall that `covariance = scale @ scale.T`.

        Additional leading dimensions (if any) will index batches.

        Args:
          loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
            implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
            `b >= 0` and `k` is the event size.
          scale: Instance of `LinearOperator` with same `dtype` as `loc` and shape
            `[B1, ..., Bb, k, k]`.
          validate_args: Python `bool`, default `False`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are
            invalid, correct behavior is not guaranteed.
          allow_nan_stats: Python `bool`, default `True`. If `False`, raise an
            exception if a statistic (e.g. mean/mode/etc...) is undefined for any
            batch member If `True`, batch members with valid parameters leading to
            undefined statistics will return NaN for this statistic.
          name: The name to give Ops created by the initializer.

        Raises:
          ValueError: if `scale` is unspecified.
          TypeError: if not `scale.dtype.is_floating`


return self.bijector.shift<EOL>

The `loc` `Tensor` in `Y = scale @ X + loc`.


return self.bijector.scale<EOL>

The `scale` `LinearOperator` in `Y = scale @ X + loc`.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>scale = distribution_util.make_diag_scale(<EOL>loc=loc,<EOL>scale_diag=scale_diag,<EOL>scale_identity_multiplier=scale_identity_multiplier,<EOL>validate_args=False,<EOL>assert_positive=False)<EOL><DEDENT><DEDENT>super(VectorExponentialDiag, self).__init__(<EOL>loc=loc,<EOL>scale=scale,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL>self._parameters = parameters<EOL>

Construct Vector Exponential distribution supported on a subset of `R^k`.

        The `batch_shape` is the broadcast shape between `loc` and `scale`
        arguments.

        The `event_shape` is given by last dimension of the matrix implied by
        `scale`. The last dimension of `loc` (if provided) must broadcast with this.

        Recall that `covariance = scale @ scale.T`.

        ```none
        scale = diag(scale_diag + scale_identity_multiplier * ones(k))
        ```

        where:

        * `scale_diag.shape = [k]`, and,
        * `scale_identity_multiplier.shape = []`.

        Additional leading dimensions (if any) will index batches.

        If both `scale_diag` and `scale_identity_multiplier` are `None`, then
        `scale` is the Identity matrix.

        Args:
          loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
            implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
            `b >= 0` and `k` is the event size.
          scale_diag: Non-zero, floating-point `Tensor` representing a diagonal
            matrix added to `scale`. May have shape `[B1, ..., Bb, k]`, `b >= 0`,
            and characterizes `b`-batches of `k x k` diagonal matrices added to
            `scale`. When both `scale_identity_multiplier` and `scale_diag` are
            `None` then `scale` is the `Identity`.
          scale_identity_multiplier: Non-zero, floating-point `Tensor` representing
            a scaled-identity-matrix added to `scale`. May have shape
            `[B1, ..., Bb]`, `b >= 0`, and characterizes `b`-batches of scaled
            `k x k` identity matrices added to `scale`. When both
            `scale_identity_multiplier` and `scale_diag` are `None` then `scale` is
            the `Identity`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`,
            statistics (e.g., mean, mode, variance) use the value "`NaN`" to
            indicate the result is undefined. When `False`, an exception is raised
            if one or more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          ValueError: if at most `scale_identity_multiplier` is specified.


if not callable(fn):<EOL><INDENT>raise TypeError("<STR_LIT>".format(fn))<EOL><DEDENT>return types.FunctionType(<EOL>code=fn.__code__, globals=fn.__globals__,<EOL>name=fn.__name__, argdefs=fn.__defaults__,<EOL>closure=fn.__closure__)<EOL>

Create a deep copy of fn.

    Args:
      fn: a callable

    Returns:
      A `FunctionType`: a deep copy of fn.

    Raises:
      TypeError: if `fn` is not a callable.


old_str = old_str or "<STR_LIT>"<EOL>old_str_lines = old_str.split("<STR_LIT:\n>")<EOL>append_str = "<STR_LIT:\n>".join("<STR_LIT>" % line for line in append_str.split("<STR_LIT:\n>"))<EOL>has_args_ix = [<EOL>ix for ix, line in enumerate(old_str_lines)<EOL>if line.strip().lower() == "<STR_LIT>"]<EOL>if has_args_ix:<EOL><INDENT>final_args_ix = has_args_ix[-<NUM_LIT:1>]<EOL>return ("<STR_LIT:\n>".join(old_str_lines[:final_args_ix])<EOL>+ "<STR_LIT>" + append_str + "<STR_LIT>"<EOL>+ "<STR_LIT:\n>".join(old_str_lines[final_args_ix:]))<EOL><DEDENT>else:<EOL><INDENT>return old_str + "<STR_LIT>" + append_str<EOL><DEDENT>

Update old_str by inserting append_str just before the "Args:" section.


def _convert_to_tensor(value, dtype=None, dtype_hint=None, name=None):

if (tf.nest.is_nested(dtype) or<EOL>tf.nest.is_nested(dtype_hint)):<EOL><INDENT>if dtype is None:<EOL><INDENT>fn = lambda v, pd: tf.convert_to_tensor(v, dtype_hint=pd, name=name)<EOL>return tf.nest.map_structure(fn, value, dtype_hint)<EOL><DEDENT>elif dtype_hint is None:<EOL><INDENT>fn = lambda v, d: tf.convert_to_tensor(v, dtype=d, name=name)<EOL>return tf.nest.map_structure(fn, value, dtype_hint)<EOL><DEDENT>else:<EOL><INDENT>fn = lambda v, d, pd: tf.convert_to_tensor(  <EOL>v, dtype=d, dtype_hint=pd, name=name)<EOL>return tf.nest.map_structure(fn, value, dtype, dtype_hint)<EOL><DEDENT><DEDENT>return tf.convert_to_tensor(<EOL>value=value, dtype=dtype, dtype_hint=dtype_hint, name=name)<EOL>

Converts the given `value` to a (structure of) `Tensor`.

    This function converts Python objects of various types to a (structure of)
    `Tensor` objects. It accepts `Tensor` objects, numpy arrays, Python lists, and
    Python scalars. For example:

    Args:
      value: An object whose structure matches that of `dtype ` and/or
        `dtype_hint` and for which each leaf has a registered `Tensor` conversion
        function.
      dtype: Optional (structure of) element type for the returned tensor. If
        missing, the type is inferred from the type of `value`.
      dtype_hint: Optional (structure of) element type for the returned tensor,
        used when dtype is None. In some cases, a caller may not have a dtype in
        mind when converting to a tensor, so dtype_hint can be used as a soft
        preference.  If the conversion to `dtype_hint` is not possible, this
        argument has no effect.
      name: Optional name to use if a new `Tensor` is created.

    Returns:
      tensor: A (structure of) `Tensor` based on `value`.

    Raises:
      TypeError: If no conversion function is registered for `value` to `dtype`.
      RuntimeError: If a registered conversion function returns an invalid value.
      ValueError: If the `value` is a tensor not of given `dtype` in graph mode.


return {k: v for k, v in dict_.items() if v is not val}<EOL>

Removes `dict` keys which have have `self` as value.


<EOL>if isinstance(x, dict):<EOL><INDENT>return _PrettyDict({<EOL>k: _recursively_replace_dict_for_pretty_dict(v)<EOL>for k, v in x.items()})<EOL><DEDENT>if (isinstance(x, collections.Sequence) and<EOL>not isinstance(x, six.string_types)):<EOL><INDENT>args = (_recursively_replace_dict_for_pretty_dict(x_) for x_ in x)<EOL>is_named_tuple = (isinstance(x, tuple) and<EOL>hasattr(x, "<STR_LIT>") and<EOL>hasattr(x, "<STR_LIT>"))<EOL>return type(x)(*args) if is_named_tuple else type(x)(args)<EOL><DEDENT>if isinstance(x, collections.Mapping):<EOL><INDENT>return type(x)(**{k: _recursively_replace_dict_for_pretty_dict(v)<EOL>for k, v in x.items()})<EOL><DEDENT>return x<EOL>

Recursively replace `dict`s with `_PrettyDict`.


def __new__(mcs, classname, baseclasses, attrs):

if not baseclasses:  <EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>which_base = [<EOL>base for base in baseclasses<EOL>if base == _BaseDistribution or issubclass(base, Distribution)]<EOL>base = which_base[<NUM_LIT:0>] if which_base else None<EOL>if base is None or base == _BaseDistribution:<EOL><INDENT>return super(_DistributionMeta, mcs).__new__(<EOL>mcs, classname, baseclasses, attrs)<EOL><DEDENT>if not issubclass(base, Distribution):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>classname, base.__name__))<EOL><DEDENT>for attr in _DISTRIBUTION_PUBLIC_METHOD_WRAPPERS:<EOL><INDENT>if attr in attrs:<EOL><INDENT>continue<EOL><DEDENT>special_attr = "<STR_LIT>".format(attr)<EOL>class_attr_value = attrs.get(attr, None)<EOL>base_attr_value = getattr(base, attr, None)<EOL>if not base_attr_value:<EOL><INDENT>raise AttributeError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(base.__name__, attr))<EOL><DEDENT>class_special_attr_value = attrs.get(special_attr, None)<EOL>if class_special_attr_value is None:<EOL><INDENT>continue<EOL><DEDENT>class_special_attr_docstring = tf_inspect.getdoc(class_special_attr_value)<EOL>if not class_special_attr_docstring:<EOL><INDENT>continue<EOL><DEDENT>class_attr_value = _copy_fn(base_attr_value)<EOL>class_attr_docstring = tf_inspect.getdoc(base_attr_value)<EOL>if class_attr_docstring is None:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>".format(<EOL>base.__name__, attr))<EOL><DEDENT>class_attr_value.__doc__ = _update_docstring(<EOL>class_attr_value.__doc__,<EOL>"<STR_LIT>".format(<EOL>classname, class_special_attr_docstring))<EOL>attrs[attr] = class_attr_value<EOL><DEDENT>default_init = attrs.get("<STR_LIT>", None)<EOL>if default_init is None:<EOL><INDENT>return super(_DistributionMeta, mcs).__new__(<EOL>mcs, classname, baseclasses, attrs)<EOL><DEDENT>@decorator.decorator<EOL>def wrapped_init(wrapped, self_, *args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>del wrapped<EOL>self_._parameters = None<EOL>default_init(self_, *args, **kwargs)<EOL>if self_._parameters is None:<EOL><INDENT>dummy_self = tuple()<EOL>self_._parameters = lambda: (  <EOL>_remove_dict_keys_with_value(<EOL>inspect.getcallargs(default_init, dummy_self, *args, **kwargs),<EOL>dummy_self))<EOL><DEDENT>elif hasattr(self_._parameters, "<STR_LIT>"):<EOL><INDENT>self_._parameters = _remove_dict_keys_with_value(<EOL>self_._parameters, self_)<EOL><DEDENT><DEDENT>attrs["<STR_LIT>"] = wrapped_init(default_init)  <EOL>return super(_DistributionMeta, mcs).__new__(<EOL>mcs, classname, baseclasses, attrs)<EOL>

Control the creation of subclasses of the Distribution class.

        The main purpose of this method is to properly propagate docstrings
        from private Distribution methods, like `_log_prob`, into their
        public wrappers as inherited by the Distribution base class
        (e.g. `log_prob`).

        Args:
          classname: The name of the subclass being created.
          baseclasses: A tuple of parent classes.
          attrs: A dict mapping new attributes to their values.

        Returns:
          The class object.

        Raises:
          TypeError: If `Distribution` is not a subclass of `BaseDistribution`, or
            the new class is derived via multiple inheritance and the first
            parent class is not a subclass of `BaseDistribution`.
          AttributeError:  If `Distribution` does not implement e.g. `log_prob`.
          ValueError:  If a `Distribution` public method lacks a docstring.


graph_parents = [] if graph_parents is None else graph_parents<EOL>for i, t in enumerate(graph_parents):<EOL><INDENT>if t is None or not tf.is_tensor(t):<EOL><INDENT>raise ValueError("<STR_LIT>" % (i, t))<EOL><DEDENT><DEDENT>if not name or name[-<NUM_LIT:1>] != "<STR_LIT:/>":  <EOL><INDENT>non_unique_name = name or type(self).__name__<EOL>with tf.name_scope(non_unique_name) as name:<EOL><INDENT>pass<EOL><DEDENT><DEDENT>self._dtype = dtype<EOL>self._reparameterization_type = reparameterization_type<EOL>self._allow_nan_stats = allow_nan_stats<EOL>self._validate_args = validate_args<EOL>self._parameters = parameters<EOL>self._parameters_sanitized = False<EOL>self._graph_parents = graph_parents<EOL>self._name = name<EOL>

Constructs the `Distribution`.

        **This is a private method for subclass use.**

        Args:
          dtype: The type of the event samples. `None` implies no type-enforcement.
          reparameterization_type: Instance of `ReparameterizationType`.
            If `tfd.FULLY_REPARAMETERIZED`, this
            `Distribution` can be reparameterized in terms of some standard
            distribution with a function whose Jacobian is constant for the support
            of the standard distribution. If `tfd.NOT_REPARAMETERIZED`,
            then no such reparameterization is available.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          parameters: Python `dict` of parameters used to instantiate this
            `Distribution`.
          graph_parents: Python `list` of graph prerequisites of this
            `Distribution`.
          name: Python `str` name prefixed to Ops created by this class. Default:
            subclass name.

        Raises:
          ValueError: if any member of graph_parents is `None` or not a `Tensor`.


@classmethod<EOL><INDENT>def param_shapes(cls, sample_shape, name="<STR_LIT>"):<DEDENT>

with tf.name_scope(name):<EOL><INDENT>return cls._param_shapes(sample_shape)<EOL><DEDENT>

Shapes of parameters given the desired shape of a call to `sample()`.

        This is a class method that describes what key/value arguments are required
        to instantiate the given `Distribution` so that a particular shape is
        returned for that instance's call to `sample()`.

        Subclasses should override class method `_param_shapes`.

        Args:
          sample_shape: `Tensor` or python list/tuple. Desired shape of a call to
            `sample()`.
          name: name to prepend ops with.

        Returns:
          `dict` of parameter name to `Tensor` shapes.


if isinstance(sample_shape, tf.TensorShape):<EOL><INDENT>if not tensorshape_util.is_fully_defined(sample_shape):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>sample_shape = tensorshape_util.as_list(sample_shape)<EOL><DEDENT>params = cls.param_shapes(sample_shape)<EOL>static_params = {}<EOL>for name, shape in params.items():<EOL><INDENT>static_shape = tf.get_static_value(shape)<EOL>if static_shape is None:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>static_params[name] = tf.TensorShape(static_shape)<EOL><DEDENT>return static_params<EOL>

param_shapes with static (i.e. `TensorShape`) shapes.

        This is a class method that describes what key/value arguments are required
        to instantiate the given `Distribution` so that a particular shape is
        returned for that instance's call to `sample()`. Assumes that the sample's
        shape is known statically.

        Subclasses should override class method `_param_shapes` to return
        constant-valued tensors when constant values are fed.

        Args:
          sample_shape: `TensorShape` or python list/tuple. Desired shape of a call
            to `sample()`.

        Returns:
          `dict` of parameter name to `TensorShape`.

        Raises:
          ValueError: if `sample_shape` is a `TensorShape` and is not fully defined.


return self._name if hasattr(self, "<STR_LIT>") else None<EOL>

Name prepended to all ops created by this `Distribution`.


return self._dtype if hasattr(self, "<STR_LIT>") else None<EOL>

The `DType` of `Tensor`s handled by this `Distribution`.


<EOL>if (not hasattr(self, "<STR_LIT>") or<EOL>not self._parameters_sanitized):<EOL><INDENT>if callable(self._parameters):<EOL><INDENT>self._parameters = self._parameters()<EOL><DEDENT>self._parameters = {<EOL>k: v for k, v in self._parameters.items()<EOL>if not k.startswith("<STR_LIT>") and v is not self}<EOL>self._parameters_sanitized = True<EOL><DEDENT>return self._parameters<EOL>

Dictionary of parameters used to instantiate this `Distribution`.


raise NotImplementedError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(cls))<EOL>

Returns a dict mapping constructor argument names to per-event rank.

        Distributions may implement this method to provide support for slicing
        (`__getitem__`) on the batch axes.

        Examples: Normal has scalar parameters, so would return
        `{'loc': 0, 'scale': 0}`. On the other hand, MultivariateNormalTriL has
        vector loc and matrix scale, so returns `{'loc': 1, 'scale_tril': 2}`. When
        a distribution accepts multiple parameterizations, either all possible
        parameters may be specified by the dict, e.g. Bernoulli returns
        `{'logits': 0, 'probs': 0}`, or if convenient only the parameters relevant
        to this instance may be specified.

        Parameter dtypes are inferred from Tensor attributes on the distribution
        where available, e.g. `bernoulli.probs`, 'mvn.scale_tril', falling back with
        a warning to the dtype of the distribution.

        Returns:
          params_event_ndims: Per-event parameter ranks, a `str->int dict`.


return slicing.batch_slice(self, self._params_event_ndims(), {}, slices)<EOL>

Slices the batch axes of this distribution, returning a new instance.

        ```python
        b = tfd.Bernoulli(logits=tf.zeros([3, 5, 7, 9]))
        b.batch_shape  # => [3, 5, 7, 9]
        b2 = b[:, tf.newaxis, ..., -2:, 1::2]
        b2.batch_shape  # => [3, 1, 5, 2, 4]

        x = tf.random.normal([5, 3, 2, 2])
        cov = tf.matmul(x, x, transpose_b=True)
        chol = tf.cholesky(cov)
        loc = tf.random.normal([4, 1, 3, 1])
        mvn = tfd.MultivariateNormalTriL(loc, chol)
        mvn.batch_shape  # => [4, 5, 3]
        mvn.event_shape  # => [2]
        mvn2 = mvn[:, 3:, ..., ::-1, tf.newaxis]
        mvn2.batch_shape  # => [4, 2, 3, 1]
        mvn2.event_shape  # => [2]
        ```

        Args:
          slices: slices from the [] operator

        Returns:
          dist: A new `tfd.Distribution` instance with sliced parameters.


return self._reparameterization_type<EOL>

Describes how samples from the distribution are reparameterized.

        Currently this is one of the static instances
        `tfd.FULLY_REPARAMETERIZED` or `tfd.NOT_REPARAMETERIZED`.

        Returns:
          An instance of `ReparameterizationType`.


return self._allow_nan_stats<EOL>

Python `bool` describing behavior when a stat is undefined.

        Stats return +/- infinity when it makes sense. E.g., the variance of a
        Cauchy distribution is infinity. However, sometimes the statistic is
        undefined, e.g., if a distribution's pdf does not achieve a maximum within
        the support of the distribution, the mode is undefined. If the mean is
        undefined, then by definition the variance is undefined. E.g. the mean for
        Student's T for df = 1 is undefined (no clear way to say it is either + or -
        infinity), so the variance = E[(X - mean)**2] is also undefined.

        Returns:
          allow_nan_stats: Python `bool`.


return self._validate_args<EOL>

Python `bool` indicating possibly expensive checks are enabled.


try:<EOL><INDENT>return slicing.batch_slice(self, self._params_event_ndims(),<EOL>override_parameters_kwargs, Ellipsis)<EOL><DEDENT>except NotImplementedError:<EOL><INDENT>parameters = dict(self.parameters, **override_parameters_kwargs)<EOL>d = type(self)(**parameters)<EOL>d._parameters = parameters<EOL>d._parameters_sanitized = True<EOL>return d<EOL><DEDENT>

Creates a deep copy of the distribution.

        Note: the copy distribution may continue to depend on the original
        initialization arguments.

        Args:
          **override_parameters_kwargs: String/value dictionary of initialization
            arguments to override with new values.

        Returns:
          distribution: A new instance of `type(self)` initialized from the union
            of self.parameters and override_parameters_kwargs, i.e.,
            `dict(self.parameters, **override_parameters_kwargs)`.


Shape of a single sample from a single event index as a 1-D `Tensor`.

        The batch dimensions are indexes into independent, non-identical
        parameterizations of this distribution.

        Args:
          name: name to give to the op

        Returns:
          batch_shape: `Tensor`.


return tf.TensorShape(self._batch_shape())<EOL>

Shape of a single sample from a single event index as a `TensorShape`.

        May be partially defined or unknown.

        The batch dimensions are indexes into independent, non-identical
        parameterizations of this distribution.

        Returns:
          batch_shape: `TensorShape`, possibly unknown.


Shape of a single sample from a single batch as a 1-D int32 `Tensor`.

        Args:
          name: name to give to the op

        Returns:
          event_shape: `Tensor`.


return tf.TensorShape(self._event_shape())<EOL>

Shape of a single sample from a single batch as a `TensorShape`.

        May be partially defined or unknown.

        Returns:
          event_shape: `TensorShape`, possibly unknown.


with self._name_scope(name):<EOL><INDENT>return tf.convert_to_tensor(<EOL>value=self._is_scalar_helper(self.event_shape,<EOL>self.event_shape_tensor),<EOL>name="<STR_LIT>")<EOL><DEDENT>

Indicates that `event_shape == []`.

        Args:
          name: Python `str` prepended to names of ops created by this function.

        Returns:
          is_scalar_event: `bool` scalar `Tensor`.


with self._name_scope(name):<EOL><INDENT>return tf.convert_to_tensor(<EOL>value=self._is_scalar_helper(self.batch_shape,<EOL>self.batch_shape_tensor),<EOL>name="<STR_LIT>")<EOL><DEDENT>

Indicates that `batch_shape == []`.

        Args:
          name: Python `str` prepended to names of ops created by this function.

        Returns:
          is_scalar_batch: `bool` scalar `Tensor`.


def _call_sample_n(self, sample_shape, seed, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>sample_shape = tf.convert_to_tensor(<EOL>value=sample_shape, dtype=tf.int32, name="<STR_LIT>")<EOL>sample_shape, n = self._expand_sample_shape_to_vector(<EOL>sample_shape, "<STR_LIT>")<EOL>samples = self._sample_n(n, seed, **kwargs)<EOL>batch_event_shape = tf.shape(input=samples)[<NUM_LIT:1>:]<EOL>final_shape = tf.concat([sample_shape, batch_event_shape], <NUM_LIT:0>)<EOL>samples = tf.reshape(samples, final_shape)<EOL>samples = self._set_sample_static_shape(samples, sample_shape)<EOL>return samples<EOL><DEDENT>


def sample(self, sample_shape=(), seed=None, name="<STR_LIT>", **kwargs):

return self._call_sample_n(sample_shape, seed, name, **kwargs)<EOL>

Generate samples of the specified shape.

        Note that a call to `sample()` without arguments will generate a single
        sample.

        Args:
          sample_shape: 0D or 1D `int32` `Tensor`. Shape of the generated samples.
          seed: Python integer seed for RNG
          name: name to give to the op.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          samples: a `Tensor` with prepended dimensions `sample_shape`.


def _call_log_prob(self, value, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>value = _convert_to_tensor(<EOL>value, name="<STR_LIT:value>", dtype_hint=self.dtype)<EOL>if hasattr(self, "<STR_LIT>"):<EOL><INDENT>return self._log_prob(value, **kwargs)<EOL><DEDENT>if hasattr(self, "<STR_LIT>"):<EOL><INDENT>return tf.math.log(self._prob(value, **kwargs))<EOL><DEDENT>raise NotImplementedError("<STR_LIT>".format(<EOL>type(self).__name__))<EOL><DEDENT>


def log_prob(self, value, name="<STR_LIT>", **kwargs):

return self._call_log_prob(value, name, **kwargs)<EOL>

Log probability density/mass function.

        Args:
          value: `float` or `double` `Tensor`.
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          log_prob: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
            values of type `self.dtype`.


def _call_prob(self, value, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>value = _convert_to_tensor(<EOL>value, name="<STR_LIT:value>", dtype_hint=self.dtype)<EOL>if hasattr(self, "<STR_LIT>"):<EOL><INDENT>return self._prob(value, **kwargs)<EOL><DEDENT>if hasattr(self, "<STR_LIT>"):<EOL><INDENT>return tf.exp(self._log_prob(value, **kwargs))<EOL><DEDENT>raise NotImplementedError("<STR_LIT>".format(<EOL>type(self).__name__))<EOL><DEDENT>


def prob(self, value, name="<STR_LIT>", **kwargs):

return self._call_prob(value, name, **kwargs)<EOL>

Probability density/mass function.

        Args:
          value: `float` or `double` `Tensor`.
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          prob: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
            values of type `self.dtype`.


def _call_log_cdf(self, value, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>value = _convert_to_tensor(<EOL>value, name="<STR_LIT:value>", dtype_hint=self.dtype)<EOL>if hasattr(self, "<STR_LIT>"):<EOL><INDENT>return self._log_cdf(value, **kwargs)<EOL><DEDENT>if hasattr(self, "<STR_LIT>"):<EOL><INDENT>return tf.math.log(self._cdf(value, **kwargs))<EOL><DEDENT>raise NotImplementedError("<STR_LIT>".format(<EOL>type(self).__name__))<EOL><DEDENT>


def log_cdf(self, value, name="<STR_LIT>", **kwargs):

return self._call_log_cdf(value, name, **kwargs)<EOL>

Log cumulative distribution function.

        Given random variable `X`, the cumulative distribution function `cdf` is:

        ```none
        log_cdf(x) := Log[ P[X <= x] ]
        ```

        Often, a numerical approximation can be used for `log_cdf(x)` that yields
        a more accurate answer than simply taking the logarithm of the `cdf` when
        `x << -1`.

        Args:
          value: `float` or `double` `Tensor`.
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          logcdf: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
            values of type `self.dtype`.


def _call_cdf(self, value, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>value = _convert_to_tensor(<EOL>value, name="<STR_LIT:value>", dtype_hint=self.dtype)<EOL>if hasattr(self, "<STR_LIT>"):<EOL><INDENT>return self._cdf(value, **kwargs)<EOL><DEDENT>if hasattr(self, "<STR_LIT>"):<EOL><INDENT>return tf.exp(self._log_cdf(value, **kwargs))<EOL><DEDENT>raise NotImplementedError("<STR_LIT>".format(<EOL>type(self).__name__))<EOL><DEDENT>


def cdf(self, value, name="<STR_LIT>", **kwargs):

return self._call_cdf(value, name, **kwargs)<EOL>

Cumulative distribution function.

        Given random variable `X`, the cumulative distribution function `cdf` is:

        ```none
        cdf(x) := P[X <= x]
        ```

        Args:
          value: `float` or `double` `Tensor`.
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          cdf: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
            values of type `self.dtype`.


def _call_log_survival_function(self, value, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>value = _convert_to_tensor(<EOL>value, name="<STR_LIT:value>", dtype_hint=self.dtype)<EOL>try:<EOL><INDENT>return self._log_survival_function(value, **kwargs)<EOL><DEDENT>except NotImplementedError as original_exception:<EOL><INDENT>try:<EOL><INDENT>return tf.math.log1p(-self.cdf(value, **kwargs))<EOL><DEDENT>except NotImplementedError:<EOL><INDENT>raise original_exception<EOL><DEDENT><DEDENT><DEDENT>


def log_survival_function(self, value, name="<STR_LIT>",<EOL>**kwargs):

return self._call_log_survival_function(value, name, **kwargs)<EOL>

Log survival function.

        Given random variable `X`, the survival function is defined:

        ```none
        log_survival_function(x) = Log[ P[X > x] ]
                                 = Log[ 1 - P[X <= x] ]
                                 = Log[ 1 - cdf(x) ]
        ```

        Typically, different numerical approximations can be used for the log
        survival function, which are more accurate than `1 - cdf(x)` when `x >> 1`.

        Args:
          value: `float` or `double` `Tensor`.
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          `Tensor` of shape `sample_shape(x) + self.batch_shape` with values of type
            `self.dtype`.


def _call_survival_function(self, value, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>value = _convert_to_tensor(<EOL>value, name="<STR_LIT:value>", dtype_hint=self.dtype)<EOL>try:<EOL><INDENT>return self._survival_function(value, **kwargs)<EOL><DEDENT>except NotImplementedError as original_exception:<EOL><INDENT>try:<EOL><INDENT>return <NUM_LIT:1.> - self.cdf(value, **kwargs)<EOL><DEDENT>except NotImplementedError:<EOL><INDENT>raise original_exception<EOL><DEDENT><DEDENT><DEDENT>


def survival_function(self, value, name="<STR_LIT>", **kwargs):

return self._call_survival_function(value, name, **kwargs)<EOL>

Survival function.

        Given random variable `X`, the survival function is defined:

        ```none
        survival_function(x) = P[X > x]
                             = 1 - P[X <= x]
                             = 1 - cdf(x).
        ```

        Args:
          value: `float` or `double` `Tensor`.
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          `Tensor` of shape `sample_shape(x) + self.batch_shape` with values of type
            `self.dtype`.


def entropy(self, name="<STR_LIT>", **kwargs):

with self._name_scope(name):<EOL><INDENT>return self._entropy(**kwargs)<EOL><DEDENT>


def mean(self, name="<STR_LIT>", **kwargs):

with self._name_scope(name):<EOL><INDENT>return self._mean(**kwargs)<EOL><DEDENT>


def quantile(self, value, name="<STR_LIT>", **kwargs):

return self._call_quantile(value, name, **kwargs)<EOL>

Quantile function. Aka "inverse cdf" or "percent point function".

        Given random variable `X` and `p in [0, 1]`, the `quantile` is:

        ```none
        quantile(p) := x such that P[X <= x] == p
        ```

        Args:
          value: `float` or `double` `Tensor`.
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          quantile: a `Tensor` of shape `sample_shape(x) + self.batch_shape` with
            values of type `self.dtype`.


def variance(self, name="<STR_LIT>", **kwargs):

with self._name_scope(name):<EOL><INDENT>try:<EOL><INDENT>return self._variance(**kwargs)<EOL><DEDENT>except NotImplementedError as original_exception:<EOL><INDENT>try:<EOL><INDENT>return tf.square(self._stddev(**kwargs))<EOL><DEDENT>except NotImplementedError:<EOL><INDENT>raise original_exception<EOL><DEDENT><DEDENT><DEDENT>

Variance.

        Variance is defined as,

        ```none
        Var = E[(X - E[X])**2]
        ```

        where `X` is the random variable associated with this distribution, `E`
        denotes expectation, and `Var.shape = batch_shape + event_shape`.

        Args:
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          variance: Floating-point `Tensor` with shape identical to
            `batch_shape + event_shape`, i.e., the same shape as `self.mean()`.


def stddev(self, name="<STR_LIT>", **kwargs):

with self._name_scope(name):<EOL><INDENT>try:<EOL><INDENT>return self._stddev(**kwargs)<EOL><DEDENT>except NotImplementedError as original_exception:<EOL><INDENT>try:<EOL><INDENT>return tf.sqrt(self._variance(**kwargs))<EOL><DEDENT>except NotImplementedError:<EOL><INDENT>raise original_exception<EOL><DEDENT><DEDENT><DEDENT>

Standard deviation.

        Standard deviation is defined as,

        ```none
        stddev = E[(X - E[X])**2]**0.5
        ```

        where `X` is the random variable associated with this distribution, `E`
        denotes expectation, and `stddev.shape = batch_shape + event_shape`.

        Args:
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          stddev: Floating-point `Tensor` with shape identical to
            `batch_shape + event_shape`, i.e., the same shape as `self.mean()`.


def covariance(self, name="<STR_LIT>", **kwargs):

with self._name_scope(name):<EOL><INDENT>return self._covariance(**kwargs)<EOL><DEDENT>

Covariance.

        Covariance is (possibly) defined only for non-scalar-event distributions.

        For example, for a length-`k`, vector-valued distribution, it is calculated
        as,

        ```none
        Cov[i, j] = Covariance(X_i, X_j) = E[(X_i - E[X_i]) (X_j - E[X_j])]
        ```

        where `Cov` is a (batch of) `k x k` matrix, `0 <= (i, j) < k`, and `E`
        denotes expectation.

        Alternatively, for non-vector, multivariate distributions (e.g.,
        matrix-valued, Wishart), `Covariance` shall return a (batch of) matrices
        under some vectorization of the events, i.e.,

        ```none
        Cov[i, j] = Covariance(Vec(X)_i, Vec(X)_j) = [as above]
        ```

        where `Cov` is a (batch of) `k' x k'` matrices,
        `0 <= (i, j) < k' = reduce_prod(event_shape)`, and `Vec` is some function
        mapping indices of this distribution's event dimensions to indices of a
        length-`k'` vector.

        Args:
          name: Python `str` prepended to names of ops created by this function.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          covariance: Floating-point `Tensor` with shape `[B1, ..., Bn, k', k']`
            where the first `n` dimensions are batch coordinates and
            `k' = reduce_prod(self.event_shape)`.


def mode(self, name="<STR_LIT>", **kwargs):

with self._name_scope(name):<EOL><INDENT>return self._mode(**kwargs)<EOL><DEDENT>


def cross_entropy(self, other, name="<STR_LIT>"):

with self._name_scope(name):<EOL><INDENT>return self._cross_entropy(other)<EOL><DEDENT>

Computes the (Shannon) cross entropy.

        Denote this distribution (`self`) by `P` and the `other` distribution by
        `Q`. Assuming `P, Q` are absolutely continuous with respect to
        one another and permit densities `p(x) dr(x)` and `q(x) dr(x)`, (Shannon)
        cross entropy is defined as:

        ```none
        H[P, Q] = E_p[-log q(X)] = -int_F p(x) log q(x) dr(x)
        ```

        where `F` denotes the support of the random variable `X ~ P`.

        Args:
          other: `tfp.distributions.Distribution` instance.
          name: Python `str` prepended to names of ops created by this function.

        Returns:
          cross_entropy: `self.dtype` `Tensor` with shape `[B1, ..., Bn]`
            representing `n` different calculations of (Shannon) cross entropy.


def kl_divergence(self, other, name="<STR_LIT>"):

with self._name_scope(name):<EOL><INDENT>return self._kl_divergence(other)<EOL><DEDENT>

Computes the Kullback--Leibler divergence.

        Denote this distribution (`self`) by `p` and the `other` distribution by
        `q`. Assuming `p, q` are absolutely continuous with respect to reference
        measure `r`, the KL divergence is defined as:

        ```none
        KL[p, q] = E_p[log(p(X)/q(X))]
                 = -int_F p(x) log q(x) dr(x) + int_F p(x) log p(x) dr(x)
                 = H[p, q] - H[p]
        ```

        where `F` denotes the support of the random variable `X ~ p`, `H[., .]`
        denotes (Shannon) cross entropy, and `H[.]` denotes (Shannon) entropy.

        Args:
          other: `tfp.distributions.Distribution` instance.
          name: Python `str` prepended to names of ops created by this function.

        Returns:
          kl_divergence: `self.dtype` `Tensor` with shape `[B1, ..., Bn]`
            representing `n` different calculations of the Kullback-Leibler
            divergence.


with tf.name_scope(self.name):<EOL><INDENT>with tf.name_scope(name) as scope:<EOL><INDENT>yield scope<EOL><DEDENT><DEDENT>

Helper function to standardize op scope.


def _expand_sample_shape_to_vector(self, x, name):

x_static_val = tf.get_static_value(x)<EOL>if x_static_val is None:<EOL><INDENT>prod = tf.reduce_prod(input_tensor=x)<EOL><DEDENT>else:<EOL><INDENT>prod = np.prod(x_static_val, dtype=dtype_util.as_numpy_dtype(x.dtype))<EOL><DEDENT>x = distribution_util.expand_to_vector(x, tensor_name=name)<EOL>return x, prod<EOL>

Helper to `sample` which ensures input is 1D.


def _set_sample_static_shape(self, x, sample_shape):

<EOL>sample_shape = tf.TensorShape(tf.get_static_value(sample_shape))<EOL>ndims = tensorshape_util.rank(x.shape)<EOL>sample_ndims = tensorshape_util.rank(sample_shape)<EOL>batch_ndims = tensorshape_util.rank(self.batch_shape)<EOL>event_ndims = tensorshape_util.rank(self.event_shape)<EOL>if (ndims is None and<EOL>sample_ndims is not None and<EOL>batch_ndims is not None and<EOL>event_ndims is not None):<EOL><INDENT>ndims = sample_ndims + batch_ndims + event_ndims<EOL>tensorshape_util.set_shape(x, [None] * ndims)<EOL><DEDENT>if ndims is not None and sample_ndims is not None:<EOL><INDENT>shape = tensorshape_util.concatenate(sample_shape,<EOL>[None] * (ndims - sample_ndims))<EOL>tensorshape_util.set_shape(x, shape)<EOL><DEDENT>if ndims is not None and event_ndims is not None:<EOL><INDENT>shape = tf.TensorShape(<EOL>[None]*(ndims - event_ndims)).concatenate(self.event_shape)<EOL>tensorshape_util.set_shape(x, shape)<EOL><DEDENT>if batch_ndims is not None:<EOL><INDENT>if ndims is not None:<EOL><INDENT>if sample_ndims is None and event_ndims is not None:<EOL><INDENT>sample_ndims = ndims - batch_ndims - event_ndims<EOL><DEDENT>elif event_ndims is None and sample_ndims is not None:<EOL><INDENT>event_ndims = ndims - batch_ndims - sample_ndims<EOL><DEDENT><DEDENT>if sample_ndims is not None and event_ndims is not None:<EOL><INDENT>shape = tf.TensorShape([None]*sample_ndims).concatenate(<EOL>self.batch_shape).concatenate([None]*event_ndims)<EOL>tensorshape_util.set_shape(x, shape)<EOL><DEDENT><DEDENT>return x<EOL>

Helper to `sample`; sets static shape info.


def _is_scalar_helper(self, static_shape, dynamic_shape_fn):

if tensorshape_util.rank(static_shape) is not None:<EOL><INDENT>return tensorshape_util.rank(static_shape) == <NUM_LIT:0><EOL><DEDENT>shape = dynamic_shape_fn()<EOL>if tf.compat.dimension_value(shape.shape[<NUM_LIT:0>]) is not None:<EOL><INDENT>return tensorshape_util.as_list(shape.shape) == [<NUM_LIT:0>]<EOL><DEDENT>return tf.equal(tf.shape(input=shape)[<NUM_LIT:0>], <NUM_LIT:0>)<EOL>

Implementation for `is_scalar_batch` and `is_scalar_event`.


return tf.reduce_logsumexp(input_tensor=vs[..., tf.newaxis] + ms, axis=-<NUM_LIT:2>)<EOL>

Multiply tensor of vectors by matrices assuming values stored are logs.


return tf.reduce_logsumexp(input_tensor=ms + vs[..., tf.newaxis, :], axis=-<NUM_LIT:1>)<EOL>

Multiply tensor of matrices by vectors assuming values stored are logs.


return tf.reduce_sum(input_tensor=vs[..., tf.newaxis] * ms, axis=-<NUM_LIT:2>)<EOL>


states = tf.reshape(tf.range(num_states),<EOL>tf.concat([[num_states],<EOL>tf.ones_like(dist.batch_shape_tensor())],<EOL>axis=<NUM_LIT:0>))<EOL>return distribution_util.move_dimension(dist.log_prob(states), <NUM_LIT:0>, -<NUM_LIT:1>)<EOL>

Tabulate log probabilities from a batch of distributions.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._runtime_assertions = []  <EOL>if num_steps < <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>".format(num_steps))<EOL><DEDENT>self._initial_distribution = initial_distribution<EOL>self._observation_distribution = observation_distribution<EOL>self._transition_distribution = transition_distribution<EOL>if (initial_distribution.event_shape is not None and<EOL>tensorshape_util.rank(initial_distribution.event_shape) != <NUM_LIT:0>):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>elif validate_args:<EOL><INDENT>self._runtime_assertions += [<EOL>assert_util.assert_equal(<EOL>tf.shape(input=initial_distribution.event_shape_tensor())[<NUM_LIT:0>],<EOL><NUM_LIT:0>,<EOL>message="<STR_LIT>"<EOL>"<STR_LIT>")<EOL>]<EOL><DEDENT>if (transition_distribution.event_shape is not None and<EOL>tensorshape_util.rank(transition_distribution.event_shape) != <NUM_LIT:0>):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>elif validate_args:<EOL><INDENT>self._runtime_assertions += [<EOL>assert_util.assert_equal(<EOL>tf.shape(input=transition_distribution.event_shape_tensor())[<NUM_LIT:0>],<EOL><NUM_LIT:0>,<EOL>message="<STR_LIT>"<EOL>"<STR_LIT>")<EOL>]<EOL><DEDENT>if (transition_distribution.batch_shape is not None and<EOL>tensorshape_util.rank(transition_distribution.batch_shape) == <NUM_LIT:0>):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>elif validate_args:<EOL><INDENT>self._runtime_assertions += [<EOL>assert_util.assert_greater(<EOL>tf.size(input=transition_distribution.batch_shape_tensor()),<EOL><NUM_LIT:0>,<EOL>message="<STR_LIT>"<EOL>"<STR_LIT>")<EOL>]<EOL><DEDENT>if (observation_distribution.batch_shape is not None and<EOL>tensorshape_util.rank(observation_distribution.batch_shape) == <NUM_LIT:0>):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>elif validate_args:<EOL><INDENT>self._runtime_assertions += [<EOL>assert_util.assert_greater(<EOL>tf.size(input=observation_distribution.batch_shape_tensor()),<EOL><NUM_LIT:0>,<EOL>message="<STR_LIT>"<EOL>"<STR_LIT>")<EOL>]<EOL><DEDENT>with tf.control_dependencies(self._runtime_assertions):<EOL><INDENT>self._num_states = ((transition_distribution.batch_shape and<EOL>transition_distribution.batch_shape[-<NUM_LIT:1>]) or<EOL>transition_distribution.batch_shape_tensor()[-<NUM_LIT:1>])<EOL>observation_states = ((observation_distribution.batch_shape and<EOL>observation_distribution.batch_shape[-<NUM_LIT:1>]) or<EOL>observation_distribution.batch_shape_tensor()[-<NUM_LIT:1>])<EOL><DEDENT>if (tf.is_tensor(self._num_states) or tf.is_tensor(observation_states)):<EOL><INDENT>if validate_args:<EOL><INDENT>self._runtime_assertions += [<EOL>assert_util.assert_equal(<EOL>self._num_states,<EOL>observation_states,<EOL>message="<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>")<EOL>]<EOL><DEDENT><DEDENT>elif self._num_states != observation_states:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>self._log_init = _extract_log_probs(self._num_states,<EOL>initial_distribution)<EOL>self._log_trans = _extract_log_probs(self._num_states,<EOL>transition_distribution)<EOL>self._num_steps = num_steps<EOL>self._num_states = tf.shape(input=self._log_init)[-<NUM_LIT:1>]<EOL>self._underlying_event_rank = tf.size(<EOL>input=self._observation_distribution.event_shape_tensor())<EOL>self.static_event_shape = tf.TensorShape(<EOL>[num_steps]).concatenate(self._observation_distribution.event_shape)<EOL>with tf.control_dependencies(self._runtime_assertions):<EOL><INDENT>self.static_batch_shape = tf.broadcast_static_shape(<EOL>self._initial_distribution.batch_shape,<EOL>tf.broadcast_static_shape(<EOL>self._transition_distribution.batch_shape[:-<NUM_LIT:1>],<EOL>self._observation_distribution.batch_shape[:-<NUM_LIT:1>]))<EOL><DEDENT>super(HiddenMarkovModel, self).__init__(<EOL>dtype=self._observation_distribution.dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=(self._initial_distribution._graph_parents +<EOL>self._transition_distribution._graph_parents +<EOL>self._observation_distribution._graph_parents),<EOL>name=name)<EOL>self._parameters = parameters<EOL><DEDENT>

Initialize hidden Markov model.

        Args:
          initial_distribution: A `Categorical`-like instance.
            Determines probability of first hidden state in Markov chain.
            The number of categories must match the number of categories of
            `transition_distribution` as well as both the rightmost batch
            dimension of `transition_distribution` and the rightmost batch
            dimension of `observation_distribution`.
          transition_distribution: A `Categorical`-like instance.
            The rightmost batch dimension indexes the probability distribution
            of each hidden state conditioned on the previous hidden state.
          observation_distribution: A `tfp.distributions.Distribution`-like
            instance.  The rightmost batch dimension indexes the distribution
            of each observation conditioned on the corresponding hidden state.
          num_steps: The number of steps taken in Markov chain. A python `int`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
            Default value: `True`.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: "HiddenMarkovModel".

        Raises:
          ValueError: if `num_steps` is not at least 1.
          ValueError: if `initial_distribution` does not have scalar `event_shape`.
          ValueError: if `transition_distribution` does not have scalar
            `event_shape.`
          ValueError: if `transition_distribution` and `observation_distribution`
            are fully defined but don't have matching rightmost dimension.


initial_log_probs = tf.broadcast_to(self._log_init,<EOL>tf.concat([self.batch_shape_tensor(),<EOL>[self._num_states]],<EOL>axis=<NUM_LIT:0>))<EOL>if self._num_steps > <NUM_LIT:1>:<EOL><INDENT>transition_log_probs = self._log_trans<EOL>def forward_step(log_probs, _):<EOL><INDENT>return _log_vector_matrix(log_probs, transition_log_probs)<EOL><DEDENT>dummy_index = tf.zeros(self._num_steps - <NUM_LIT:1>, dtype=tf.float32)<EOL>forward_log_probs = tf.scan(forward_step, dummy_index,<EOL>initializer=initial_log_probs,<EOL>name="<STR_LIT>")<EOL>forward_log_probs = tf.concat([[initial_log_probs], forward_log_probs],<EOL>axis=<NUM_LIT:0>)<EOL><DEDENT>else:<EOL><INDENT>forward_log_probs = initial_log_probs[tf.newaxis, ...]<EOL><DEDENT>return tf.exp(forward_log_probs)<EOL>

Compute marginal pdf for each individual observable.


def posterior_marginals(self, observations, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>with tf.control_dependencies(self._runtime_assertions):<EOL><INDENT>observation_tensor_shape = tf.shape(input=observations)<EOL>with self._observation_shape_preconditions(observation_tensor_shape):<EOL><INDENT>observation_batch_shape = observation_tensor_shape[<EOL>:-<NUM_LIT:1> - self._underlying_event_rank]<EOL>observation_event_shape = observation_tensor_shape[<EOL>-<NUM_LIT:1> - self._underlying_event_rank:]<EOL>batch_shape = tf.broadcast_dynamic_shape(observation_batch_shape,<EOL>self.batch_shape_tensor())<EOL>log_init = tf.broadcast_to(self._log_init,<EOL>tf.concat([batch_shape,<EOL>[self._num_states]],<EOL>axis=<NUM_LIT:0>))<EOL>log_transition = self._log_trans<EOL>observations = tf.broadcast_to(observations,<EOL>tf.concat([batch_shape,<EOL>observation_event_shape],<EOL>axis=<NUM_LIT:0>))<EOL>observation_rank = tf.rank(observations)<EOL>underlying_event_rank = self._underlying_event_rank<EOL>observations = distribution_util.move_dimension(<EOL>observations, observation_rank - underlying_event_rank - <NUM_LIT:1>, <NUM_LIT:0>)<EOL>observations = tf.expand_dims(<EOL>observations,<EOL>observation_rank - underlying_event_rank)<EOL>observation_log_probs = self._observation_distribution.log_prob(<EOL>observations)<EOL>log_adjoint_prob = tf.zeros_like(log_init)<EOL>def forward_step(log_previous_step, log_prob_observation):<EOL><INDENT>return _log_vector_matrix(log_previous_step,<EOL>log_transition) + log_prob_observation<EOL><DEDENT>log_prob = log_init + observation_log_probs[<NUM_LIT:0>]<EOL>forward_log_probs = tf.scan(forward_step, observation_log_probs[<NUM_LIT:1>:],<EOL>initializer=log_prob,<EOL>name="<STR_LIT>")<EOL>forward_log_probs = tf.concat([[log_prob], forward_log_probs], axis=<NUM_LIT:0>)<EOL>def backward_step(log_previous_step, log_prob_observation):<EOL><INDENT>return _log_matrix_vector(log_transition,<EOL>log_prob_observation + log_previous_step)<EOL><DEDENT>backward_log_adjoint_probs = tf.scan(<EOL>backward_step,<EOL>observation_log_probs[<NUM_LIT:1>:],<EOL>initializer=log_adjoint_prob,<EOL>reverse=True,<EOL>name="<STR_LIT>")<EOL>total_log_prob = tf.reduce_logsumexp(<EOL>input_tensor=forward_log_probs[-<NUM_LIT:1>], axis=-<NUM_LIT:1>)<EOL>backward_log_adjoint_probs = tf.concat([backward_log_adjoint_probs,<EOL>[log_adjoint_prob]], axis=<NUM_LIT:0>)<EOL>log_likelihoods = forward_log_probs + backward_log_adjoint_probs<EOL>marginal_log_probs = distribution_util.move_dimension(<EOL>log_likelihoods - total_log_prob[..., tf.newaxis], <NUM_LIT:0>, -<NUM_LIT:2>)<EOL>return categorical.Categorical(logits=marginal_log_probs)<EOL><DEDENT><DEDENT><DEDENT>

Compute marginal posterior distribution for each state.

        This function computes, for each time step, the marginal
        conditional probability that the hidden Markov model was in
        each possible state given the observations that were made
        at each time step.
        So if the hidden states are `z[0],...,z[num_steps - 1]` and
        the observations are `x[0], ..., x[num_steps - 1]`, then
        this function computes `P(z[i] | x[0], ..., x[num_steps - 1])`
        for all `i` from `0` to `num_steps - 1`.

        This operation is sometimes called smoothing. It uses a form
        of the forward-backward algorithm.

        Note: the behavior of this function is undefined if the
        `observations` argument represents impossible observations
        from the model.

        Args:
          observations: A tensor representing a batch of observations
            made on the hidden Markov model.  The rightmost dimension of this tensor
            gives the steps in a sequence of observations from a single sample from
            the hidden Markov model. The size of this dimension should match the
            `num_steps` parameter of the hidden Markov model object. The other
            dimensions are the dimensions of the batch and these are broadcast with
            the hidden Markov model's parameters.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: "HiddenMarkovModel".

        Returns:
          posterior_marginal: A `Categorical` distribution object representing the
            marginal probability of the hidden Markov model being in each state at
            each step. The rightmost dimension of the `Categorical` distributions
            batch will equal the `num_steps` parameter providing one marginal
            distribution for each step. The other dimensions are the dimensions
            corresponding to the batch of observations.

        Raises:
          ValueError: if rightmost dimension of `observations` does not
          have size `num_steps`.