def posterior_mode(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>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_prob = log_init + observation_log_probs[<NUM_LIT:0>]<EOL>if self._num_steps == <NUM_LIT:1>:<EOL><INDENT>most_likely_end = tf.argmax(input=log_prob, axis=-<NUM_LIT:1>)<EOL>return most_likely_end[..., tf.newaxis]<EOL><DEDENT>def forward_step(previous_step_pair, log_prob_observation):<EOL><INDENT>log_prob_previous = previous_step_pair[<NUM_LIT:0>]<EOL>log_prob = (log_prob_previous[..., tf.newaxis] +<EOL>self._log_trans +<EOL>log_prob_observation[..., tf.newaxis, :])<EOL>most_likely_given_successor = tf.argmax(input=log_prob, axis=-<NUM_LIT:2>)<EOL>max_log_p_given_successor = tf.reduce_max(input_tensor=log_prob,<EOL>axis=-<NUM_LIT:2>)<EOL>return (max_log_p_given_successor, most_likely_given_successor)<EOL><DEDENT>forward_log_probs, all_most_likely_given_successor = tf.scan(<EOL>forward_step,<EOL>observation_log_probs[<NUM_LIT:1>:],<EOL>initializer=(log_prob,<EOL>tf.zeros(tf.shape(input=log_init), dtype=tf.int64)),<EOL>name="<STR_LIT>")<EOL>most_likely_end = tf.argmax(input=forward_log_probs[-<NUM_LIT:1>], axis=-<NUM_LIT:1>)<EOL>def backward_step(most_likely_successor, most_likely_given_successor):<EOL><INDENT>return tf.reduce_sum(<EOL>input_tensor=(most_likely_given_successor *<EOL>tf.one_hot(most_likely_successor,<EOL>self._num_states,<EOL>dtype=tf.int64)),<EOL>axis=-<NUM_LIT:1>)<EOL><DEDENT>backward_scan = tf.scan(<EOL>backward_step,<EOL>all_most_likely_given_successor,<EOL>most_likely_end,<EOL>reverse=True)<EOL>most_likely_sequences = tf.concat([backward_scan, [most_likely_end]],<EOL>axis=<NUM_LIT:0>)<EOL>return distribution_util.move_dimension(most_likely_sequences, <NUM_LIT:0>, -<NUM_LIT:1>)<EOL><DEDENT><DEDENT><DEDENT>

Compute maximum likelihood sequence of hidden states.

        When this function is provided with a sequence of observations
        `x[0], ..., x[num_steps - 1]`, it returns the sequence of hidden
        states `z[0], ..., z[num_steps - 1]`, drawn from the underlying
        Markov chain, that is most likely to yield those observations.

        It uses the [Viterbi algorithm](
        https://en.wikipedia.org/wiki/Viterbi_algorithm).

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

        Note: if there isn't a unique most likely sequence then one
        of the equally most likely sequences is chosen.

        Args:
          observations: A tensor representing a batch of observations made on the
            hidden Markov model.  The rightmost dimensions of this tensor correspond
            to the dimensions of the observation distributions of the underlying
            Markov chain.  The next dimension from the right indexes 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_mode: A `Tensor` representing the most likely sequence of hidden
            states. The rightmost dimension of this tensor will equal the
            `num_steps` parameter providing one hidden state for each step. The
            other dimensions are those of the batch.

        Raises:
          ValueError: if the `observations` tensor does not consist of
          sequences of `num_steps` observations.

        #### Examples

        ```python
        tfd = tfp.distributions

        # A simple weather model.

        # Represent a cold day with 0 and a hot day with 1.
        # Suppose the first day of a sequence has a 0.8 chance of being cold.

        initial_distribution = tfd.Categorical(probs=[0.8, 0.2])

        # Suppose a cold day has a 30% chance of being followed by a hot day
        # and a hot day has a 20% chance of being followed by a cold day.

        transition_distribution = tfd.Categorical(probs=[[0.7, 0.3],
                                                         [0.2, 0.8]])

        # Suppose additionally that on each day the temperature is
        # normally distributed with mean and standard deviation 0 and 5 on
        # a cold day and mean and standard deviation 15 and 10 on a hot day.

        observation_distribution = tfd.Normal(loc=[0., 15.], scale=[5., 10.])

        # This gives the hidden Markov model:

        model = tfd.HiddenMarkovModel(
            initial_distribution=initial_distribution,
            transition_distribution=transition_distribution,
            observation_distribution=observation_distribution,
            num_steps=7)

        # Suppose we observe gradually rising temperatures over a week:
        temps = [-2., 0., 2., 4., 6., 8., 10.]

        # We can now compute the most probable sequence of hidden states:

        model.posterior_mode(temps)

        # The result is [0 0 0 0 0 1 1] telling us that the transition
        # from "cold" to "hot" most likely happened between the
        # 5th and 6th days.
        ```


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale], 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)<EOL>self._scale = tf.identity(scale)<EOL>dtype_util.assert_same_float_dtype([self._loc, self._scale])<EOL><DEDENT><DEDENT>super(Cauchy, 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 Cauchy distributions.

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

        Args:
          loc: Floating point tensor; the modes of the distribution(s).
          scale: Floating point tensor; the locations 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 `loc` and `scale` have different `dtype`.


return self._loc<EOL>

Distribution parameter for the location.


return self._scale<EOL>


with tf.name_scope("<STR_LIT>"):<EOL><INDENT>return (x - self.loc) / self.scale<EOL><DEDENT>


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

Reconstruct input `x` from a its normalized version.


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.distribution.df<EOL>

Degrees of freedom in these Student's t distribution(s).


return self.bijector.shift<EOL>

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


return self.bijector.scale<EOL>

Dense (batch) covariance matrix, if available.


def normal_conjugates_known_scale_posterior(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))<EOL>

