<EOL>c0 = self.total_concentration[..., tf.newaxis]<EOL>return tf.sqrt((<NUM_LIT:1.> + c0 / self.total_count[..., tf.newaxis]) / (<NUM_LIT:1.> + c0))<EOL>

Helper to `_covariance` and `_variance` which computes a shared scale.


<EOL>if tf.executing_eagerly():<EOL><INDENT>tf.compat.v1.set_random_seed(seed)<EOL>return None<EOL><DEDENT>return seed<EOL>

Helper which uses graph seed if using TFE.


def _mixture_stddev_np(pi_vector, mu_vector, sigma_vector):

pi_vector = np.expand_dims(pi_vector, axis=<NUM_LIT:1>)<EOL>mean_wa = np.matmul(pi_vector, np.expand_dims(mu_vector, axis=<NUM_LIT:2>))<EOL>var_wa = np.matmul(pi_vector, np.expand_dims(sigma_vector**<NUM_LIT:2>, axis=<NUM_LIT:2>))<EOL>mid_term = np.matmul(pi_vector, np.expand_dims(mu_vector**<NUM_LIT:2>, axis=<NUM_LIT:2>))<EOL>mixture_variance = (<EOL>np.squeeze(var_wa) + np.squeeze(mid_term) - np.squeeze(mean_wa**<NUM_LIT:2>))<EOL>return np.sqrt(mixture_variance)<EOL>

Computes the standard deviation of a univariate mixture distribution.

    Acts upon `np.array`s (not `tf.Tensor`s).

    Args:
      pi_vector: A `np.array` of mixture weights. Shape `[batch, components]`.
      mu_vector: A `np.array` of means. Shape `[batch, components]`
      sigma_vector: A `np.array` of stddevs. Shape `[batch, components]`.

    Returns:
      A `np.array` containing the batch of standard deviations.


def sample_distributions(self, sample_shape=(), seed=None, value=None,<EOL>name='<STR_LIT>'):

with self._name_scope(name):<EOL><INDENT>ds, xs = self._call_flat_sample_distributions(sample_shape, seed, value)<EOL>return self._unflatten(ds), self._unflatten(xs)<EOL><DEDENT>

Generate samples and the (random) distributions.

        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 generating random numbers.
          value: `list` of `Tensor`s in `distribution_fn` order to use to
            parameterize other ("downstream") distribution makers.
            Default value: `None` (i.e., draw a sample from each distribution).
          name: name prepended to ops created by this function.
            Default value: `"sample_distributions"`.

        Returns:
          distributions: a `tuple` of `Distribution` instances for each of
            `distribution_fn`.
          samples: a `tuple` of `Tensor`s with prepended dimensions `sample_shape`
            for each of `distribution_fn`.


<EOL>if self._most_recently_built_distributions is None:<EOL><INDENT>return None<EOL><DEDENT>return self._unflatten(<EOL>tf.TensorShape(None) if d is None else d.event_shape<EOL>for d in self._most_recently_built_distributions)<EOL>

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

        May be partially defined or unknown.

        Returns:
          event_shape: `tuple` of `TensorShape`s representing the `event_shape` for
            each distribution in `model`.


with self._name_scope(name):<EOL><INDENT>return self._unflatten(self._map_attr_over_dists('<STR_LIT>')[<NUM_LIT:1>])<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` for each distribution in `model`.


return self._input_output_cholesky<EOL>

Boolean indicating if `Tensor` input/outputs are Cholesky factorized.


with self._name_scope(name):<EOL><INDENT>return (self.df * self.scale_operator.log_abs_determinant() +<EOL><NUM_LIT:0.5> * self.df * self.dimension * math.log(<NUM_LIT>) +<EOL>self._multi_lgamma(<NUM_LIT:0.5> * self.df, self.dimension))<EOL><DEDENT>

Computes the log normalizing constant, log(Z).


return self._df<EOL>

Wishart distribution degree(s) of freedom.


parameters = dict(locals())<EOL>self._input_output_cholesky = input_output_cholesky<EOL>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>if not dtype_util.is_floating(scale_operator.dtype):<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>" %<EOL>scale_operator.dtype)<EOL><DEDENT>if not scale_operator.is_square:<EOL><INDENT>print(scale_operator.to_dense().eval())<EOL>raise ValueError("<STR_LIT>")<EOL><DEDENT>self._scale_operator = scale_operator<EOL>self._df = tf.convert_to_tensor(<EOL>value=df, dtype=scale_operator.dtype, name="<STR_LIT>")<EOL>dtype_util.assert_same_float_dtype([self._df, self._scale_operator])<EOL>if tf.compat.dimension_value(self._scale_operator.shape[-<NUM_LIT:1>]) is None:<EOL><INDENT>self._dimension = tf.cast(<EOL>self._scale_operator.domain_dimension_tensor(),<EOL>dtype=self._scale_operator.dtype,<EOL>name="<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>self._dimension = tf.convert_to_tensor(<EOL>value=tf.compat.dimension_value(self._scale_operator.shape[-<NUM_LIT:1>]),<EOL>dtype=self._scale_operator.dtype,<EOL>name="<STR_LIT>")<EOL><DEDENT>df_val = tf.get_static_value(self._df)<EOL>dim_val = tf.get_static_value(self._dimension)<EOL>if df_val is not None and dim_val is not None:<EOL><INDENT>df_val = np.asarray(df_val)<EOL>if not df_val.shape:<EOL><INDENT>df_val = [df_val]<EOL><DEDENT>if np.any(df_val < dim_val):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>"<EOL>% (df_val, dim_val))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions = assert_util.assert_less_equal(<EOL>self._dimension,<EOL>self._df,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>" % (self._dimension, self._df)))<EOL>self._df = distribution_util.with_dependencies(<EOL>[assertions], self._df)<EOL><DEDENT><DEDENT><DEDENT>super(_WishartLinearOperator, self).__init__(<EOL>dtype=self._scale_operator.dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>parameters=parameters,<EOL>graph_parents=(<EOL>[self._df, self._dimension] + self._scale_operator.graph_parents),<EOL>name=name)<EOL>

Construct Wishart distributions.

        Args:
          df: `float` or `double` tensor, the degrees of freedom of the
            distribution(s). `df` must be greater than or equal to `k`.
          scale_operator: `float` or `double` instance of `LinearOperator`.
          input_output_cholesky: Python `bool`. If `True`, functions whose input or
            output have the semantics of samples assume inputs are in Cholesky form
            and return outputs in Cholesky form. In particular, if this flag is
            `True`, input to `log_prob` is presumed of Cholesky form and output from
            `sample`, `mean`, and `mode` are of Cholesky form.  Setting this
            argument to `True` is purely a computational optimization and does not
            change the underlying distribution; for instance, `mean` returns the
            Cholesky of the mean, not the mean of Cholesky factors. The `variance`
            and `stddev` methods are unaffected by this flag.
            Default value: `False` (i.e., input/output does not have Cholesky
            semantics).
          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:
          TypeError: if scale is not floating-type
          TypeError: if scale.dtype != df.dtype
          ValueError: if df < k, where scale operator event shape is
            `(k, k)`


def _multi_digamma(self, a, p, name="<STR_LIT>"):

with self._name_scope(name):<EOL><INDENT>seq = self._multi_gamma_sequence(a, p)<EOL>return tf.reduce_sum(input_tensor=tf.math.digamma(seq), axis=[-<NUM_LIT:1>])<EOL><DEDENT>

Computes the multivariate digamma function; Psi_p(a).


return self._scale_operator<EOL>

Wishart distribution scale matrix as an Linear Operator.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale], preferred_dtype=tf.float32)<EOL>loc = tf.convert_to_tensor(value=loc, name="<STR_LIT>", dtype=dtype)<EOL>scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>", dtype=dtype)<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_positive(scale)] if validate_args else []):<EOL><INDENT>self._loc = tf.identity(loc, name="<STR_LIT>")<EOL>self._scale = tf.identity(scale, name="<STR_LIT>")<EOL><DEDENT>dtype_util.assert_same_float_dtype([self._loc, self._scale])<EOL><DEDENT>super(HalfCauchy, self).__init__(<EOL>dtype=self._scale.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=[self._loc, self._scale],<EOL>name=name)<EOL>

Construct a half-Cauchy distribution with `loc` and `scale`.

        Args:
          loc: Floating-point `Tensor`; the location(s) of the distribution(s).
          scale: Floating-point `Tensor`; the scale(s) of the distribution(s).
            Must contain only positive values.
          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` (i.e. do not validate args).
          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: 'HalfCauchy'.

        Raises:
          TypeError: if `loc` and `scale` have different `dtype`.