Posterior Normal distribution with 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 posterior" is
    the distribution of the unknown `loc`.

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

    Returns a posterior (also Normal) distribution object, with parameters
    `(loc', scale'**2)`, where:

    ```
    mu ~ N(mu', sigma'**2)
    sigma'**2 = 1/(1/sigma0**2 + n/sigma**2),
    mu' = (mu0/sigma0**2 + s/sigma**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 posterior distribution object for the unknown observation
      mean `loc`.

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


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.


def scipy_trunc_norm_dist(loc, scale, low, high):

a = (low - loc) / scale<EOL>b = (high - loc) / scale<EOL>return stats.truncnorm(a, b, loc=loc, scale=scale)<EOL>

Construct a scipy.stats.truncnorm for the (scalar) parameters given.

    Note: scipy's definition of the parameters is slightly different.
    https://github.com/scipy/scipy/issues/7591

    Args:
      loc: Params describing distribution (doesn't support batch)
      scale:
      low:
      high:

    Returns:
      scipy frozen distribution.


def assertEmpiricalDistributionsEqual(self, sample_a, sample_b, rtol=<NUM_LIT>,<EOL>atol=<NUM_LIT>):

self.assertAllFinite(sample_a)<EOL>self.assertAllFinite(sample_b)<EOL>lb = np.min([sample_a, sample_b])<EOL>ub = np.max([sample_a, sample_b])<EOL>hist_a = np.histogram(sample_a, range=(lb, ub), bins=<NUM_LIT:30>, density=True)[<NUM_LIT:0>]<EOL>hist_b = np.histogram(sample_b, range=(lb, ub), bins=<NUM_LIT:30>, density=True)[<NUM_LIT:0>]<EOL>self.assertAllClose(hist_a, hist_b, rtol=rtol, atol=atol)<EOL>

Assert the empirical distribution of two set of samples is similar.

        Args:
          sample_a: Flat numpy array of samples from dist a.
          sample_b: Flat numpy array of samples from dist b.
          rtol: Relative tolerances in the histogram comparison.
          atol: Absolute tolerances in the histogram comparison.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[loc, scale_diag, scale_identity_multiplier, skewness, tailweight],<EOL>tf.float32)<EOL>loc = loc if loc is None else tf.convert_to_tensor(<EOL>value=loc, name="<STR_LIT>", dtype=dtype)<EOL>tailweight = <NUM_LIT:1.> if tailweight is None else tailweight<EOL>has_default_skewness = skewness is None<EOL>skewness = <NUM_LIT:0.> if skewness is None else skewness<EOL>scale_linop = distribution_util.make_diag_scale(<EOL>loc=loc,<EOL>scale_diag=scale_diag,<EOL>scale_identity_multiplier=scale_identity_multiplier,<EOL>validate_args=False,<EOL>assert_positive=False,<EOL>dtype=dtype)<EOL>batch_shape, event_shape = distribution_util.shapes_from_loc_and_scale(<EOL>loc, scale_linop)<EOL>scale_diag_part = scale_linop.diag_part()<EOL>dtype = scale_diag_part.dtype<EOL>if distribution is None:<EOL><INDENT>distribution = normal.Normal(<EOL>loc=tf.zeros([], dtype=dtype),<EOL>scale=tf.ones([], dtype=dtype),<EOL>allow_nan_stats=allow_nan_stats)<EOL><DEDENT>else:<EOL><INDENT>asserts = distribution_util.maybe_check_scalar_distribution(<EOL>distribution, dtype, validate_args)<EOL>if asserts:<EOL><INDENT>scale_diag_part = distribution_util.with_dependencies(<EOL>asserts, scale_diag_part)<EOL><DEDENT><DEDENT>skewness = tf.convert_to_tensor(<EOL>value=skewness, dtype=dtype, name="<STR_LIT>")<EOL>tailweight = tf.convert_to_tensor(<EOL>value=tailweight, dtype=dtype, name="<STR_LIT>")<EOL>f = sinh_arcsinh_bijector.SinhArcsinh(<EOL>skewness=skewness, tailweight=tailweight)<EOL>if has_default_skewness:<EOL><INDENT>f_noskew = f<EOL><DEDENT>else:<EOL><INDENT>f_noskew = sinh_arcsinh_bijector.SinhArcsinh(<EOL>skewness=dtype_util.as_numpy_dtype(skewness.dtype)(<NUM_LIT:0.>),<EOL>tailweight=tailweight)<EOL><DEDENT>c = <NUM_LIT:2> * scale_diag_part / f_noskew.forward(<EOL>tf.convert_to_tensor(value=<NUM_LIT:2>, dtype=dtype))<EOL>affine = affine_bijector.Affine(<EOL>shift=loc, scale_diag=c, validate_args=validate_args)<EOL>bijector = chain_bijector.Chain([affine, f])<EOL>super(VectorSinhArcsinhDiag, self).__init__(<EOL>distribution=distribution,<EOL>bijector=bijector,<EOL>batch_shape=batch_shape,<EOL>event_shape=event_shape,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>self._parameters = parameters<EOL>self._loc = loc<EOL>self._scale = scale_linop<EOL>self._tailweight = tailweight<EOL>self._skewness = skewness<EOL>

Construct VectorSinhArcsinhDiag distribution on `R^k`.

        The arguments `scale_diag` and `scale_identity_multiplier` combine to
        define the diagonal `scale` referred to in this class docstring:

        ```none
        scale = diag(scale_diag + scale_identity_multiplier * ones(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

        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_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 scale-identity-matrix added to `scale`. May have shape
            `[B1, ..., Bb]`, `b >= 0`, and characterizes `b`-batches of scale
            `k x k` identity matrices added to `scale`. When both
            `scale_identity_multiplier` and `scale_diag` are `None` then `scale`
            is the `Identity`.
          skewness:  Skewness parameter.  floating-point `Tensor` with shape
            broadcastable with `event_shape`.
          tailweight:  Tailweight parameter.  floating-point `Tensor` with shape
            broadcastable with `event_shape`.
          distribution: `tf.Distribution`-like instance. Distribution from which `k`
            iid samples are used as input to transformation `F`.  Default is
            `tfd.Normal(loc=0., scale=1.)`.
            Must be a scalar-batch, scalar-event distribution.  Typically
            `distribution.reparameterization_type = FULLY_REPARAMETERIZED` or it is
            a function of non-trainable parameters. WARNING: If you backprop through
            a VectorSinhArcsinhDiag sample and `distribution` is not
            `FULLY_REPARAMETERIZED` yet is a function of trainable variables, then
            the gradient will be incorrect!
          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.


return self._loc<EOL>

The `loc` in `Y := loc + scale @ F(Z) * (2 / F(2)).


return self._scale<EOL>

The `LinearOperator` `scale` in `Y := loc + scale @ F(Z) * (2 / F(2)).


return self._tailweight<EOL>

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


return self._skewness<EOL>

Controls the skewness.  `Skewness > 0` means right skew.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale, skewness, tailweight],<EOL>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>tailweight = <NUM_LIT:1.> if tailweight is None else tailweight<EOL>has_default_skewness = skewness is None<EOL>skewness = <NUM_LIT:0.> if skewness is None else skewness<EOL>tailweight = tf.convert_to_tensor(<EOL>value=tailweight, name="<STR_LIT>", dtype=dtype)<EOL>skewness = tf.convert_to_tensor(<EOL>value=skewness, name="<STR_LIT>", dtype=dtype)<EOL>batch_shape = distribution_util.get_broadcast_shape(<EOL>loc, scale, tailweight, skewness)<EOL>if distribution is None:<EOL><INDENT>distribution = normal.Normal(<EOL>loc=tf.zeros([], dtype=dtype),<EOL>scale=tf.ones([], dtype=dtype),<EOL>allow_nan_stats=allow_nan_stats)<EOL><DEDENT>else:<EOL><INDENT>asserts = distribution_util.maybe_check_scalar_distribution(<EOL>distribution, dtype, validate_args)<EOL>if asserts:<EOL><INDENT>loc = distribution_util.with_dependencies(asserts, loc)<EOL><DEDENT><DEDENT>f = sinh_arcsinh_bijector.SinhArcsinh(<EOL>skewness=skewness, tailweight=tailweight)<EOL>if has_default_skewness:<EOL><INDENT>f_noskew = f<EOL><DEDENT>else:<EOL><INDENT>f_noskew = sinh_arcsinh_bijector.SinhArcsinh(<EOL>skewness=dtype_util.as_numpy_dtype(skewness.dtype)(<NUM_LIT:0>),<EOL>tailweight=tailweight)<EOL><DEDENT>c = <NUM_LIT:2> * scale / f_noskew.forward(<EOL>tf.convert_to_tensor(value=<NUM_LIT:2>, dtype=dtype))<EOL>affine = affine_scalar_bijector.AffineScalar(<EOL>shift=loc,<EOL>scale=c,<EOL>validate_args=validate_args)<EOL>bijector = chain_bijector.Chain([affine, f])<EOL>super(SinhArcsinh, self).__init__(<EOL>distribution=distribution,<EOL>bijector=bijector,<EOL>batch_shape=batch_shape,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>self._parameters = parameters<EOL>self._loc = loc<EOL>self._scale = scale<EOL>self._tailweight = tailweight<EOL>self._skewness = skewness<EOL>

Construct SinhArcsinh distribution on `(-inf, inf)`.

        Arguments `(loc, scale, skewness, tailweight)` must have broadcastable shape
        (indexing batch dimensions).  They must all have the same `dtype`.

        Args:
          loc: Floating-point `Tensor`.
          scale:  `Tensor` of same `dtype` as `loc`.
          skewness:  Skewness parameter.  Default is `0.0` (no skew).
          tailweight:  Tailweight parameter. Default is `1.0` (unchanged tailweight)
          distribution: `tf.Distribution`-like instance. Distribution that is
            transformed to produce this distribution.
            Default is `tfd.Normal(0., 1.)`.
            Must be a scalar-batch, scalar-event distribution.  Typically
            `distribution.reparameterization_type = FULLY_REPARAMETERIZED` or it is
            a function of non-trainable parameters. WARNING: If you backprop through
            a `SinhArcsinh` sample and `distribution` is not
            `FULLY_REPARAMETERIZED` yet is a function of trainable variables, then
            the gradient will be incorrect!
          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._loc<EOL>


return self._scale<EOL>


return self._tailweight<EOL>


return self._skewness<EOL>


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

Construct Student's t distributions.

        The distributions have degree of freedom `df`, mean `loc`, and scale
        `scale`.

        The parameters `df`, `loc`, and `scale` must be shaped in a way that
        supports broadcasting (e.g. `df + loc + scale` is a valid operation).

        Args:
          df: Floating-point `Tensor`. The degrees of freedom of the
            distribution(s). `df` must contain only positive values.
          loc: Floating-point `Tensor`. The mean(s) of the distribution(s).
          scale: Floating-point `Tensor`. The scaling factor(s) for the
            distribution(s). Note that `scale` is not technically the standard
            deviation of this distribution but has semantics more similar to
            standard deviation than variance.
          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 loc and scale are different dtypes.


return self._df<EOL>


return self._scale<EOL>

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


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([total_count, logits, probs],<EOL>preferred_dtype=tf.float32)<EOL>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>logits, probs, validate_args=validate_args, name=name, dtype=dtype)<EOL>total_count = tf.convert_to_tensor(<EOL>value=total_count, name="<STR_LIT>", dtype=dtype)<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_positive(total_count)] if validate_args else []):<EOL><INDENT>self._total_count = tf.identity(total_count, name="<STR_LIT>")<EOL><DEDENT><DEDENT>super(NegativeBinomial, self).__init__(<EOL>dtype=self._probs.dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._total_count, self._probs, self._logits],<EOL>name=name)<EOL>