with tf.name_scope("<STR_LIT>"):<EOL><INDENT>return z * self.scale + self.loc<EOL><DEDENT>

Reconstruct input `x` from a its normalized version.


def _extend_support_with_default_value(self, x, f, default_value):

with tf.name_scope("<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, dtype=self.dtype, name="<STR_LIT:x>")<EOL>loc = self.loc + tf.zeros_like(self.scale) + tf.zeros_like(x)<EOL>x = x + tf.zeros_like(loc)<EOL>y = f(tf.where(x < loc, self._inv_z(<NUM_LIT:0.5>) + tf.zeros_like(x), x))<EOL>if default_value == <NUM_LIT:0.>:<EOL><INDENT>default_value = tf.zeros_like(y)<EOL><DEDENT>elif default_value == <NUM_LIT:1.>:<EOL><INDENT>default_value = tf.ones_like(y)<EOL><DEDENT>else:<EOL><INDENT>default_value = tf.fill(<EOL>dims=tf.shape(input=y),<EOL>value=dtype_util.as_numpy_dtype(self.dtype)(default_value))<EOL><DEDENT>return tf.where(x < loc, default_value, y)<EOL><DEDENT>

Returns `f(x)` if x is in the support, and `default_value` otherwise.

        Given `f` which is defined on the support of this distribution
        (`x >= loc`), extend the function definition to the real line
        by defining `f(x) = default_value` for `x < loc`.

        Args:
          x: Floating-point `Tensor` to evaluate `f` at.
          f: Callable that takes in a `Tensor` and returns a `Tensor`. This
            represents the function whose domain of definition we want to extend.
          default_value: Python or numpy literal representing the value to use for
            extending the domain.
        Returns:
          `Tensor` representing an extension of `f(x)`.


return self._scale<EOL>


@functools.wraps(f)<EOL>def _check_arg_and_apply_f(*args, **kwargs):<EOL><INDENT>dist = args[<NUM_LIT:0>]<EOL>x = args[<NUM_LIT:1>]<EOL>with tf.control_dependencies([<EOL>assert_util.assert_greater_equal(<EOL>x, dist.loc, message="<STR_LIT>")<EOL>] if dist.validate_args else []):<EOL><INDENT>return f(*args, **kwargs)<EOL><DEDENT><DEDENT>return _check_arg_and_apply_f<EOL>

Decorator function for argument bounds checking.

    This decorator is meant to be used with methods that require the first
    argument to be in the support of the distribution. If `validate_args` is
    `True`, the method is wrapped with an assertion that the first argument is
    greater than or equal to `loc`, since the support of the half-Cauchy
    distribution is given by `[loc, infinity)`.

    Args:
      f: method to be decorated.

    Returns:
      Returns a decorated method that, when `validate_args` attribute of the class
      is `True`, will assert that all elements in the first argument are within
      the support of the distribution before executing the original method.


return self._probs<EOL>


@kullback_leibler.RegisterKL(OneHotCategorical, OneHotCategorical)<EOL>def _kl_categorical_categorical(a, b, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return tf.reduce_sum(<EOL>input_tensor=tf.nn.softmax(a.logits) *<EOL>(tf.nn.log_softmax(a.logits) - tf.nn.log_softmax(b.logits)),<EOL>axis=-<NUM_LIT:1>)<EOL><DEDENT>

Calculate the batched KL divergence KL(a || b) with a, b OneHotCategorical.

    Args:
      a: instance of a OneHotCategorical distribution object.
      b: instance of a OneHotCategorical distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_categorical_categorical".

    Returns:
      Batchwise KL(a || b)


def _explore(v):<EOL><INDENT>if v.depth < <NUM_LIT:0>:<EOL><INDENT>v.depth = ((<NUM_LIT:1> + max([-<NUM_LIT:1>] + [_explore(annotated_graph[u])<EOL>for u in v.parents]))<EOL>if v.parents else <NUM_LIT:0>)<EOL><DEDENT>return v.depth<EOL><DEDENT>annotated_graph = {k: _Node(k, v) for k, v in g.items()}<EOL>for v in annotated_graph.values():<EOL><INDENT>_explore(v)<EOL><DEDENT>return annotated_graph<EOL>

Computes the number of edges on longest path from node to root.


if not _is_dict_like(model):<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>type(model).__name__))<EOL><DEDENT>[<EOL>self._dist_fn,<EOL>self._dist_fn_wrapped,<EOL>self._dist_fn_args,<EOL>self._dist_fn_name,  <EOL>] = _prob_chain_rule_flatten(model)<EOL>

Creates `dist_fn`, `dist_fn_wrapped`, `dist_fn_args`, `dist_fn_name`.


def _log_unnorm_prob(self, x, name=None):

with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if self.input_output_cholesky:<EOL><INDENT>logdet = <NUM_LIT> * tf.reduce_sum(<EOL>input_tensor=tf.math.log(tf.linalg.diag_part(x)), axis=[-<NUM_LIT:1>])<EOL><DEDENT>else:<EOL><INDENT>_, logdet = tf.linalg.slogdet(x)<EOL><DEDENT>answer = (self.concentration - <NUM_LIT:1.>) * logdet<EOL>return answer<EOL><DEDENT>

Returns the unnormalized log density of an LKJ distribution.

        Args:
          x: `float` or `double` `Tensor` of correlation matrices.  The shape of `x`
            must be `B + [D, D]`, where `B` broadcasts with the shape of
            `concentration`.
          name: Python `str` name prefixed to Ops created by this function.

        Returns:
          log_p: A Tensor of the unnormalized log density of each matrix element of
            `x`, with respect to an LKJ distribution with parameter the
            corresponding element of `concentration`.


return self._concentration<EOL>


def _sample_n(self, num_samples, seed=None, name=None):

if self.dimension < <NUM_LIT:0>:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>')<EOL><DEDENT>seed = seed_stream.SeedStream(seed, '<STR_LIT>')<EOL>with tf.name_scope('<STR_LIT>' or name):<EOL><INDENT>if not dtype_util.is_floating(self.concentration.dtype):<EOL><INDENT>raise TypeError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT:{}>'.format(dtype_util.name(self.concentration.dtype)))<EOL><DEDENT>concentration = _replicate(num_samples, self.concentration)<EOL>concentration_shape = tf.shape(input=concentration)<EOL>if self.dimension <= <NUM_LIT:1>:<EOL><INDENT>shape = tf.concat([<EOL>concentration_shape, [self.dimension, self.dimension]], axis=<NUM_LIT:0>)<EOL>return tf.ones(shape=shape, dtype=self.concentration.dtype)<EOL><DEDENT>beta_conc = concentration + (self.dimension - <NUM_LIT>) / <NUM_LIT><EOL>beta_dist = beta.Beta(concentration1=beta_conc, concentration0=beta_conc)<EOL>corr12 = <NUM_LIT> * beta_dist.sample(seed=seed()) - <NUM_LIT:1.><EOL>first_row = tf.concat([<EOL>tf.ones_like(corr12)[..., tf.newaxis],<EOL>tf.zeros_like(corr12)[..., tf.newaxis]], axis=-<NUM_LIT:1>)<EOL>second_row = tf.concat([<EOL>corr12[..., tf.newaxis],<EOL>tf.sqrt(<NUM_LIT:1> - corr12**<NUM_LIT:2>)[..., tf.newaxis]], axis=-<NUM_LIT:1>)<EOL>chol_result = tf.concat([<EOL>first_row[..., tf.newaxis, :],<EOL>second_row[..., tf.newaxis, :]], axis=-<NUM_LIT:2>)<EOL>for n in range(<NUM_LIT:2>, self.dimension):<EOL><INDENT>beta_conc -= <NUM_LIT:0.5><EOL>norm = beta.Beta(<EOL>concentration1=n/<NUM_LIT>,<EOL>concentration0=beta_conc<EOL>).sample(seed=seed())<EOL>distance = tf.sqrt(norm)[..., tf.newaxis]<EOL>direction = _uniform_unit_norm(<EOL>n, concentration_shape, self.concentration.dtype, seed)<EOL>raw_correlation = distance * direction  <EOL>new_row = tf.concat(<EOL>[raw_correlation, tf.sqrt(<NUM_LIT:1.> - norm[..., tf.newaxis])], axis=-<NUM_LIT:1>)<EOL>chol_result = tf.concat([<EOL>chol_result,<EOL>tf.zeros_like(chol_result[..., <NUM_LIT:0>][..., tf.newaxis])], axis=-<NUM_LIT:1>)<EOL>chol_result = tf.concat(<EOL>[chol_result, new_row[..., tf.newaxis, :]], axis=-<NUM_LIT:2>)<EOL><DEDENT>if self.input_output_cholesky:<EOL><INDENT>return chol_result<EOL><DEDENT>result = tf.matmul(chol_result, chol_result, transpose_b=True)<EOL>result = tf.linalg.set_diag(<EOL>result,<EOL>tf.ones(shape=tf.shape(input=result)[:-<NUM_LIT:1>], dtype=result.dtype))<EOL>return result<EOL><DEDENT>

Returns a Tensor of samples from an LKJ distribution.

        Args:
          num_samples: Python `int`. The number of samples to draw.
          seed: Python integer seed for RNG
          name: Python `str` name prefixed to Ops created by this function.

        Returns:
          samples: A Tensor of correlation matrices with shape `[n, B, D, D]`,
            where `B` is the shape of the `concentration` parameter, and `D`
            is the `dimension`.

        Raises:
          ValueError: If `dimension` is negative.