Construct NegativeBinomial distributions.

        Args:
          total_count: Non-negative floating-point `Tensor` with shape
            broadcastable to `[B1,..., Bb]` with `b >= 0` and the same dtype as
            `probs` or `logits`. Defines this as a batch of `N1 x ... x Nm`
            different Negative Binomial distributions. In practice, this represents
            the number of negative Bernoulli trials to stop at (the `total_count`
            of failures), but this is still a valid distribution when
            `total_count` is a non-integer.
          logits: Floating-point `Tensor` with shape broadcastable to
            `[B1, ..., Bb]` where `b >= 0` indicates the number of batch dimensions.
            Each entry represents logits for the probability of success for
            independent Negative Binomial distributions and must be in the open
            interval `(-inf, inf)`. Only one of `logits` or `probs` should be
            specified.
          probs: Positive floating-point `Tensor` with shape broadcastable to
            `[B1, ..., Bb]` where `b >= 0` indicates the number of batch dimensions.
            Each entry represents the probability of success for independent
            Negative Binomial distributions and must be in the open interval
            `(0, 1)`. Only one of `logits` or `probs` should be 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.


return self._total_count<EOL>


return self._logits<EOL>


return self._probs<EOL>


def _fn(self, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>x = getattr(self.distribution, attr)(**kwargs)<EOL>shape = prefer_static.concat([<EOL>self.distribution.batch_shape_tensor(),<EOL>prefer_static.ones(prefer_static.rank_from_shape(self.sample_shape),<EOL>dtype=self.sample_shape.dtype),<EOL>self.distribution.event_shape_tensor(),<EOL>], axis=<NUM_LIT:0>)<EOL>x = tf.reshape(x, shape=shape)<EOL>shape = prefer_static.concat([<EOL>self.distribution.batch_shape_tensor(),<EOL>self.sample_shape,<EOL>self.distribution.event_shape_tensor(),<EOL>], axis=<NUM_LIT:0>)<EOL>return tf.broadcast_to(x, shape)<EOL><DEDENT>return _fn<EOL>

Factory for implementing summary statistics, eg, mean, stddev, mode.


@kullback_leibler.RegisterKL(Sample, Sample)<EOL>def _kl_sample(a, b, name='<STR_LIT>'):

assertions = []<EOL>a_ss = tf.get_static_value(a.sample_shape)<EOL>b_ss = tf.get_static_value(b.sample_shape)<EOL>msg = '<STR_LIT>'<EOL>if a_ss is not None and b_ss is not None:<EOL><INDENT>if not np.array_equal(a_ss, b_ss):<EOL><INDENT>raise ValueError(msg)<EOL><DEDENT><DEDENT>elif a.validate_args or b.validate_args:<EOL><INDENT>assertions.append(assert_util.assert_equal(<EOL>a.sample_shape, b.sample_shape, message=msg))<EOL><DEDENT>with tf.control_dependencies(assertions):<EOL><INDENT>kl = kullback_leibler.kl_divergence(<EOL>a.distribution, b.distribution, name=name)<EOL>n = prefer_static.reduce_prod(a.sample_shape)<EOL>return tf.cast(x=n, dtype=kl.dtype) * kl<EOL><DEDENT>

Batched KL divergence `KL(a || b)` for Sample distributions.

    We can leverage the fact that:

    ```
    KL(Sample(a) || Sample(b)) = sum(KL(a || b))
    ```

    where the sum is over the `sample_shape` dims.

    Args:
      a: Instance of `Sample` distribution.
      b: Instance of `Sample` distribution.
      name: (optional) name to use for created ops.
        Default value: `"kl_sample"`'.

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

    Raises:
      ValueError: If the `sample_shape` of `a` and `b` don't match.


parameters = dict(locals())<EOL>name = name or '<STR_LIT>' + distribution.name<EOL>self._distribution = distribution<EOL>with tf.name_scope(name) as name:<EOL><INDENT>sample_shape = distribution_util.expand_to_vector(tf.convert_to_tensor(<EOL>value=sample_shape, dtype_hint=tf.int32, name='<STR_LIT>'))<EOL>self._sample_shape = sample_shape<EOL>super(Sample, self).__init__(<EOL>dtype=self._distribution.dtype,<EOL>reparameterization_type=self._distribution.reparameterization_type,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=self._distribution.allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=([sample_shape] + distribution._graph_parents),  <EOL>name=name)<EOL><DEDENT>

Construct the `Sample` distribution.

        Args:
          distribution: The base distribution instance to transform. Typically an
            instance of `Distribution`.
          sample_shape: `int` scalar or vector `Tensor` representing the shape of a
            single sample.
          validate_args: Python `bool`.  Whether to validate input with asserts.
            If `validate_args` is `False`, and the inputs are invalid,
            correct behavior is not guaranteed.
          name: The name for ops managed by the distribution.
            Default value: `None` (i.e., `'Sample' + distribution.name`).


@kullback_leibler.RegisterKL(Dirichlet, Dirichlet)<EOL>def _kl_dirichlet_dirichlet(d1, d2, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>digamma_sum_d1 = tf.math.digamma(<EOL>tf.reduce_sum(input_tensor=d1.concentration, axis=-<NUM_LIT:1>, keepdims=True))<EOL>digamma_diff = tf.math.digamma(d1.concentration) - digamma_sum_d1<EOL>concentration_diff = d1.concentration - d2.concentration<EOL>return (<EOL>tf.reduce_sum(input_tensor=concentration_diff * digamma_diff, axis=-<NUM_LIT:1>) -<EOL>tf.math.lbeta(d1.concentration) + tf.math.lbeta(d2.concentration))<EOL><DEDENT>

Batchwise KL divergence KL(d1 || d2) with d1 and d2 Dirichlet.

    Args:
      d1: instance of a Dirichlet distribution object.
      d2: instance of a Dirichlet distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_dirichlet_dirichlet".

    Returns:
      Batchwise KL(d1 || d2)


Initialize a batch of Dirichlet distributions.

        Args:
          concentration: Positive floating-point `Tensor` indicating mean number
            of class occurrences; aka "alpha". Implies `self.dtype`, and
            `self.batch_shape`, `self.event_shape`, i.e., if
            `concentration.shape = [N1, N2, ..., Nm, k]` then
            `batch_shape = [N1, N2, ..., Nm]` and
            `event_shape = [k]`.
          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._concentration<EOL>

Concentration parameter; expected counts for that coordinate.


return self._total_concentration<EOL>

Sum of last dim of concentration parameter.


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

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


def _maybe_assert_valid_concentration(self, concentration, validate_args):

Checks the validity of the concentration parameter.


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>


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>dtype = dtype_util.common_dtype([loc, covariance_matrix], tf.float32)<EOL>loc = loc if loc is None else tf.convert_to_tensor(<EOL>value=loc, name="<STR_LIT>", dtype=dtype)<EOL>if covariance_matrix is None:<EOL><INDENT>scale_tril = None<EOL><DEDENT>else:<EOL><INDENT>covariance_matrix = tf.convert_to_tensor(<EOL>value=covariance_matrix, name="<STR_LIT>", dtype=dtype)<EOL>if validate_args:<EOL><INDENT>covariance_matrix = distribution_util.with_dependencies([<EOL>assert_util.assert_near(<EOL>covariance_matrix,<EOL>tf.linalg.matrix_transpose(covariance_matrix),<EOL>message="<STR_LIT>")<EOL>], covariance_matrix)<EOL><DEDENT>scale_tril = tf.linalg.cholesky(covariance_matrix)<EOL><DEDENT>super(MultivariateNormalFullCovariance, self).__init__(<EOL>loc=loc,<EOL>scale_tril=scale_tril,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL><DEDENT><DEDENT>self._parameters = parameters<EOL>

Construct Multivariate Normal distribution on `R^k`.

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

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

        A non-batch `covariance_matrix` matrix is a `k x k` symmetric positive
        definite matrix.  In other words it is (real) symmetric with all eigenvalues
        strictly positive.

        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.
          covariance_matrix: Floating-point, symmetric positive definite `Tensor` of
            same `dtype` as `loc`.  The strict upper triangle of `covariance_matrix`
            is ignored, so if `covariance_matrix` is not symmetric no error will be
            raised (unless `validate_args is True`).  `covariance_matrix` 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 `covariance_matrix` are specified.


def _is_empty_observation_data(<EOL>feature_ndims, observation_index_points, observations):

<EOL>if observation_index_points is None and observations is None:<EOL><INDENT>return True<EOL><DEDENT>num_obs = tf.compat.dimension_value(<EOL>observation_index_points.shape[-(feature_ndims + <NUM_LIT:1>)])<EOL>if num_obs is not None and num_obs == <NUM_LIT:0>:<EOL><INDENT>return True<EOL><DEDENT>return False<EOL>

Returns `True` if given observation data is empty.

    Emptiness means either
      1. Both `observation_index_points` and `observations` are `None`, or
      2. the "number of observations" shape is 0. The shape of
      `observation_index_points` is `[..., N, f1, ..., fF]`, where `N` is the
      number of observations and the `f`s are feature dims. Thus, we look at the
      shape element just to the left of the leftmost feature dim. If that shape is
      zero, we consider the data empty.

    We don't check the shape of observations; validations are checked elsewhere in
    the calling code, to ensure these shapes are consistent.

    Args:
      feature_ndims: the number of feature dims, as reported by the GP kernel.
      observation_index_points: the observation data locations in the index set.
      observations: the observation data.

    Returns:
      is_empty: True if the data were deemed to be empty.


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.


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.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[df, index_points, jitter], tf.float32)<EOL>df = tf.convert_to_tensor(value=df, dtype=dtype, name='<STR_LIT>')<EOL>index_points = tf.convert_to_tensor(<EOL>value=index_points, dtype=dtype, name='<STR_LIT>')<EOL>jitter = tf.convert_to_tensor(value=jitter, dtype=dtype, name='<STR_LIT>')<EOL>with tf.control_dependencies([<EOL>assert_util.assert_greater(<EOL>df, tf.cast(<NUM_LIT>, df.dtype), message='<STR_LIT>')<EOL>] if validate_args else []):<EOL><INDENT>self._df = tf.identity(df)<EOL><DEDENT>self._kernel = kernel<EOL>self._index_points = index_points<EOL>if mean_fn is None:<EOL><INDENT>mean_fn = lambda x: tf.zeros([<NUM_LIT:1>], dtype=dtype)<EOL><DEDENT>else:<EOL><INDENT>if not callable(mean_fn):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT><DEDENT>self._mean_fn = mean_fn<EOL>self._jitter = jitter<EOL>with tf.name_scope('<STR_LIT>'):<EOL><INDENT>kernel_matrix = _add_diagonal_shift(<EOL>kernel.matrix(self.index_points, self.index_points),<EOL>jitter)<EOL>self._covariance_matrix = kernel_matrix<EOL>scale = tf.linalg.LinearOperatorLowerTriangular(<EOL>tf.linalg.cholesky(<EOL>((self.df - <NUM_LIT:2>) / self.df)[..., tf.newaxis, tf.newaxis] *<EOL>kernel_matrix),<EOL>is_non_singular=True,<EOL>name='<STR_LIT>')<EOL>super(StudentTProcess, self).__init__(<EOL>df=df,<EOL>loc=mean_fn(index_points),<EOL>scale=scale,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL>self._parameters = parameters<EOL>self._graph_parents = [index_points, jitter]<EOL><DEDENT><DEDENT>

Instantiate a StudentTProcess Distribution.

        Args:
          df: Positive Floating-point `Tensor` representing the degrees of freedom.
            Must be greater than 2.
          kernel: `PositiveSemidefiniteKernel`-like instance representing the
            TP's covariance function.
          index_points: `float` `Tensor` representing finite (batch of) vector(s) of
            points in the index set over which the TP is defined. Shape has the form
            `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
            dimensions and must equal `kernel.feature_ndims` and `e` is the number
            (size) of index points in each batch. Ultimately this distribution
            corresponds to a `e`-dimensional multivariate Student's T. The batch
            shape must be broadcastable with `kernel.batch_shape` and any batch dims
            yielded by `mean_fn`.
          mean_fn: Python `callable` that acts on `index_points` to produce a (batch
            of) vector(s) of mean values at `index_points`. Takes a `Tensor` of
            shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
            broadcastable with `[b1, ..., bB]`. Default value: `None` implies
            constant zero function.
          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: "StudentTProcess".

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


<EOL>z = tf.convert_to_tensor(value=z)<EOL>wrap = lambda result: tf.debugging.check_numerics(result, '<STR_LIT>'.format(v<EOL>))<EOL>if float(v) >= <NUM_LIT:2>:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>')<EOL><DEDENT>cache = cache or {}<EOL>safe_z = tf.where(z > <NUM_LIT:0>, z, tf.ones_like(z))<EOL>if v in cache:<EOL><INDENT>return wrap(cache[v])<EOL><DEDENT>if v == <NUM_LIT:0>:<EOL><INDENT>cache[v] = tf.math.bessel_i0e(z)<EOL><DEDENT>elif v == <NUM_LIT:1>:<EOL><INDENT>cache[v] = tf.math.bessel_i1e(z)<EOL><DEDENT>elif v == <NUM_LIT:0.5>:<EOL><INDENT>sinhe = lambda x: (tf.exp(x - tf.abs(x)) - tf.exp(-x - tf.abs(x))) / <NUM_LIT:2><EOL>cache[v] = (<EOL>np.sqrt(<NUM_LIT:2> / np.pi) * sinhe(z) *<EOL>tf.where(z > <NUM_LIT:0>, tf.math.rsqrt(safe_z), tf.ones_like(safe_z)))<EOL><DEDENT>elif v == -<NUM_LIT:0.5>:<EOL><INDENT>coshe = lambda x: (tf.exp(x - tf.abs(x)) + tf.exp(-x - tf.abs(x))) / <NUM_LIT:2><EOL>cache[v] = (<EOL>np.sqrt(<NUM_LIT:2> / np.pi) * coshe(z) *<EOL>tf.where(z > <NUM_LIT:0>, tf.math.rsqrt(safe_z), tf.ones_like(safe_z)))<EOL><DEDENT>if v <= <NUM_LIT:1>:<EOL><INDENT>return wrap(cache[v])<EOL><DEDENT>cache[v] = (_bessel_ive(v - <NUM_LIT:2>, z, cache) -<EOL>(<NUM_LIT:2> * (v - <NUM_LIT:1>)) * _bessel_ive(v - <NUM_LIT:1>, z, cache) / z)<EOL>return wrap(cache[v])<EOL>

Computes I_v(z)*exp(-abs(z)) using a recurrence relation, where z > 0.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([mean_direction, concentration],<EOL>tf.float32)<EOL>mean_direction = tf.convert_to_tensor(<EOL>value=mean_direction, name='<STR_LIT>', dtype=dtype)<EOL>concentration = tf.convert_to_tensor(<EOL>value=concentration, name='<STR_LIT>', dtype=dtype)<EOL>assertions = [<EOL>assert_util.assert_non_negative(<EOL>concentration, message='<STR_LIT>'),<EOL>assert_util.assert_greater(<EOL>tf.shape(input=mean_direction)[-<NUM_LIT:1>],<EOL><NUM_LIT:1>,<EOL>message='<STR_LIT>'),<EOL>assert_util.assert_near(<EOL><NUM_LIT:1.>,<EOL>tf.linalg.norm(tensor=mean_direction, axis=-<NUM_LIT:1>),<EOL>message='<STR_LIT>')<EOL>] if validate_args else []<EOL>static_event_dim = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(mean_direction.shape, <NUM_LIT:1>)[-<NUM_LIT:1>])<EOL>if static_event_dim is not None and static_event_dim > <NUM_LIT:5>:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>elif validate_args:<EOL><INDENT>assertions += [<EOL>assert_util.assert_less_equal(<EOL>tf.shape(input=mean_direction)[-<NUM_LIT:1>],<EOL><NUM_LIT:5>,<EOL>message='<STR_LIT>')<EOL>]<EOL><DEDENT>with tf.control_dependencies(assertions):<EOL><INDENT>self._mean_direction = tf.identity(mean_direction)<EOL>self._concentration = tf.identity(concentration)<EOL><DEDENT>dtype_util.assert_same_float_dtype(<EOL>[self._mean_direction, self._concentration])<EOL>reparameterization_type = (<EOL>reparameterization.FULLY_REPARAMETERIZED<EOL>if static_event_dim == <NUM_LIT:3> else<EOL>reparameterization.NOT_REPARAMETERIZED)<EOL>super(VonMisesFisher, self).__init__(<EOL>dtype=self._concentration.dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization_type,<EOL>parameters=parameters,<EOL>graph_parents=[self._mean_direction, self._concentration],<EOL>name=name)<EOL><DEDENT>

Creates a new `VonMisesFisher` instance.

        Args:
          mean_direction: Floating-point `Tensor` with shape [B1, ... Bn, D].
            A unit vector indicating the mode of the distribution, or the
            unit-normalized direction of the mean. (This is *not* in general the
            mean of the distribution; the mean is not generally in the support of
            the distribution.) NOTE: `D` is currently restricted to <= 5.
          concentration: Floating-point `Tensor` having batch shape [B1, ... Bn]
            broadcastable with `mean_direction`. The level of concentration of
            samples around the `mean_direction`. `concentration=0` indicates a
            uniform distribution over the unit hypersphere, and `concentration=+inf`
            indicates a `Deterministic` distribution (delta function) at
            `mean_direction`.
          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: For known-bad arguments, i.e. unsupported event dimension.


return self._mean_direction<EOL>


return self._concentration<EOL>


event_dim = tf.compat.dimension_value(self.event_shape[<NUM_LIT:0>])<EOL>if event_dim is None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>safe_conc = tf.where(self.concentration > <NUM_LIT:0>,<EOL>self.concentration,<EOL>tf.ones_like(self.concentration))<EOL>safe_lognorm = ((event_dim / <NUM_LIT:2> - <NUM_LIT:1>) * tf.math.log(safe_conc) -<EOL>(event_dim / <NUM_LIT:2>) * np.log(<NUM_LIT:2> * np.pi) -<EOL>tf.math.log(_bessel_ive(event_dim / <NUM_LIT:2> - <NUM_LIT:1>, safe_conc)) -<EOL>tf.abs(safe_conc))<EOL>log_nsphere_surface_area = (<EOL>np.log(<NUM_LIT>) + (event_dim / <NUM_LIT:2>) * np.log(np.pi) -<EOL>tf.math.lgamma(tf.cast(event_dim / <NUM_LIT:2>, self.dtype)))<EOL>return tf.where(self.concentration > <NUM_LIT:0>,<EOL>-safe_lognorm,<EOL>log_nsphere_surface_area * tf.ones_like(safe_lognorm))<EOL>

Computes the log-normalizer of the distribution.


Check counts for proper shape, values, then return tensor version.


return (self.mean_direction +<EOL>tf.zeros_like(self.concentration)[..., tf.newaxis])<EOL>

The mode of the von Mises-Fisher distribution is the mean direction.


event_dim = (<EOL>tf.compat.dimension_value(self.event_shape[<NUM_LIT:0>]) or<EOL>self._event_shape_tensor()[<NUM_LIT:0>])<EOL>basis = tf.concat([[<NUM_LIT:1.>], tf.zeros([event_dim - <NUM_LIT:1>], dtype=self.dtype)],<EOL>axis=<NUM_LIT:0>),<EOL>u = tf.nn.l2_normalize(basis - self.mean_direction, axis=-<NUM_LIT:1>)<EOL>return samples - <NUM_LIT:2> * tf.reduce_sum(<EOL>input_tensor=samples * u, axis=-<NUM_LIT:1>, keepdims=True) * u<EOL>

Applies a Householder rotation to `samples`.


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>


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._rate = tf.convert_to_tensor(<EOL>value=rate,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([rate], preferred_dtype=tf.float32))<EOL><DEDENT>super(Exponential, self).__init__(<EOL>concentration=<NUM_LIT:1.>,<EOL>rate=self._rate,<EOL>allow_nan_stats=allow_nan_stats,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>self._parameters = parameters<EOL>self._graph_parents += [self._rate]<EOL>

Construct Exponential distribution with parameter `rate`.

        Args:
          rate: Floating point tensor, equivalent to `1 / mean`. 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.


def _get_tol(tol, dtype, validate_args):

if tol is None:<EOL><INDENT>return tf.convert_to_tensor(value=<NUM_LIT:0>, dtype=dtype)<EOL><DEDENT>tol = tf.convert_to_tensor(value=tol, dtype=dtype)<EOL>if validate_args:<EOL><INDENT>tol = distribution_util.with_dependencies([<EOL>assert_util.assert_non_negative(<EOL>tol, message="<STR_LIT>")<EOL>], tol)<EOL><DEDENT>return tol<EOL>

Gets a Tensor of type `dtype`, 0 if `tol` is None, validation optional.


@kullback_leibler.RegisterKL(_BaseDeterministic, distribution.Distribution)<EOL>def _kl_deterministic_distribution(a, b, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return -b.log_prob(a.loc)<EOL><DEDENT>

Calculate the batched KL divergence `KL(a || b)` with `a` Deterministic.

    Args:
      a: instance of a Deterministic distribution object.
      b: instance of a Distribution distribution object.
      name: (optional) Name to use for created operations. Default is
        "kl_deterministic_distribution".

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


with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, atol, rtol],<EOL>preferred_dtype=tf.float32)<EOL>loc = tf.convert_to_tensor(value=loc, name="<STR_LIT>", dtype=dtype)<EOL>if is_vector and validate_args:<EOL><INDENT>msg = "<STR_LIT>"<EOL>if tensorshape_util.rank(loc.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(loc.shape) < <NUM_LIT:1>:<EOL><INDENT>raise ValueError(msg)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>loc = distribution_util.with_dependencies(<EOL>[assert_util.assert_rank_at_least(loc, <NUM_LIT:1>, message=msg)], loc)<EOL><DEDENT><DEDENT>self._loc = loc<EOL>self._atol = _get_tol(atol, self._loc.dtype, validate_args)<EOL>self._rtol = _get_tol(rtol, self._loc.dtype, validate_args)<EOL>super(_BaseDeterministic, self).__init__(<EOL>dtype=self._loc.dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._loc, self._atol, self._rtol],<EOL>name=name)<EOL>if rtol is None:<EOL><INDENT>self._slack = self.atol<EOL><DEDENT>else:<EOL><INDENT>self._slack = self.atol + self.rtol * tf.abs(self.loc)<EOL><DEDENT><DEDENT>

Initialize a batch of `_BaseDeterministic` distributions.

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

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

        Args:
          loc: Numeric `Tensor`.  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`.
          is_vector:  Python `bool`.  If `True`, this is for `VectorDeterministic`,
            else `Deterministic`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          parameters: Dict of locals to facilitate copy construction.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          ValueError:  If `loc` is a scalar.


return self._loc<EOL>

Point (or batch of points) at which this distribution is supported.


return self._atol<EOL>

Absolute tolerance for comparing points to `self.loc`.


return self._rtol<EOL>

Relative tolerance for comparing points to `self.loc`.


Initialize a scalar `Deterministic` distribution.

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

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

        Args:
          loc: Numeric `Tensor` of shape `[B1, ..., Bb]`, with `b >= 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.


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.


return self._event_size<EOL>

Scalar `int32` tensor: the number of classes.


return self._temperature<EOL>

Batchwise temperature tensor of a RelaxedCategorical.


return self._logits<EOL>


return self._probs<EOL>


Initialize RelaxedOneHotCategorical using class log-probabilities.

        Args:
          temperature: An 0-D `Tensor`, representing the temperature
            of a set of RelaxedOneHotCategorical distributions. The temperature
            should be positive.
          logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities
            of a set of RelaxedOneHotCategorical 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 RelaxedOneHotCategorical 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: Unused in this distribution.
          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: A name for this distribution (optional).


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

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return (tf.math.log(b.scale) - tf.math.log(a.scale) + np.euler_gamma *<EOL>(a.scale / b.scale - <NUM_LIT:1.>) +<EOL>tf.math.expm1((b.loc - a.loc) / b.scale +<EOL>tf.math.lgamma(a.scale / b.scale + <NUM_LIT:1.>)) +<EOL>(a.loc - b.loc) / b.scale)<EOL><DEDENT>

Calculate the batched KL divergence KL(a || b) with a and b Gumbel.

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

    Returns:
      Batchwise KL(a || b)


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>loc = tf.identity(loc, name="<STR_LIT>")<EOL>scale = tf.identity(scale, name="<STR_LIT>")<EOL>dtype_util.assert_same_float_dtype([loc, scale])<EOL>self._gumbel_bijector = gumbel_bijector.Gumbel(<EOL>loc=loc, scale=scale, validate_args=validate_args)<EOL><DEDENT>super(Gumbel, self).__init__(<EOL>distribution=uniform.Uniform(<EOL>low=np.finfo(dtype_util.as_numpy_dtype(dtype)).tiny,<EOL>high=tf.ones([], dtype=loc.dtype),<EOL>allow_nan_stats=allow_nan_stats),<EOL>bijector=invert_bijector.Invert(self._gumbel_bijector),<EOL>batch_shape=distribution_util.get_broadcast_shape(loc, scale),<EOL>parameters=parameters,<EOL>name=name)<EOL><DEDENT>

Construct Gumbel distributions with location and scale `loc` and `scale`.

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

        Args:
          loc: Floating point tensor, the means of the distribution(s).
          scale: Floating point tensor, the scales of the distribution(s).
            scale 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`.
          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: `'Gumbel'`.

        Raises:
          TypeError: if loc and scale are different dtypes.


return self._gumbel_bijector.loc<EOL>

Distribution parameter for the location.


return self._gumbel_bijector.scale<EOL>


hierarchy_a = tf_inspect.getmro(type_a)<EOL>hierarchy_b = tf_inspect.getmro(type_b)<EOL>dist_to_children = None<EOL>kl_fn = None<EOL>for mro_to_a, parent_a in enumerate(hierarchy_a):<EOL><INDENT>for mro_to_b, parent_b in enumerate(hierarchy_b):<EOL><INDENT>candidate_dist = mro_to_a + mro_to_b<EOL>candidate_kl_fn = _DIVERGENCES.get((parent_a, parent_b), None)<EOL>if not kl_fn or (candidate_kl_fn and candidate_dist < dist_to_children):<EOL><INDENT>dist_to_children = candidate_dist<EOL>kl_fn = candidate_kl_fn<EOL><DEDENT><DEDENT><DEDENT>return kl_fn<EOL>

Get the KL function registered for classes a and b.


def kl_divergence(distribution_a, distribution_b,<EOL>allow_nan_stats=True, name=None):

kl_fn = _registered_kl(type(distribution_a), type(distribution_b))<EOL>if kl_fn is None:<EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>type(distribution_a).__name__, type(distribution_b).__name__))<EOL><DEDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>kl_t = kl_fn(distribution_a, distribution_b, name=name)<EOL>if allow_nan_stats:<EOL><INDENT>return kl_t<EOL><DEDENT>kl_t = tf.identity(kl_t, name="<STR_LIT>")<EOL>with tf.control_dependencies([<EOL>tf.Assert(<EOL>tf.logical_not(tf.reduce_any(input_tensor=tf.math.is_nan(kl_t))),<EOL>[<EOL>("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>distribution_a.name, distribution_b.name)),<EOL>kl_t<EOL>])<EOL>]):<EOL><INDENT>return tf.identity(kl_t, name="<STR_LIT>")<EOL><DEDENT><DEDENT>

Get the KL-divergence KL(distribution_a || distribution_b).

    If there is no KL method registered specifically for `type(distribution_a)`
    and `type(distribution_b)`, then the class hierarchies of these types are
    searched.

    If one KL method is registered between any pairs of classes in these two
    parent hierarchies, it is used.

    If more than one such registered method exists, the method whose registered
    classes have the shortest sum MRO paths to the input types is used.

    If more than one such shortest path exists, the first method
    identified in the search is used (favoring a shorter MRO distance to
    `type(distribution_a)`).

    Args:
      distribution_a: The first distribution.
      distribution_b: The second distribution.
      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.

    Returns:
      A Tensor with the batchwise KL-divergence between `distribution_a`
      and `distribution_b`.

    Raises:
      NotImplementedError: If no KL method is defined for distribution types
        of `distribution_a` and `distribution_b`.


def cross_entropy(ref, other,<EOL>allow_nan_stats=True, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return ref.entropy() + kl_divergence(<EOL>ref, other, allow_nan_stats=allow_nan_stats)<EOL><DEDENT>

Computes the (Shannon) cross entropy.

    Denote two distributions by `P` (`ref`) and `Q` (`other`). Assuming `P, Q`
    are absolutely continuous with respect to one another and permit densities
    `p(x) dr(x)` and `q(x) dr(x)`, (Shanon) cross entropy is defined as:

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

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

    Args:
      ref: `tfd.Distribution` instance.
      other: `tfd.Distribution` instance.
      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` prepended to names of ops created by this function.

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


def __init__(self, dist_cls_a, dist_cls_b):

self._key = (dist_cls_a, dist_cls_b)<EOL>

Initialize the KL registrar.

        Args:
          dist_cls_a: the class of the first argument of the KL divergence.
          dist_cls_b: the class of the second argument of the KL divergence.


if not callable(kl_fn):<EOL><INDENT>raise TypeError("<STR_LIT>" % kl_fn)<EOL><DEDENT>if self._key in _DIVERGENCES:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>% (self._key[<NUM_LIT:0>].__name__, self._key[<NUM_LIT:1>].__name__,<EOL>_DIVERGENCES[self._key]))<EOL><DEDENT>_DIVERGENCES[self._key] = kl_fn<EOL>return kl_fn<EOL>

Perform the KL registration.

        Args:
          kl_fn: The function to use for the KL divergence.

        Returns:
          kl_fn

        Raises:
          TypeError: if kl_fn is not a callable.
          ValueError: if a KL divergence function has already been registered for
            the given argument classes.


@kullback_leibler.RegisterKL(Gamma, Gamma)<EOL>def _kl_gamma_gamma(g0, g1, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return (((g0.concentration - g1.concentration) *<EOL>tf.math.digamma(g0.concentration)) +<EOL>tf.math.lgamma(g1.concentration) -<EOL>tf.math.lgamma(g0.concentration) +<EOL>g1.concentration * tf.math.log(g0.rate) -<EOL>g1.concentration * tf.math.log(g1.rate) + g0.concentration *<EOL>(g1.rate / g0.rate - <NUM_LIT:1.>))<EOL><DEDENT>

Calculate the batched KL divergence KL(g0 || g1) with g0 and g1 Gamma.

    Args:
      g0: instance of a Gamma distribution object.
      g1: instance of a Gamma distribution object.
      name: (optional) Name to use for created operations. Default is
        "kl_gamma_gamma".

    Returns:
      kl_gamma_gamma: `Tensor`. The batchwise KL(g0 || g1).


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._concentration<EOL>


return self._rate<EOL>


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale], tf.float32)<EOL>super(LogNormal, self).__init__(<EOL>distribution=normal.Normal(<EOL>loc=tf.convert_to_tensor(value=loc, name="<STR_LIT>", dtype=dtype),<EOL>scale=tf.convert_to_tensor(<EOL>value=scale, name="<STR_LIT>", dtype=dtype)),<EOL>bijector=exp_bijector.Exp(),<EOL>validate_args=validate_args,<EOL>parameters=parameters,<EOL>name=name)<EOL><DEDENT>

Construct a log-normal distribution.

        The LogNormal distribution models positive-valued random variables
        whose logarithm is normally distributed with mean `loc` and
        standard deviation `scale`. It is constructed as the exponential
        transformation of a Normal distribution.

        Args:
          loc: Floating-point `Tensor`; the means of the underlying
            Normal distribution(s).
          scale: Floating-point `Tensor`; the stddevs of the underlying
            Normal distribution(s).
          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.


return self.distribution.loc<EOL>

Distribution parameter for the pre-transformed mean.


return self.distribution.scale<EOL>

Distribution parameter for the pre-transformed standard deviation.


z = x - y<EOL>return z[..., tf.newaxis, :] * z[..., tf.newaxis]<EOL>

Convenience function analogous to tf.squared_difference.


if tf.executing_eagerly():<EOL><INDENT>with tf.GradientTape() as tape:<EOL><INDENT>tape.watch(x)<EOL>value = f(x)<EOL><DEDENT>batch_jacobian = tape.batch_jacobian(value, x)<EOL><DEDENT>else:<EOL><INDENT>value = f(x)<EOL>batch_jacobian = gradients.batch_jacobian(value, x)<EOL><DEDENT>return value, batch_jacobian<EOL>

Enables uniform interface to value and batch jacobian calculation.

    Works in both eager and graph modes.

    Arguments:
      f: The scalar function to evaluate.
      x: The value at which to compute the value and the batch jacobian.

    Returns:
      A tuple (f(x), J(x)), where J(x) is the batch jacobian.


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.