return self._input_output_cholesky<EOL>


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`.


@kullback_leibler.RegisterKL(GaussianProcess, normal.Normal)<EOL>def _kl_gp_normal(gp, n, name=None):

return _kl_gp_compatible(gp, n, name or '<STR_LIT>')<EOL>

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

    Args:
      gp: instance of a GaussianProcess distribution object.
      n: instance of a Normal distribution object.
      name: (optional) Name to use for created operations.
        default is 'kl_gp_normal'.

    Returns:
      Batchwise KL(gp || n)


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.


return self._interpolated_affine<EOL>

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


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


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.


return self._endpoint_affine<EOL>

Affine transformation for each of `K` components.


return self._grid<EOL>

Grid of mixing probabilities, one for 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.


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.


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.


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`


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 _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`.


<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`.


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.


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_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 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 _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.


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

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


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 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 _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 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 __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.


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

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


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.


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`.


return self._scale<EOL>


return self.bijector.scale<EOL>

Dense (batch) covariance matrix, if available.


parameters = dict(locals())<EOL>graph_parents = [df, loc, scale_identity_multiplier, scale_diag,<EOL>scale_tril, scale_perturb_factor, scale_perturb_diag]<EOL>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>dtype = dtype_util.common_dtype(graph_parents, tf.float32)<EOL>df = tf.convert_to_tensor(value=df, name="<STR_LIT>", dtype=dtype)<EOL>affine = affine_bijector.Affine(<EOL>shift=loc,<EOL>scale_identity_multiplier=scale_identity_multiplier,<EOL>scale_diag=scale_diag,<EOL>scale_tril=scale_tril,<EOL>scale_perturb_factor=scale_perturb_factor,<EOL>scale_perturb_diag=scale_perturb_diag,<EOL>validate_args=validate_args,<EOL>dtype=dtype)<EOL>distribution = student_t.StudentT(<EOL>df=df,<EOL>loc=tf.zeros([], dtype=affine.dtype),<EOL>scale=tf.ones([], dtype=affine.dtype))<EOL>batch_shape, override_event_shape = (<EOL>distribution_util.shapes_from_loc_and_scale(<EOL>affine.shift, affine.scale))<EOL>override_batch_shape = distribution_util.pick_vector(<EOL>distribution.is_scalar_batch(), batch_shape,<EOL>tf.constant([], dtype=tf.int32))<EOL>super(_VectorStudentT, self).__init__(<EOL>distribution=distribution,<EOL>bijector=affine,<EOL>batch_shape=override_batch_shape,<EOL>event_shape=override_event_shape,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>self._parameters = parameters<EOL><DEDENT><DEDENT>

Instantiates the vector Student's t-distributions on `R^k`.

        The `batch_shape` is the broadcast between `df.batch_shape` and
        `Affine.batch_shape` where `Affine` is constructed from `loc` and
        `scale_*` arguments.

        The `event_shape` is the event shape of `Affine.event_shape`.

        Args:
          df: Floating-point `Tensor`. The degrees of freedom of the
            distribution(s). `df` must contain only positive values. Must be
            scalar if `loc`, `scale_*` imply non-scalar batch_shape or must have the
            same `batch_shape` implied by `loc`, `scale_*`.
          loc: Floating-point `Tensor`. If this is set to `None`, no `loc` is
            applied.
          scale_identity_multiplier: floating point rank 0 `Tensor` representing a
            scaling done to the identity matrix. When `scale_identity_multiplier =
            scale_diag=scale_tril = None` then `scale += IdentityMatrix`. Otherwise
            no scaled-identity-matrix is added to `scale`.
          scale_diag: Floating-point `Tensor` representing the diagonal matrix.
            `scale_diag` has shape [N1, N2, ..., k], which represents a k x k
            diagonal matrix. When `None` no diagonal term is added to `scale`.
          scale_tril: Floating-point `Tensor` representing the diagonal matrix.
            `scale_diag` has shape [N1, N2, ..., k, k], which represents a k x k
            lower triangular matrix. When `None` no `scale_tril` term is added to
            `scale`. The upper triangular elements above the diagonal are ignored.
          scale_perturb_factor: Floating-point `Tensor` representing factor matrix
            with last two dimensions of shape `(k, r)`. When `None`, no rank-r
            update is added to `scale`.
          scale_perturb_diag: Floating-point `Tensor` representing the diagonal
            matrix. `scale_perturb_diag` has shape [N1, N2, ..., r], which
            represents an r x r Diagonal matrix. When `None` low rank updates will
            take the form `scale_perturb_factor * scale_perturb_factor.T`.
          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.


return self.bijector.shift<EOL>

Locations of these Student's t distribution(s).


def normal_conjugates_known_scale_predictive(prior, scale, s, n):

if not isinstance(prior, normal.Normal):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if s.dtype != prior.dtype:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>"<EOL>% (s.dtype, prior.dtype))<EOL><DEDENT>n = tf.cast(n, prior.dtype)<EOL>scale0_2 = tf.square(prior.scale)<EOL>scale_2 = tf.square(scale)<EOL>scalep_2 = <NUM_LIT:1.0>/(<NUM_LIT:1>/scale0_2 + n/scale_2)<EOL>return normal.Normal(<EOL>loc=(prior.loc / scale0_2 + s / scale_2) * scalep_2,<EOL>scale=tf.sqrt(scalep_2 + scale_2))<EOL>

Posterior predictive Normal distribution w. conjugate prior on the mean.

    This model assumes that `n` observations (with sum `s`) come from a
    Normal with unknown mean `loc` (described by the Normal `prior`)
    and known variance `scale**2`. The "known scale predictive"
    is the distribution of new observations, conditioned on the existing
    observations and our prior.

    Accepts a prior Normal distribution object, having parameters
    `loc0` and `scale0`, as well as known `scale` values of the predictive
    distribution(s) (also assumed Normal),
    and statistical estimates `s` (the sum(s) of the observations) and
    `n` (the number(s) of observations).

    Calculates the Normal distribution(s) `p(x | sigma**2)`:

    ```
    p(x | sigma**2) = int N(x | mu, sigma**2)N(mu | prior.loc, prior.scale**2) dmu
                    = N(x | prior.loc, 1 / (sigma**2 + prior.scale**2))
    ```

    Returns the predictive posterior distribution object, with parameters
    `(loc', scale'**2)`, where:

    ```
    sigma_n**2 = 1/(1/sigma0**2 + n/sigma**2),
    mu' = (mu0/sigma0**2 + s/sigma**2) * sigma_n**2.
    sigma'**2 = sigma_n**2 + sigma**2,
    ```

    Distribution parameters from `prior`, as well as `scale`, `s`, and `n`.
    will broadcast in the case of multidimensional sets of parameters.

    Args:
      prior: `Normal` object of type `dtype`:
        the prior distribution having parameters `(loc0, scale0)`.
      scale: tensor of type `dtype`, taking values `scale > 0`.
        The known stddev parameter(s).
      s: Tensor of type `dtype`. The sum(s) of observations.
      n: Tensor of type `int`. The number(s) of observations.

    Returns:
      A new Normal predictive distribution object.

    Raises:
      TypeError: if dtype of `s` does not match `dtype`, or `prior` is not a
        Normal object.


return self._tailweight<EOL>

Controls the tail decay.  `tailweight > 1` means faster than Normal.


return self._scale<EOL>

Scaling factors of these Student's t distribution(s).


return tf.math.rsqrt(<NUM_LIT:1.> + self.total_concentration[..., tf.newaxis])<EOL>

Helper to `_covariance` and `_variance` which computes a shared scale.


if not self.validate_args:<EOL><INDENT>return x<EOL><DEDENT>return distribution_util.with_dependencies([<EOL>assert_util.assert_positive(x, message="<STR_LIT>"),<EOL>assert_util.assert_near(<EOL>tf.ones([], dtype=self.dtype),<EOL>tf.reduce_sum(input_tensor=x, axis=-<NUM_LIT:1>),<EOL>message="<STR_LIT>"),<EOL>], x)<EOL>


return self._concentration<EOL>

Concentration parameter; expected counts for that coordinate.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([<EOL>index_points, observation_index_points, observations,<EOL>observation_noise_variance, predictive_noise_variance, jitter<EOL>], 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>observation_index_points = (None if observation_index_points is None else<EOL>tf.convert_to_tensor(<EOL>value=observation_index_points,<EOL>dtype=dtype,<EOL>name='<STR_LIT>'))<EOL>observations = (None if observations is None else tf.convert_to_tensor(<EOL>value=observations, 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>predictive_noise_variance = (<EOL>observation_noise_variance<EOL>if predictive_noise_variance is None else tf.convert_to_tensor(<EOL>value=predictive_noise_variance,<EOL>dtype=dtype,<EOL>name='<STR_LIT>'))<EOL>jitter = tf.convert_to_tensor(value=jitter, dtype=dtype, name='<STR_LIT>')<EOL>if (observation_index_points is None) != (observations is None):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>observations, observation_index_points))<EOL><DEDENT>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._name = name<EOL>self._observation_index_points = observation_index_points<EOL>self._observations = observations<EOL>self._observation_noise_variance = observation_noise_variance<EOL>self._predictive_noise_variance = predictive_noise_variance<EOL>self._jitter = jitter<EOL>self._validate_args = validate_args<EOL>with tf.name_scope('<STR_LIT>'):<EOL><INDENT>conditional_kernel = tfpk.SchurComplement(<EOL>base_kernel=kernel,<EOL>fixed_inputs=observation_index_points,<EOL>diag_shift=jitter + observation_noise_variance)<EOL>if _is_empty_observation_data(<EOL>feature_ndims=kernel.feature_ndims,<EOL>observation_index_points=observation_index_points,<EOL>observations=observations):<EOL><INDENT>conditional_mean_fn = mean_fn<EOL><DEDENT>else:<EOL><INDENT>_validate_observation_data(<EOL>kernel=kernel,<EOL>observation_index_points=observation_index_points,<EOL>observations=observations)<EOL>def conditional_mean_fn(x):<EOL><INDENT>k_x_obs_linop = tf.linalg.LinearOperatorFullMatrix(<EOL>kernel.matrix(x, observation_index_points))<EOL>chol_linop = tf.linalg.LinearOperatorLowerTriangular(<EOL>conditional_kernel.divisor_matrix_cholesky)<EOL>diff = observations - mean_fn(observation_index_points)<EOL>return mean_fn(x) + k_x_obs_linop.matvec(<EOL>chol_linop.solvevec(chol_linop.solvevec(diff), adjoint=True))<EOL><DEDENT><DEDENT>graph_parents = [observation_noise_variance, jitter]<EOL>def _maybe_append(x):<EOL><INDENT>if x is not None:<EOL><INDENT>graph_parents.append(x)<EOL><DEDENT><DEDENT>_maybe_append(index_points)<EOL>_maybe_append(observation_index_points)<EOL>_maybe_append(observations)<EOL>super(GaussianProcessRegressionModel, self).__init__(<EOL>kernel=conditional_kernel,<EOL>mean_fn=conditional_mean_fn,<EOL>index_points=index_points,<EOL>jitter=jitter,<EOL>observation_noise_variance=predictive_noise_variance,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats, name=name)<EOL>self._parameters = parameters<EOL>self._graph_parents = graph_parents<EOL><DEDENT><DEDENT>

Construct a GaussianProcessRegressionModel instance.

        Args:
          kernel: `PositiveSemidefiniteKernel`-like instance representing the
            GP's covariance function.
          index_points: `float` `Tensor` representing finite collection, or batch of
            collections, 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 an `e`-dimensional multivariate normal. The
            batch shape must be broadcastable with `kernel.batch_shape` and any
            batch dims yielded by `mean_fn`.
          observation_index_points: `float` `Tensor` representing finite collection,
            or batch of collections, of points in the index set for which some data
            has been observed. 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. `[b1, ..., bB, e]` must be broadcastable with the shape of
            `observations`, and `[b1, ..., bB]` must be broadcastable with the
            shapes of all other batched parameters (`kernel.batch_shape`,
            `index_points`, etc). The default value is `None`, which corresponds to
            the empty set of observations, and simply results in the prior
            predictive model (a GP with noise of variance
            `predictive_noise_variance`).
          observations: `float` `Tensor` representing collection, or batch of
            collections, of observations corresponding to
            `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
            must be brodcastable with the batch and example shapes of
            `observation_index_points`. The batch shape `[b1, ..., bB]` must be
            broadcastable with the shapes of all other batched parameters
            (`kernel.batch_shape`, `index_points`, etc.). The default value is
            `None`, which corresponds to the empty set of observations, and simply
            results in the prior predictive model (a GP with noise of variance
            `predictive_noise_variance`).
          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.`
          predictive_noise_variance: `float` `Tensor` representing the variance in
            the posterior predictive model. If `None`, we simply re-use
            `observation_noise_variance` for the posterior predictive noise. If set
            explicitly, however, we use this value. This allows us, for example, to
            omit predictive noise variance (by setting this to zero) to obtain
            noiseless posterior predictions of function values, conditioned on noisy
            observations.
          mean_fn: Python `callable` that acts on `index_points` to produce a
            collection, or batch of collections, 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 the constant zero function.
          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: 'GaussianProcessRegressionModel'.

        Raises:
          ValueError: if either
            - only one of `observations` and `observation_index_points` is given, or
            - `mean_fn` is not `None` and not callable.


def _validate_observation_data(<EOL>kernel, observation_index_points, observations):

<EOL>ndims = kernel.feature_ndims<EOL>if (tensorshape_util.is_fully_defined(<EOL>observation_index_points.shape[:-ndims]) and<EOL>tensorshape_util.is_fully_defined(observations.shape)):<EOL><INDENT>index_point_count = observation_index_points.shape[:-ndims]<EOL>observation_count = observations.shape<EOL>try:<EOL><INDENT>tf.broadcast_static_shape(index_point_count, observation_count)<EOL><DEDENT>except ValueError:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>index_point_count, observation_count))<EOL><DEDENT><DEDENT>

Ensure that observation data and locations have consistent shapes.

    This basically means that the batch shapes are broadcastable. We can only
    ensure this when those shapes are fully statically defined.

    Args:
      kernel: The GP kernel.
      observation_index_points: the observation data locations in the index set.
      observations: the observation data.

    Raises:
      ValueError: if the observations' batch shapes are not broadcastable.


seed = seed_stream.SeedStream(seed, salt='<STR_LIT>')<EOL>u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=<NUM_LIT:0>)<EOL>z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype)<EOL>safe_conc = tf.where(self.concentration > <NUM_LIT:0>,<EOL>self.concentration,<EOL>tf.ones_like(self.concentration))<EOL>safe_z = tf.where(z > <NUM_LIT:0>, z, tf.ones_like(z))<EOL>safe_u = <NUM_LIT:1> + tf.reduce_logsumexp(<EOL>input_tensor=[<EOL>tf.math.log(safe_z),<EOL>tf.math.log1p(-safe_z) - <NUM_LIT:2> * safe_conc<EOL>],<EOL>axis=<NUM_LIT:0>) / safe_conc<EOL>u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u,<EOL><NUM_LIT:2> * z - <NUM_LIT:1>)<EOL>u = tf.where(tf.equal(z, <NUM_LIT:0>), -tf.ones_like(u), u)<EOL>if not self._allow_nan_stats:<EOL><INDENT>u = tf.debugging.check_numerics(u, '<STR_LIT>')<EOL><DEDENT>return u[..., tf.newaxis]<EOL>


Initialize a `VectorDeterministic` distribution on `R^k`, for `k >= 0`.

        Note that there is only one point in `R^0`, the "point" `[]`.  So if `k = 0`
        then `self.prob([]) == 1`.

        The `atol` and `rtol` parameters allow for some slack in `pmf`
        computations, e.g. due to floating-point error.

        ```
        pmf(x; loc)
          = 1, if All[Abs(x - loc) <= atol + rtol * Abs(loc)],
          = 0, otherwise
        ```

        Args:
          loc: Numeric `Tensor` of shape `[B1, ..., Bb, k]`, with `b >= 0`, `k >= 0`
            The point (or batch of points) on which this distribution is supported.
          atol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
            shape.  The absolute tolerance for comparing closeness to `loc`.
            Default is `0`.
          rtol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
            shape.  The relative tolerance for comparing closeness to `loc`.
            Default is `0`.
          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.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([logits, probs, temperature], tf.float32)<EOL>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>name=name,<EOL>logits=logits,<EOL>probs=probs,<EOL>validate_args=validate_args,<EOL>multidimensional=True,<EOL>dtype=dtype)<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_positive(temperature)] if validate_args else []):<EOL><INDENT>self._temperature = tf.convert_to_tensor(<EOL>value=temperature, name="<STR_LIT>", dtype=dtype)<EOL>self._temperature_2d = tf.reshape(<EOL>self._temperature, [-<NUM_LIT:1>, <NUM_LIT:1>], name="<STR_LIT>")<EOL><DEDENT>logits_shape_static = tensorshape_util.with_rank_at_least(<EOL>self._logits.shape, <NUM_LIT:1>)<EOL>if tensorshape_util.rank(logits_shape_static) is not None:<EOL><INDENT>self._batch_rank = tf.convert_to_tensor(<EOL>value=tensorshape_util.rank(logits_shape_static) - <NUM_LIT:1>,<EOL>dtype=tf.int32,<EOL>name="<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>self._batch_rank = tf.rank(self._logits) - <NUM_LIT:1><EOL><DEDENT><DEDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>self._event_size = tf.shape(input=self._logits)[-<NUM_LIT:1>]<EOL><DEDENT><DEDENT>super(ExpRelaxedOneHotCategorical, 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=[self._logits, self._probs, self._temperature],<EOL>name=name)<EOL>

Initialize ExpRelaxedOneHotCategorical using class log-probabilities.

        Args:
          temperature: An 0-D `Tensor`, representing the temperature
            of a set of ExpRelaxedCategorical distributions. The temperature should
            be positive.
          logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities
            of a set of ExpRelaxedCategorical distributions. The first
            `N - 1` dimensions index into a batch of independent distributions and
            the last dimension represents a vector of logits for each class. Only
            one of `logits` or `probs` should be passed in.
          probs: An N-D `Tensor`, `N >= 1`, representing the probabilities
            of a set of ExpRelaxedCategorical distributions. The first
            `N - 1` dimensions index into a batch of independent distributions and
            the last dimension represents a vector of probabilities for each
            class. Only one of `logits` or `probs` should be passed in.
          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.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([concentration, rate], tf.float32)<EOL>concentration = tf.convert_to_tensor(<EOL>value=concentration, name="<STR_LIT>", dtype=dtype)<EOL>rate = tf.convert_to_tensor(value=rate, name="<STR_LIT>", dtype=dtype)<EOL>with tf.control_dependencies([<EOL>assert_util.assert_positive(concentration),<EOL>assert_util.assert_positive(rate),<EOL>] if validate_args else []):<EOL><INDENT>self._concentration = tf.identity(concentration)<EOL>self._rate = tf.identity(rate)<EOL>dtype_util.assert_same_float_dtype([self._concentration, self._rate])<EOL><DEDENT><DEDENT>super(Gamma, self).__init__(<EOL>dtype=dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>parameters=parameters,<EOL>graph_parents=[self._concentration, self._rate],<EOL>name=name)<EOL>

Construct Gamma with `concentration` and `rate` parameters.

        The parameters `concentration` and `rate` must be shaped in a way that
        supports broadcasting (e.g. `concentration + rate` is a valid operation).

        Args:
          concentration: Floating point tensor, the concentration params of the
            distribution(s). Must contain only positive values.
          rate: Floating point tensor, the inverse scale params of the
            distribution(s). Must contain only positive values.
          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:
          TypeError: if `concentration` and `rate` are different dtypes.


return self.distribution.scale<EOL>

Distribution parameter for the pre-transformed standard deviation.


if tensorshape_util.rank(x.shape) is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>if isinstance(self._components_distribution, independent.Independent):<EOL><INDENT>univariate_components = self._components_distribution.distribution<EOL><DEDENT>else:<EOL><INDENT>univariate_components = self._components_distribution<EOL><DEDENT>with tf.control_dependencies([<EOL>assert_util.assert_equal(<EOL>univariate_components.is_scalar_event(),<EOL>True,<EOL>message="<STR_LIT>")<EOL>]):<EOL><INDENT>x_padded = self._pad_sample_dims(x)  <EOL>log_prob_x = univariate_components.log_prob(x_padded)  <EOL>cdf_x = univariate_components.cdf(x_padded)  <EOL>cumsum_log_prob_x = tf.reshape(<EOL>tf.math.cumsum(<EOL>tf.reshape(log_prob_x, [-<NUM_LIT:1>, self._event_size]),<EOL>exclusive=True,<EOL>axis=-<NUM_LIT:1>),<EOL>tf.shape(input=log_prob_x))  <EOL>logits_mix_prob = distribution_utils.pad_mixture_dimensions(<EOL>self.mixture_distribution.logits, self, self.mixture_distribution,<EOL>self._event_ndims)  <EOL>log_posterior_weights_x = logits_mix_prob + cumsum_log_prob_x<EOL>component_axis = tensorshape_util.rank(x.shape) - self._event_ndims<EOL>posterior_weights_x = tf.nn.softmax(log_posterior_weights_x,<EOL>axis=component_axis)<EOL>return tf.reduce_sum(<EOL>input_tensor=posterior_weights_x * cdf_x, axis=component_axis)<EOL><DEDENT>

Performs distributional transform of the mixture samples.

        Distributional transform removes the parameters from samples of a
        multivariate distribution by applying conditional CDFs:
          (F(x_1), F(x_2 | x1_), ..., F(x_d | x_1, ..., x_d-1))
        (the indexing is over the "flattened" event dimensions).
        The result is a sample of product of Uniform[0, 1] distributions.

        We assume that the components are factorized, so the conditional CDFs become
          F(x_i | x_1, ..., x_i-1) = sum_k w_i^k F_k (x_i),
        where w_i^k is the posterior mixture weight: for i > 0
          w_i^k = w_k prob_k(x_1, ..., x_i-1) / sum_k' w_k' prob_k'(x_1, ..., x_i-1)
        and w_0^k = w_k is the mixture probability of the k-th component.

        Arguments:
          x: Sample of mixture distribution

        Returns:
          Result of the distributional transform


def grad(dy):<EOL><INDENT>return array_ops.prevent_gradient(<EOL>dy, message="<STR_LIT>")<EOL><DEDENT>return tf.identity(x), grad<EOL>

Disables computation of the second derivatives for a tensor.

    NB: you need to apply a non-identity function to the output tensor for the
    exception to be raised.

    Arguments:
      x: A tensor.

    Returns:
      A tensor with the same value and the same derivative as x, but that raises
      LookupError when trying to compute the second derivatives.


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.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale], tf.float32)<EOL>if loc is not None:<EOL><INDENT>loc = tf.convert_to_tensor(value=loc, name="<STR_LIT>", dtype=dtype)<EOL><DEDENT>if scale is not None:<EOL><INDENT>scale = tf.convert_to_tensor(value=scale, dtype=dtype, name="<STR_LIT>")<EOL><DEDENT>self._quadrature_grid, self._quadrature_probs = tuple(quadrature_fn(<EOL>loc, scale, quadrature_size, validate_args))<EOL>dt = self._quadrature_grid.dtype<EOL>if not dtype_util.base_equal(dt, self._quadrature_probs.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>dtype_util.name(dt),<EOL>dtype_util.name(self._quadrature_probs.dtype)))<EOL><DEDENT>self._distribution = poisson.Poisson(<EOL>log_rate=self._quadrature_grid,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats)<EOL>self._mixture_distribution = categorical.Categorical(<EOL>logits=tf.math.log(self._quadrature_probs),<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats)<EOL>self._loc = loc<EOL>self._scale = scale<EOL>self._quadrature_size = quadrature_size<EOL>super(PoissonLogNormalQuadratureCompound, self).__init__(<EOL>dtype=dt,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[loc, scale],<EOL>name=name)<EOL><DEDENT>

Constructs the PoissonLogNormalQuadratureCompound`.

        Note: `probs` returned by (optional) `quadrature_fn` are presumed to be
        either a length-`quadrature_size` vector or a batch of vectors in 1-to-1
        correspondence with the returned `grid`. (I.e., broadcasting is only
        partially supported.)

        Args:
          loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
            the LogNormal prior.
          scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
            the LogNormal prior.
          quadrature_size: Python `int` scalar representing the number of quadrature
            points.
          quadrature_fn: Python callable taking `loc`, `scale`,
            `quadrature_size`, `validate_args` and returning `tuple(grid, probs)`
            representing the LogNormal grid and corresponding normalized weight.
            normalized) weight.
            Default value: `quadrature_scheme_lognormal_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:
          TypeError: if `quadrature_grid` and `quadrature_probs` have different base
            `dtype`.


return self._scale<EOL>


expanded_x = x<EOL>for _ in range(tensorshape_util.rank(self.event_shape)):<EOL><INDENT>expanded_x = tf.expand_dims(expanded_x, -<NUM_LIT:1>)<EOL><DEDENT>return expanded_x<EOL>

Expand the rank of x up to static_event_rank times for broadcasting.

        The static event rank was checked to not be None at construction time.

        Args:
          x: A tensor to expand.
        Returns:
          The expanded tensor.


with self._name_scope(name):<EOL><INDENT>with tf.control_dependencies(self._assertions):<EOL><INDENT>distribution_entropies = [d.entropy() for d in self.components]<EOL>cat_probs = self._cat_probs(log_probs=False)<EOL>partial_entropies = [<EOL>c_p * m for (c_p, m) in zip(cat_probs, distribution_entropies)<EOL>]<EOL>return tf.add_n(partial_entropies)<EOL><DEDENT><DEDENT>

r"""A lower bound on the entropy of this mixture model.

        The bound below is not always very tight, and its usefulness depends
        on the mixture probabilities and the components in use.

        A lower bound is useful for ELBO when the `Mixture` is the variational
        distribution:

        \\(
        \log p(x) >= ELBO = \int q(z) \log p(x, z) dz + H[q]
        \\)

        where \\( p \\) is the prior distribution, \\( q \\) is the variational,
        and \\( H[q] \\) is the entropy of \\( q \\). If there is a lower bound
        \\( G[q] \\) such that \\( H[q] \geq G[q] \\) then it can be used in
        place of \\( H[q] \\).

        For a mixture of distributions \\( q(Z) = \sum_i c_i q_i(Z) \\) with
        \\( \sum_i c_i = 1 \\), by the concavity of \\( f(x) = -x \log x \\), a
        simple lower bound is:

        \\(
        \begin{align}
        H[q] & = - \int q(z) \log q(z) dz \\\
           & = - \int (\sum_i c_i q_i(z)) \log(\sum_i c_i q_i(z)) dz \\\
           & \geq - \sum_i c_i \int q_i(z) \log q_i(z) dz \\\
           & = \sum_i c_i H[q_i]
        \end{align}
        \\)

        This is the term we calculate below for \\( G[q] \\).

        Args:
          name: A name for this operation (optional).

        Returns:
          A lower bound on the Mixture's entropy.


shape = tensorshape_util.concatenate(self.batch_shape, self.event_shape)<EOL>has_static_shape = tensorshape_util.is_fully_defined(shape)<EOL>if not has_static_shape:<EOL><INDENT>shape = tf.concat([<EOL>self.batch_shape_tensor(),<EOL>self.event_shape_tensor(),<EOL>], <NUM_LIT:0>)<EOL><DEDENT>return shape<EOL>


parameters = dict(locals())<EOL>def _convert_to_tensor(x, name, dtype=None):<EOL><INDENT>return None if x is None else tf.convert_to_tensor(<EOL>value=x, name=name, dtype=dtype)<EOL><DEDENT>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>dtype = dtype_util.common_dtype([<EOL>loc, scale_diag, scale_identity_multiplier, scale_perturb_factor,<EOL>scale_perturb_diag<EOL>], tf.float32)<EOL>has_low_rank = (scale_perturb_factor is not None or<EOL>scale_perturb_diag is not None)<EOL>scale = distribution_util.make_diag_scale(<EOL>loc=loc,<EOL>scale_diag=scale_diag,<EOL>scale_identity_multiplier=scale_identity_multiplier,<EOL>validate_args=validate_args,<EOL>assert_positive=has_low_rank,<EOL>dtype=dtype)<EOL>scale_perturb_factor = _convert_to_tensor(<EOL>scale_perturb_factor, name="<STR_LIT>", dtype=dtype)<EOL>scale_perturb_diag = _convert_to_tensor(<EOL>scale_perturb_diag, name="<STR_LIT>", dtype=dtype)<EOL>if has_low_rank:<EOL><INDENT>scale = tf.linalg.LinearOperatorLowRankUpdate(<EOL>scale,<EOL>u=scale_perturb_factor,<EOL>diag_update=scale_perturb_diag,<EOL>is_diag_update_positive=scale_perturb_diag is None,<EOL>is_non_singular=True,  <EOL>is_self_adjoint=True,<EOL>is_positive_definite=True,<EOL>is_square=True)<EOL><DEDENT><DEDENT><DEDENT>super(MultivariateNormalDiagPlusLowRank, 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 = diag(scale_diag + scale_identity_multiplier ones(k)) +
            scale_perturb_factor @ diag(scale_perturb_diag) @ scale_perturb_factor.T
        ```

        where:

        * `scale_diag.shape = [k]`,
        * `scale_identity_multiplier.shape = []`,
        * `scale_perturb_factor.shape = [k, r]`, typically `k >> r`, and,
        * `scale_perturb_diag.shape = [r]`.

        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`.
          scale_perturb_factor: Floating-point `Tensor` representing a rank-`r`
            perturbation added to `scale`. May have shape `[B1, ..., Bb, k, r]`,
            `b >= 0`, and characterizes `b`-batches of rank-`r` updates to `scale`.
            When `None`, no rank-`r` update is added to `scale`.
          scale_perturb_diag: Floating-point `Tensor` representing a diagonal matrix
            inside the rank-`r` perturbation added to `scale`. May have shape
            `[B1, ..., Bb, r]`, `b >= 0`, and characterizes `b`-batches of `r x r`
            diagonal matrices inside the perturbation added to `scale`. When
            `None`, an identity matrix is used inside the perturbation. Can only be
            specified if `scale_perturb_factor` is also specified.
          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.


def _build_iid_normal_model(self, num_timesteps, latent_size,<EOL>observation_size, transition_variance,<EOL>observation_variance):

transition_variance = self._build_placeholder(transition_variance)<EOL>observation_variance = self._build_placeholder(observation_variance)<EOL>random_orthogonal_matrix = lambda: np.linalg.qr(<EOL>np.random.randn(latent_size, latent_size))[<NUM_LIT:0>][:observation_size, :]<EOL>observation_matrix = self._build_placeholder(random_orthogonal_matrix())<EOL>model = tfd.LinearGaussianStateSpaceModel(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=self._build_placeholder(<EOL>np.zeros([latent_size, latent_size])),<EOL>transition_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.sqrt(transition_variance) *<EOL>tf.ones([latent_size], dtype=self.dtype)),<EOL>observation_matrix=observation_matrix,<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.sqrt(observation_variance) *<EOL>tf.ones([observation_size], dtype=self.dtype)),<EOL>initial_state_prior=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.sqrt(transition_variance) *<EOL>tf.ones([latent_size], dtype=self.dtype)),<EOL>validate_args=True)<EOL>return model<EOL>

Build a model whose outputs are IID normal by construction.


if not (isinstance(ndims, int) and ndims >= <NUM_LIT:0>):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>.format(ndims))<EOL><DEDENT>if ndims == <NUM_LIT:0>:<EOL><INDENT>return x<EOL><DEDENT>x = tf.convert_to_tensor(value=x)<EOL>original_shape = x.shape<EOL>new_shape = distribution_util.pad(<EOL>tf.shape(input=x), axis=<NUM_LIT:0>, back=True, value=<NUM_LIT:1>, count=ndims)<EOL>x = tf.reshape(x, new_shape)<EOL>x.set_shape(original_shape.concatenate([<NUM_LIT:1>]*ndims))<EOL>return x<EOL>

Maybe add `ndims` ones to `x.shape` on the right.

    If `ndims` is zero, this is a no-op; otherwise, we will create and return a
    new `Tensor` whose shape is that of `x` with `ndims` ones concatenated on the
    right side. If the shape of `x` is known statically, the shape of the return
    value will be as well.

    Args:
      x: The `Tensor` we'll return a reshaping of.
      ndims: Python `integer` number of ones to pad onto `x.shape`.
    Returns:
      If `ndims` is zero, `x`; otherwise, a `Tensor` whose shape is that of `x`
      with `ndims` ones concatenated on the right side. If possible, returns a
      `Tensor` whose shape is known statically.
    Raises:
      ValueError: if `ndims` is not a Python `integer` greater than or equal to
      zero.


if not kernels:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if len(set([k.feature_ndims for k in kernels])) > <NUM_LIT:1>:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>" %<EOL>str([k.feature_ndims for k in kernels]))<EOL><DEDENT>self._kernels = _flatten_summand_list(kernels)<EOL>if name is None:<EOL><INDENT>name = '<STR_LIT>'<EOL><DEDENT>super(_SumKernel, self).__init__(<EOL>feature_ndims=kernels[<NUM_LIT:0>].feature_ndims,<EOL>dtype=util.maybe_get_common_dtype(<EOL>[None if k.dtype is None else k for k in kernels]),<EOL>name=name)<EOL>

Create a kernel which is the sum of `kernels`.

        The input list is 'flattened' in the sense that any entries which are also
        of type `_SumKernel` will have their list of kernels appended to this
        instance's list of kernels. This will reduce the stack depth when actually
        evaluating the sum over kernel applications.

        Args:
          kernels: Python `list` of `PositiveSemidefiniteKernel` instances.
          name: Python `str` name prefixed to Ops created by this class.
        Raises:
          ValueError: `kernels` is an empty list, or `kernels` don't all have the
          same `feature_ndims`.


def __init__(self, feature_ndims, dtype=None, name=None):

if not (isinstance(feature_ndims, int) and feature_ndims > <NUM_LIT:0>):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>' +<EOL>'<STR_LIT>'.format(feature_ndims))<EOL><DEDENT>self._feature_ndims = feature_ndims<EOL>self._dtype = dtype<EOL>if not name or name[-<NUM_LIT:1>] != '<STR_LIT:/>':  <EOL><INDENT>name = tf.compat.v1.name_scope(name or type(self).__name__).name<EOL><DEDENT>self._name = name<EOL>

Construct a PositiveSemidefiniteKernel (subclass) instance.

        Args:
          feature_ndims: Python `integer` indicating the number of dims (the rank)
            of the feature space this kernel acts on.
          dtype: `DType` on which this kernel operates.
          name: Python `str` name prefixed to Ops created by this class. Default:
            subclass name.

        Raises:
          ValueError: if `feature_ndims` is not an integer greater than 0
        Inputs to PositiveSemidefiniteKernel methods partition into 3 pieces:

        ```none
        [b1, ..., bB, e, f1, ..., fF]
        '----------'  |  '---------'
             |        |       '-- Feature dimensions
             |        '-- Example dimension (`matrix`-only)
             '-- Batch dimensions
        ```

        The `feature_ndims` argument declares how many of the right-most shape
        dimensions belong to the feature dimensions. This enables us to predict
        which shape dimensions will be 'reduced' away during kernel computation.