<EOL>s_shape = tensors[<NUM_LIT:0>].shape<EOL>for t in tensors[<NUM_LIT:1>:]:<EOL><INDENT>s_shape = tf.broadcast_static_shape(s_shape, t.shape)<EOL><DEDENT>if tensorshape_util.is_fully_defined(s_shape):<EOL><INDENT>return tensorshape_util.as_list(s_shape)<EOL><DEDENT>d_shape = tf.shape(input=tensors[<NUM_LIT:0>])<EOL>for t in tensors[<NUM_LIT:1>:]:<EOL><INDENT>d_shape = tf.broadcast_dynamic_shape(d_shape, tf.shape(input=t))<EOL><DEDENT>return d_shape<EOL>

Get broadcast shape as a Python list of integers (preferred) or `Tensor`.

    Args:
      *tensors:  One or more `Tensor` objects (already converted!).

    Returns:
      broadcast shape:  Python list (if shapes determined statically), otherwise
        an `int32` `Tensor`.


return {<EOL>tf.bool: True,<EOL>tf.uint8: True,<EOL>tf.uint16: True,<EOL>}.get(dt.base_dtype, False)<EOL>

Helper returning True if dtype is known to be unsigned.


def _maybe_attach_assertion(x):<EOL><INDENT>if not validate_args:<EOL><INDENT>return x<EOL><DEDENT>if assert_positive:<EOL><INDENT>return with_dependencies([<EOL>assert_util.assert_positive(<EOL>tf.linalg.diag_part(x), message="<STR_LIT>"),<EOL>], x)<EOL><DEDENT>return with_dependencies([<EOL>assert_util.assert_none_equal(<EOL>tf.linalg.diag_part(x),<EOL>tf.zeros([], x.dtype),<EOL>message="<STR_LIT>"),<EOL>], x)<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[loc, scale_tril, scale_diag, scale_identity_multiplier],<EOL>preferred_dtype=tf.float32)<EOL>loc = _convert_to_tensor(loc, name="<STR_LIT>", dtype=dtype)<EOL>scale_tril = _convert_to_tensor(scale_tril, name="<STR_LIT>", dtype=dtype)<EOL>scale_diag = _convert_to_tensor(scale_diag, name="<STR_LIT>", dtype=dtype)<EOL>scale_identity_multiplier = _convert_to_tensor(<EOL>scale_identity_multiplier,<EOL>name="<STR_LIT>",<EOL>dtype=dtype)<EOL><DEDENT>if scale_tril is not None:<EOL><INDENT>scale_tril = tf.linalg.band_part(scale_tril, -<NUM_LIT:1>, <NUM_LIT:0>)  <EOL>tril_diag = tf.linalg.diag_part(scale_tril)<EOL>if scale_diag is not None:<EOL><INDENT>tril_diag += scale_diag<EOL><DEDENT>if scale_identity_multiplier is not None:<EOL><INDENT>tril_diag += scale_identity_multiplier[..., tf.newaxis]<EOL><DEDENT>scale_tril = tf.linalg.set_diag(scale_tril, tril_diag)<EOL>return tf.linalg.LinearOperatorLowerTriangular(<EOL>tril=_maybe_attach_assertion(scale_tril),<EOL>is_non_singular=True,<EOL>is_self_adjoint=False,<EOL>is_positive_definite=assert_positive)<EOL><DEDENT>return make_diag_scale(<EOL>loc=loc,<EOL>scale_diag=scale_diag,<EOL>scale_identity_multiplier=scale_identity_multiplier,<EOL>shape_hint=shape_hint,<EOL>validate_args=validate_args,<EOL>assert_positive=assert_positive,<EOL>name=name)<EOL>

Creates a LinearOperator representing a lower triangular matrix.

    Args:
      loc: Floating-point `Tensor`. This is used for inferring shape in the case
        where only `scale_identity_multiplier` is set.
      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 the
        LinearOperator. The upper triangular elements above the diagonal are
        ignored.
      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 the LinearOperator.
      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`.
      shape_hint: scalar integer `Tensor` representing a hint at the dimension of
        the identity matrix when only `scale_identity_multiplier` is set.
      validate_args: Python `bool` indicating whether arguments should be checked
        for correctness.
      assert_positive: Python `bool` indicating whether LinearOperator should be
        checked for being positive definite.
      name: Python `str` name given to ops managed by this object.

    Returns:
      `LinearOperator` representing a lower triangular matrix.

    Raises:
      ValueError:  If only `scale_identity_multiplier` is set and `loc` and
        `shape_hint` are both None.


static_x = tf.get_static_value(x)<EOL>if static_x is not None:<EOL><INDENT>return static_x<EOL><DEDENT>return x<EOL>

Return static value of tensor `x` if available, else `x`.

    Args:
      x: `Tensor` (already converted).

    Returns:
      Numpy array (if static value is obtainable), else `Tensor`.


with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>threshold = np.log(np.finfo(dtype_util.as_numpy_dtype(x.dtype)).eps) + <NUM_LIT><EOL>is_too_small = tf.less(x, np.exp(threshold))<EOL>is_too_large = tf.greater(x, -threshold)<EOL>too_small_value = tf.math.log(x)<EOL>too_large_value = x<EOL>x = tf.where(tf.logical_or(is_too_small, is_too_large), tf.ones_like(x), x)<EOL>y = x + tf.math.log(-tf.math.expm1(-x))  <EOL>return tf.where(is_too_small, too_small_value,<EOL>tf.where(is_too_large, too_large_value, y))<EOL><DEDENT>

Computes the inverse softplus, i.e., x = softplus_inverse(softplus(x)).

    Mathematically this op is equivalent to:

    ```none
    softplus_inverse = log(exp(x) - 1.)
    ```

    Args:
      x: `Tensor`. Non-negative (not enforced), floating-point.
      name: A name for the operation (optional).

    Returns:
      `Tensor`. Has the same type/shape as input `x`.


if dtype is None:<EOL><INDENT>dtype = dtype_util.common_dtype([probs, logits], preferred_dtype=tf.float32)<EOL><DEDENT>with tf.name_scope(name):<EOL><INDENT>if (probs is None) == (logits is None):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if probs is None:<EOL><INDENT>logits = tf.convert_to_tensor(value=logits, name="<STR_LIT>", dtype=dtype)<EOL>if not dtype_util.is_floating(logits.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if multidimensional:<EOL><INDENT>if validate_args:<EOL><INDENT>logits = embed_check_categorical_event_shape(logits)<EOL><DEDENT>return logits, tf.nn.softmax(logits, name="<STR_LIT>")<EOL><DEDENT>return logits, tf.sigmoid(logits, name="<STR_LIT>")<EOL><DEDENT>probs = tf.convert_to_tensor(value=probs, name="<STR_LIT>", dtype=dtype)<EOL>if not dtype_util.is_floating(probs.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if validate_args:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>one = tf.constant(<NUM_LIT:1.>, probs.dtype)<EOL>dependencies = [assert_util.assert_non_negative(probs)]<EOL>if multidimensional:<EOL><INDENT>probs = embed_check_categorical_event_shape(probs)<EOL>dependencies += [<EOL>assert_util.assert_near(<EOL>tf.reduce_sum(input_tensor=probs, axis=-<NUM_LIT:1>),<EOL>one,<EOL>message="<STR_LIT>")<EOL>]<EOL><DEDENT>else:<EOL><INDENT>dependencies += [<EOL>assert_util.assert_less_equal(<EOL>probs, one, message="<STR_LIT>")<EOL>]<EOL><DEDENT>probs = with_dependencies(dependencies, probs)<EOL><DEDENT><DEDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>if multidimensional:<EOL><INDENT>return tf.math.log(probs), probs<EOL><DEDENT>return tf.math.log(probs) - tf.math.log1p(-<NUM_LIT:1.> * probs), probs<EOL><DEDENT><DEDENT>

Converts logit to probabilities (or vice-versa), and returns both.

    Args:
      logits: Floating-point `Tensor` representing log-odds.
      probs: Floating-point `Tensor` representing probabilities.
      multidimensional: Python `bool`, default `False`. If `True`, represents
        whether the last dimension of `logits` or `probs`, a `[N1, N2, ...  k]`
        dimensional tensor, representing the logit or probability of `shape[-1]`
        classes.
      validate_args: Python `bool`, default `False`. When `True`, either assert `0
        <= probs <= 1` (if not `multidimensional`) or that the last dimension of
        `probs` sums to one.
      name: A name for this operation (optional).
      dtype: `tf.DType` to prefer when converting args to `Tensor`s.

    Returns:
      logits, probs: Tuple of `Tensor`s. If `probs` has an entry that is `0` or
        `1`, then the corresponding entry in the returned logit will be `-Inf` and
        `Inf` respectively.

    Raises:
      ValueError: if neither `probs` nor `logits` were passed in, or both were.


@hps.composite<EOL>def broadcasting_shapes(draw, target_shape, n):

target_shape = tf.TensorShape(target_shape)<EOL>target_rank = target_shape.ndims<EOL>result = []<EOL>current_shape = tf.TensorShape([])<EOL>for is_last in [False] * (n-<NUM_LIT:1>) + [True]:<EOL><INDENT>next_rank, force_fullsize_dim = _compute_rank_and_fullsize_reqd(<EOL>draw, target_shape, current_shape, is_last=is_last)<EOL>next_shape = target_shape[target_rank - next_rank:].as_list()<EOL>for i, force_fullsize in enumerate(force_fullsize_dim):<EOL><INDENT>if not force_fullsize and draw(hps.booleans()):<EOL><INDENT>next_shape[i] = <NUM_LIT:1><EOL><DEDENT><DEDENT>next_shape = tf.TensorShape(next_shape)<EOL>current_shape = tf.broadcast_static_shape(current_shape, next_shape)<EOL>result.append(next_shape)<EOL><DEDENT>return result<EOL>

Draws a set of `n` shapes that broadcast to `target_shape`.

    For each shape we need to choose its rank, and whether or not each axis i is 1
    or target_shape[i]. This function chooses a set of `n` shapes that have
    possibly mismatched ranks, and possibly broadcasting axes, with the promise
    that the broadcast of the set of all shapes matches `target_shape`.

    Args:
      draw: Hypothesis sampler.
      target_shape: The target (fully-defined) batch shape.
      n: `int`, the number of shapes to draw.

    Returns:
      shapes: Sequence of `tf.TensorShape` such that the set of shapes broadcast
        to `target_shape`. The shapes are fully defined.


if isinstance(x, tf.TensorShape):<EOL><INDENT>return x.dims<EOL><DEDENT>r = tf.TensorShape(x).dims<EOL>return None if r is None else list(map(tf.compat.dimension_value, r))<EOL>

Returns a list of dimension sizes, or `None` if `rank` is unknown.

    For more details, see `help(tf.TensorShape.dims)`.

    Args:
      x: object representing a shape; convertible to `tf.TensorShape`.

    Returns:
      shape_as_list: list of sizes or `None` values representing each
        dimensions size if known. A size is `tf.Dimension` if input is a
        `tf.TensorShape` and an `int` otherwise.


shape = tf.get_static_value(tensor)<EOL>if shape is not None:<EOL><INDENT>return [None if dim == -<NUM_LIT:1> else dim for dim in shape]<EOL><DEDENT>return tensor_util.constant_value_as_shape(tensor)<EOL>

A version of `constant_value()` that returns a `TensorShape`.

    This version should be used when a constant tensor value is
    interpreted as a (possibly partial) shape, e.g. in the shape
    function for `tf.reshape()`. By explicitly requesting a
    `TensorShape` as the return value, it is possible to represent
    unknown dimensions; by contrast, `constant_value()` is
    all-or-nothing.

    Args:
      tensor: The rank-0 or rank-1 Tensor to be evaluated.

    Returns:
      A `TensorShape` based on the constant value of the given `tensor`.

    Raises:
      ValueError: If the shape is rank-0 and is not statically known to be -1.


return tf.TensorShape(x).rank<EOL>

Returns the rank of this shape, or `None` if it is unspecified.

    For more details, see `help(tf.TensorShape.rank)`.

    Args:
      x: object representing a shape; convertible to `tf.TensorShape`.

    Returns:
      rank: `int` representing the number of shape dimensions.


tf.TensorShape(x).assert_is_compatible_with(other)<EOL>

Raises exception if `x` and `other` do not represent the same shape.

    This method can be used to assert that there exists a shape that both
    `x` and `other` represent.

    For more details, see `help(tf.TensorShape.assert_is_compatible_with)`.

    Args:
      x: object representing a shape; convertible to `tf.TensorShape`.
      other: object representing a shape; convertible to `tf.TensorShape`.

    Returns:
      None

    Raises:
      ValueError: If `x` and `other` do not represent the same shape.


is_nan = np.isnan(self._GetNdArray(a))<EOL>all_true = np.ones_like(is_nan, dtype=np.bool)<EOL>self.assertAllEqual(all_true, is_nan)<EOL>

Assert that every entry in a `Tensor` is NaN.

        Args:
          a: A `Tensor` whose entries must be verified as NaN.


return base_dtype(a) == base_dtype(b)<EOL>

Returns `True` if base dtypes are identical.


dtype = tf.as_dtype(dtype)<EOL>if hasattr(dtype, '<STR_LIT>'):<EOL><INDENT>return dtype.max<EOL><DEDENT>use_finfo = is_floating(dtype) or is_complex(dtype)<EOL>return np.finfo(dtype).max if use_finfo else np.iinfo(dtype).max<EOL>

Returns the maximum representable value in this data type.


dtype = tf.as_dtype(dtype)<EOL>if hasattr(dtype, '<STR_LIT>'):<EOL><INDENT>return dtype.is_bool<EOL><DEDENT>return np.dtype(dtype).kind == '<STR_LIT:b>'<EOL>

Returns whether this is a boolean data type.


with tf.name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>if dtype_util.as_numpy_dtype(x.dtype) not in [np.float32, np.float64]:<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(dtype_util.name(x.dtype)))<EOL><DEDENT>return ndtri((x + <NUM_LIT:1.>) / <NUM_LIT>) / np.sqrt(<NUM_LIT>)<EOL><DEDENT>

The inverse function for erf, the error function.

    Args:
      x: `Tensor` of type `float32`, `float64`.
      name: Python string. A name for the operation (default="erfinv").

    Returns:
      x: `Tensor` with `dtype=x.dtype`.

    Raises:
      TypeError: if `x` is not floating-type.


with tf.name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>lower_solution = -np.log(<NUM_LIT>) + x<EOL>safe_exp_neg_x = tf.exp(-tf.abs(x))<EOL>upper_solution = tf.math.log1p(-<NUM_LIT:0.5> * safe_exp_neg_x)<EOL>return tf.where(x < <NUM_LIT:0.>, lower_solution, upper_solution)<EOL><DEDENT>

Log Laplace distribution function.

    This function calculates `Log[L(x)]`, where `L(x)` is the cumulative
    distribution function of the Laplace distribution, i.e.

    ```L(x) := 0.5 * int_{-infty}^x e^{-|t|} dt```

    For numerical accuracy, `L(x)` is computed in different ways depending on `x`,

    ```
    x <= 0:
      Log[L(x)] = Log[0.5] + x, which is exact

    0 < x:
      Log[L(x)] = Log[1 - 0.5 * e^{-x}], which is exact
    ```

    Args:
      x: `Tensor` of type `float32`, `float64`.
      name: Python string. A name for the operation (default="log_ndtr").

    Returns:
      `Tensor` with `dtype=x.dtype`.

    Raises:
      TypeError: if `x.dtype` is not handled.


def log_ndtr(x, series_order=<NUM_LIT:3>, name="<STR_LIT>"):

if not isinstance(series_order, int):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if series_order < <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if series_order > <NUM_LIT:30>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>with tf.name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>if dtype_util.base_equal(x.dtype, tf.float64):<EOL><INDENT>lower_segment = LOGNDTR_FLOAT64_LOWER<EOL>upper_segment = LOGNDTR_FLOAT64_UPPER<EOL><DEDENT>elif dtype_util.base_equal(x.dtype, tf.float32):<EOL><INDENT>lower_segment = LOGNDTR_FLOAT32_LOWER<EOL>upper_segment = LOGNDTR_FLOAT32_UPPER<EOL><DEDENT>else:<EOL><INDENT>raise TypeError("<STR_LIT>" % x.dtype)<EOL><DEDENT>return tf.where(<EOL>tf.greater(x, upper_segment),<EOL>-_ndtr(-x),  <EOL>tf.where(<EOL>tf.greater(x, lower_segment),<EOL>tf.math.log(_ndtr(tf.maximum(x, lower_segment))),<EOL>_log_ndtr_lower(tf.minimum(x, lower_segment), series_order)))<EOL><DEDENT>

Log Normal distribution function.

    For details of the Normal distribution function see `ndtr`.

    This function calculates `(log o ndtr)(x)` by either calling `log(ndtr(x))` or
    using an asymptotic series. Specifically:
    - For `x > upper_segment`, use the approximation `-ndtr(-x)` based on
      `log(1-x) ~= -x, x << 1`.
    - For `lower_segment < x <= upper_segment`, use the existing `ndtr` technique
      and take a log.
    - For `x <= lower_segment`, we use the series approximation of erf to compute
      the log CDF directly.

    The `lower_segment` is set based on the precision of the input:

    ```
    lower_segment = { -20,  x.dtype=float64
                    { -10,  x.dtype=float32
    upper_segment = {   8,  x.dtype=float64
                    {   5,  x.dtype=float32
    ```

    When `x < lower_segment`, the `ndtr` asymptotic series approximation is:

    ```
       ndtr(x) = scale * (1 + sum) + R_N
       scale   = exp(-0.5 x**2) / (-x sqrt(2 pi))
       sum     = Sum{(-1)^n (2n-1)!! / (x**2)^n, n=1:N}
       R_N     = O(exp(-0.5 x**2) (2N+1)!! / |x|^{2N+3})
    ```

    where `(2n-1)!! = (2n-1) (2n-3) (2n-5) ...  (3) (1)` is a
    [double-factorial](https://en.wikipedia.org/wiki/Double_factorial).

    Args:
      x: `Tensor` of type `float32`, `float64`.
      series_order: Positive Python `integer`. Maximum depth to
        evaluate the asymptotic expansion. This is the `N` above.
      name: Python string. A name for the operation (default="log_ndtr").

    Returns:
      log_ndtr: `Tensor` with `dtype=x.dtype`.

    Raises:
      TypeError: if `x.dtype` is not handled.
      TypeError: if `series_order` is a not Python `integer.`
      ValueError:  if `series_order` is not in `[0, 30]`.


return np.prod(np.arange(n, <NUM_LIT:1>, -<NUM_LIT:2>))<EOL>

The double factorial function for small Python integer `n`.


x_2 = tf.square(x)<EOL>log_scale = -<NUM_LIT:0.5> * x_2 - tf.math.log(-x) - <NUM_LIT:0.5> * np.log(<NUM_LIT> * np.pi)<EOL>return log_scale + tf.math.log(_log_ndtr_asymptotic_series(x, series_order))<EOL>

Asymptotic expansion version of `Log[cdf(x)]`, appropriate for `x<<-1`.


q = sampling_dist_q<EOL>with tf.name_scope(name):<EOL><INDENT>z = _get_samples(q, z, n, seed)<EOL>log_values = log_f(z) + log_p(z) - q.log_prob(z)<EOL>return _logspace_mean(log_values)<EOL><DEDENT>

r"""Importance sampling with a positive function, in log-space.

    With \\(p(z) := exp^{log_p(z)}\\), and \\(f(z) = exp{log_f(z)}\\),
    this `Op` returns

    \\(Log[ n^{-1} sum_{i=1}^n [ f(z_i) p(z_i) / q(z_i) ] ],  z_i ~ q,\\)
    \\(\approx Log[ E_q[ f(Z) p(Z) / q(Z) ] ]\\)
    \\(=       Log[E_p[f(Z)]]\\)

    This integral is done in log-space with max-subtraction to better handle the
    often extreme values that `f(z) p(z) / q(z)` can take on.

    In contrast to `expectation_importance_sampler`, this `Op` returns values in
    log-space.

    User supplies either `Tensor` of samples `z`, or number of samples to draw `n`

    Args:
      log_f: Callable mapping samples from `sampling_dist_q` to `Tensors` with
        shape broadcastable to `q.batch_shape`.
        For example, `log_f` works "just like" `sampling_dist_q.log_prob`.
      log_p:  Callable mapping samples from `sampling_dist_q` to `Tensors` with
        shape broadcastable to `q.batch_shape`.
        For example, `log_p` works "just like" `q.log_prob`.
      sampling_dist_q:  The sampling distribution.
        `tfp.distributions.Distribution`.
        `float64` `dtype` recommended.
        `log_p` and `q` should be supported on the same set.
      z:  `Tensor` of samples from `q`, produced by `q.sample` for some `n`.
      n:  Integer `Tensor`.  Number of samples to generate if `z` is not provided.
      seed:  Python integer to seed the random number generator.
      name:  A name to give this `Op`.

    Returns:
      Logarithm of the importance sampling estimate.  `Tensor` with `shape` equal
        to batch shape of `q`, and `dtype` = `q.dtype`.


dtype = None<EOL>preferred_dtype = (None if preferred_dtype is None<EOL>else tf.as_dtype(preferred_dtype))<EOL>for a in tf.nest.flatten(args_list):<EOL><INDENT>if hasattr(a, '<STR_LIT>'):<EOL><INDENT>dt = tf.as_dtype(a.dtype)<EOL><DEDENT>else:<EOL><INDENT>continue<EOL><DEDENT>if dtype is None:<EOL><INDENT>dtype = dt<EOL><DEDENT>elif dtype != dt:<EOL><INDENT>raise TypeError('<STR_LIT>'.format(dtype, dt))<EOL><DEDENT><DEDENT>if dtype is None and preferred_dtype is None:<EOL><INDENT>return None<EOL><DEDENT>return (preferred_dtype if dtype is None else dtype).as_numpy_dtype<EOL>

Returns explict dtype from `args_list` if exists, else preferred_dtype.


def _reduce_logsumexp(input_tensor, axis=None, keepdims=False, name=None):

try:<EOL><INDENT>return scipy_special.logsumexp(<EOL>input_tensor, axis=_astuple(axis), keepdims=keepdims)<EOL><DEDENT>except NotImplementedError:<EOL><INDENT>m = _max_mask_non_finite(input_tensor, axis=axis, keepdims=True)<EOL>y = input_tensor - m<EOL>y = np.exp(y, out=y)<EOL>return m + np.log(np.sum(y, axis=_astuple(axis), keepdims=keepdims))<EOL><DEDENT>

Computes `log(sum(exp(input_tensor))) along the specified axis.


<EOL>if not tf.executing_eagerly():<EOL><INDENT>yield<EOL>return<EOL><DEDENT>with super(TestCase, self).assertRaisesWithPredicateMatch(<EOL>exception_type, expected_err_re_or_predicate):<EOL><INDENT>yield<EOL><DEDENT>

Returns a context manager to enclose code expected to raise an exception.

        If the exception is an OpError, the op stack is also included in the message
        predicate search.

        Args:
          exception_type: The expected type of exception that should be raised.
          expected_err_re_or_predicate: If this is callable, it should be a function
            of one argument that inspects the passed-in exception and returns True
            (success) or False (please fail the test). Otherwise, the error message
            is expected to match this regular expression partially.

        Returns:
          A context manager to surround code that is expected to raise an
          exception.


original_spec = tf_inspect.getfullargspec(original_fn)<EOL>new_spec = tf_inspect.getfullargspec(new_fn)<EOL>if original_spec != new_spec:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(<EOL>original_spec, new_spec, original_fn))<EOL><DEDENT>@decorator.decorator<EOL>def wrap(wrapped_fn, *args, **kwargs):<EOL><INDENT>del wrapped_fn<EOL>return new_fn(*args, **kwargs)<EOL><DEDENT>return wrap(original_fn)<EOL>

Wraps new_fn with the doc of original_fn.


def case(pred_fn_pairs, default=None, exclusive=False, name='<STR_LIT>'):

return control_flow_ops._case_helper(  <EOL>cond, pred_fn_pairs, default, exclusive, name, allow_python_preds=True)<EOL>

Like tf.case, except attempts to statically evaluate predicates.

    If any predicate in `pred_fn_pairs` is a bool or has a constant value, the
    associated callable will be called or omitted depending on its value.
    Otherwise this functions like tf.case.

    Args:
      pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor and a
                     callable which returns a list of tensors.
      default: Optional callable that returns a list of tensors.
      exclusive: True iff at most one predicate is allowed to evaluate to `True`.
      name: A name for this operation (optional).

    Returns:
      The tensors returned by the first pair whose predicate evaluated to True, or
      those returned by `default` if none does.

    Raises:
      TypeError: If `pred_fn_pairs` is not a list/dictionary.
      TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples.
      TypeError: If `fns[i]` is not callable for any i, or `default` is not
                 callable.


def _finish_prob_for_one_fiber(self, y, x, ildj, event_ndims,<EOL>**distribution_kwargs):

x = self._maybe_rotate_dims(x, rotate_right=True)<EOL>prob = self.distribution.prob(x, **distribution_kwargs)<EOL>if self._is_maybe_event_override:<EOL><INDENT>prob = tf.reduce_prod(input_tensor=prob, axis=self._reduce_event_indices)<EOL><DEDENT>prob *= tf.exp(tf.cast(ildj, prob.dtype))<EOL>if self._is_maybe_event_override and isinstance(event_ndims, int):<EOL><INDENT>tensorshape_util.set_shape(<EOL>prob,<EOL>tf.broadcast_static_shape(<EOL>tensorshape_util.with_rank_at_least(y.shape, <NUM_LIT:1>)[:-event_ndims],<EOL>self.batch_shape))<EOL><DEDENT>return prob<EOL>

Finish computation of prob on one element of the inverse image.


def _pick_scalar_condition(pred, cond_true, cond_false):

<EOL>pred_ = tf.get_static_value(tf.convert_to_tensor(value=pred))<EOL>if pred_ is None:<EOL><INDENT>return tf.where(pred, cond_true, cond_false)<EOL><DEDENT>return cond_true if pred_ else cond_false<EOL>

Convenience function which chooses the condition based on the predicate.


return (kwargs.get("<STR_LIT>", {}),<EOL>kwargs.get("<STR_LIT>", {}))<EOL>


with self._name_scope(name):<EOL><INDENT>return self.high - self.low<EOL><DEDENT>


def _compute_non_batch_kl(mu_a, sigma_a, mu_b, sigma_b):

<EOL>sigma_b_inv = np.linalg.inv(sigma_b)<EOL>t = np.trace(sigma_b_inv.dot(sigma_a))<EOL>q = (mu_b - mu_a).dot(sigma_b_inv).dot(mu_b - mu_a)<EOL>k = mu_a.shape[<NUM_LIT:0>]<EOL>l = np.log(np.linalg.det(sigma_b) / np.linalg.det(sigma_a))<EOL>return <NUM_LIT:0.5> * (t + q - k + l)<EOL>

Non-batch KL for N(mu_a, sigma_a), N(mu_b, sigma_b).


<EOL>scale = self.scale + tf.zeros_like(self.concentration)<EOL>is_invalid = x < scale<EOL>scale = scale + tf.zeros_like(x)<EOL>x = x + tf.zeros_like(scale)<EOL>y = f(tf.where(is_invalid, scale, x))<EOL>if alt == <NUM_LIT:0.>:<EOL><INDENT>alt = tf.zeros_like(y)<EOL><DEDENT>elif alt == <NUM_LIT:1.>:<EOL><INDENT>alt = tf.ones_like(y)<EOL><DEDENT>else:<EOL><INDENT>alt = tf.fill(<EOL>dims=tf.shape(input=y),<EOL>value=dtype_util.as_numpy_dtype(self.dtype)(alt))<EOL><DEDENT>return tf.where(is_invalid, alt, y)<EOL>

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

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

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


return self._scale<EOL>

Scale parameter and also the lower bound of the support.


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


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

Construct Multivariate Normal distribution on `R^k`.

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

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

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

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

        where:

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

        Additional leading dimensions (if any) will index batches.

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

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

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


return self._logits<EOL>


return self._probs<EOL>


return self._concentration<EOL>


return self._scale<EOL>


return self._high<EOL>


return self._peak<EOL>

Peak of the distribution. Lies in the interval.


return (self.peak - self.low) / (self.high - self.low)<EOL>


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([concentration1, concentration0],<EOL>tf.float32)<EOL>self._concentration1 = self._maybe_assert_valid_concentration(<EOL>tf.convert_to_tensor(<EOL>value=concentration1, name="<STR_LIT>", dtype=dtype),<EOL>validate_args)<EOL>self._concentration0 = self._maybe_assert_valid_concentration(<EOL>tf.convert_to_tensor(<EOL>value=concentration0, name="<STR_LIT>", dtype=dtype),<EOL>validate_args)<EOL>self._total_concentration = self._concentration1 + self._concentration0<EOL><DEDENT>super(Beta, 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=[<EOL>self._concentration1, self._concentration0,<EOL>self._total_concentration<EOL>],<EOL>name=name)<EOL>

Initialize a batch of Beta distributions.

        Args:
          concentration1: Positive floating-point `Tensor` indicating mean
            number of successes; aka "alpha". Implies `self.dtype` and
            `self.batch_shape`, i.e.,
            `concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
          concentration0: Positive floating-point `Tensor` indicating mean
            number of failures; aka "beta". Otherwise has same semantics as
            `concentration1`.
          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.


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

def delta(fn, is_property=True):<EOL><INDENT>fn1 = getattr(d1, fn)<EOL>fn2 = getattr(d2, fn)<EOL>return (fn2 - fn1) if is_property else (fn2() - fn1())<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return (delta("<STR_LIT>", is_property=False) -<EOL>tf.math.digamma(d1.concentration1) * delta("<STR_LIT>") -<EOL>tf.math.digamma(d1.concentration0) * delta("<STR_LIT>") +<EOL>(tf.math.digamma(d1.total_concentration) *<EOL>delta("<STR_LIT>")))<EOL><DEDENT>

Calculate the batchwise KL divergence KL(d1 || d2) with d1 and d2 Beta.

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

    Returns:
      Batchwise KL(d1 || d2)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>logits, probs, validate_args=validate_args, name=name)<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_positive(self._probs)] if validate_args else []):<EOL><INDENT>self._probs = tf.identity(self._probs, name="<STR_LIT>")<EOL><DEDENT><DEDENT>super(Geometric, 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._probs, self._logits],<EOL>name=name)<EOL>

Construct Geometric distributions.

        Args:
          logits: Floating-point `Tensor` with shape `[B1, ..., Bb]` where `b >= 0`
            indicates the number of batch dimensions. Each entry represents logits
            for the probability of success for independent Geometric distributions
            and must be in the range `(-inf, inf]`. Only one of `logits` or `probs`
            should be specified.
          probs: Positive floating-point `Tensor` with shape `[B1, ..., Bb]`
            where `b >= 0` indicates the number of batch dimensions. Each entry
            represents the probability of success for independent Geometric
            distributions and must be in the range `(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.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._distribution_fn = distribution_fn<EOL>self._sample0 = sample0<EOL>self._distribution0 = (distribution_fn() if sample0 is None<EOL>else distribution_fn(sample0))<EOL>if num_steps is None:<EOL><INDENT>num_steps = tensorshape_util.num_elements(<EOL>self._distribution0.event_shape)<EOL>if num_steps is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT><DEDENT>if num_steps < <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>".format(num_steps))<EOL><DEDENT>self._num_steps = num_steps<EOL><DEDENT>super(Autoregressive, self).__init__(<EOL>dtype=self._distribution0.dtype,<EOL>reparameterization_type=self._distribution0.reparameterization_type,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=self._distribution0._graph_parents,  <EOL>name=name)<EOL>

Construct an `Autoregressive` distribution.

        Args:
          distribution_fn: Python `callable` which constructs a
            `tfd.Distribution`-like instance from a `Tensor` (e.g.,
            `sample0`). The function must respect the "autoregressive property",
            i.e., there exists a permutation of event such that each coordinate is a
            diffeomorphic function of on preceding coordinates.
          sample0: Initial input to `distribution_fn`; used to
            build the distribution in `__init__` which in turn specifies this
            distribution's properties, e.g., `event_shape`, `batch_shape`, `dtype`.
            If unspecified, then `distribution_fn` should be default constructable.
          num_steps: Number of times `distribution_fn` is composed from samples,
            e.g., `num_steps=2` implies
            `distribution_fn(distribution_fn(sample0).sample(n)).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.
          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.
            Default value: "Autoregressive".

        Raises:
          ValueError: if `num_steps` and
            `num_elements(distribution_fn(sample0).event_shape)` are both `None`.
          ValueError: if `num_steps < 1`.


<EOL>return tf.cast(tf.linalg.det(x) > det_bounds, dtype=x.dtype)<EOL>

Returns whether the input matches the given determinant limit.

    Args:
      x: A floating-point `Tensor` of shape `[B1, ..., Bn, M, M]`.
      det_bounds: A floating-point `Tensor` that must broadcast to shape
        `[B1, ..., Bn]`, giving the desired lower bound on the
        determinants in `x`.

    Returns:
      mask: A floating-point `Tensor` of shape [B1, ..., Bn].  Each
        scalar is 1 if the corresponding matrix had determinant above
        the corresponding bound, otherwise 0.


def _minimum_mean(samples, envelope, low, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [samples, envelope, low]):<EOL><INDENT>dtype = dtype_util.common_dtype([samples, envelope, low], tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>envelope = tf.convert_to_tensor(<EOL>value=envelope, name='<STR_LIT>', dtype=dtype)<EOL>low = tf.convert_to_tensor(value=low, name='<STR_LIT>', dtype=dtype)<EOL>xmin = tf.reduce_min(input_tensor=samples, axis=[<NUM_LIT:0>])<EOL>msg = '<STR_LIT>'<EOL>check_op = tf.compat.v1.assert_greater_equal(xmin, low, message=msg)<EOL>with tf.control_dependencies([check_op]):<EOL><INDENT>return - _do_maximum_mean(-samples, envelope, -low)<EOL><DEDENT><DEDENT>

Returns a stochastic lower bound on the mean of a scalar distribution.

    The idea is that if the true CDF is within an `eps`-envelope of the
    empirical CDF of the samples, and the support is bounded below, then
    the mean is bounded below as well.  In symbols,

    ```none
    sup_x(|F_n(x) - F(x)|) < eps
    ```

    The 0th dimension of `samples` is interpreted as independent and
    identically distributed samples.  The remaining dimensions are
    broadcast together with `envelope` and `low`, and operated on
    separately.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `envelope` and `low`.
      envelope: Floating-point `Tensor` of sizes of admissible CDF
        envelopes (i.e., the `eps` above).
      low: Floating-point `Tensor` of lower bounds on the distributions'
        supports.  `samples >= low`.
      name: A name for this operation (optional).

    Returns:
      bound: Floating-point `Tensor` of lower bounds on the true means.

    Raises:
      InvalidArgumentError: If some `sample` is found to be smaller than
        the corresponding `low`.


def assert_true_mean_in_interval_by_dkwm(<EOL>samples, low, high, expected_low, expected_high,<EOL>false_fail_rate=<NUM_LIT>, name=None):

args_list = [samples, low, high, expected_low, expected_high, false_fail_rate]<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>args_list):<EOL><INDENT>dtype = dtype_util.common_dtype(args_list, tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>low = tf.convert_to_tensor(value=low, name='<STR_LIT>', dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name='<STR_LIT>', dtype=dtype)<EOL>expected_low = tf.convert_to_tensor(<EOL>value=expected_low, name='<STR_LIT>', dtype=dtype)<EOL>expected_high = tf.convert_to_tensor(<EOL>value=expected_high, name='<STR_LIT>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>samples = _check_shape_dominates(<EOL>samples, [low, high, expected_low, expected_high])<EOL>min_mean, max_mean = true_mean_confidence_interval_by_dkwm(<EOL>samples, low, high, false_fail_rate)<EOL>check_confidence_interval_can_intersect = tf.compat.v1.assert_greater_equal(<EOL>max_mean,<EOL>expected_low,<EOL>message='<STR_LIT>'<EOL>'<STR_LIT>')<EOL>with tf.control_dependencies([check_confidence_interval_can_intersect]):<EOL><INDENT>return tf.compat.v1.assert_less_equal(<EOL>min_mean,<EOL>expected_high,<EOL>message='<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT><DEDENT>

Asserts the mean of the given distribution is in the given interval.

    More precisely, fails if there is enough evidence (using the
    [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
    that the mean of the distribution from which the given samples are
    drawn is _outside_ the given interval with statistical significance
    `false_fail_rate` or stronger, otherwise passes.  If you also want
    to check that you are gathering enough evidence that a pass is not
    spurious, see `min_num_samples_for_dkwm_mean_test` and
    `min_discrepancy_of_true_means_detectable_by_dkwm`.

    Note that `false_fail_rate` is a total false failure rate for all
    the assertions in the batch.  As such, if the batch is nontrivial,
    the assertion will insist on stronger evidence to fail any one member.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `low` and `high`.
        The support is bounded: `low <= samples <= high`.
      low: Floating-point `Tensor` of lower bounds on the distributions'
        supports.
      high: Floating-point `Tensor` of upper bounds on the distributions'
        supports.
      expected_low: Floating-point `Tensor` of lower bounds on the
        expected true means.
      expected_high: Floating-point `Tensor` of upper bounds on the
        expected true means.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if any expected mean
        interval does not overlap with the corresponding confidence
        interval.


def assert_multivariate_true_cdf_equal_on_projections_two_sample(<EOL>samples1, samples2, num_projections, event_ndims=<NUM_LIT:1>,<EOL>false_fail_rate=<NUM_LIT>, seed=<NUM_LIT>, name=None):

<EOL>args_list = (<EOL>[samples1, samples2, num_projections, event_ndims, false_fail_rate])<EOL>strm = seed_stream.SeedStream(salt='<STR_LIT>', seed=seed)<EOL>with tf.compat.v1.name_scope(<EOL>name,<EOL>'<STR_LIT>',<EOL>args_list):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[samples1, samples2, false_fail_rate], tf.float32)<EOL>samples1 = tf.convert_to_tensor(<EOL>value=samples1, name='<STR_LIT>', dtype=dtype)<EOL>samples2 = tf.convert_to_tensor(<EOL>value=samples2, name='<STR_LIT>', dtype=dtype)<EOL>num_projections = tf.convert_to_tensor(<EOL>value=num_projections, name='<STR_LIT>')<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>tf.compat.v1.assert_scalar(false_fail_rate)  <EOL>compatible_samples = tf.compat.v1.assert_equal(<EOL>tf.shape(input=samples1)[<NUM_LIT:1>:],<EOL>tf.shape(input=samples2)[<NUM_LIT:1>:])<EOL>with tf.control_dependencies([compatible_samples]):<EOL><INDENT>event_shape = tf.shape(input=samples1)[-event_ndims:]<EOL>random_projections = _random_unit_hypersphere(<EOL>[num_projections], event_shape, dtype=dtype, seed=strm())<EOL>last_axes = list(range(-<NUM_LIT:1>, -(event_ndims+<NUM_LIT:1>), -<NUM_LIT:1>))<EOL>proj1 = tf.tensordot(samples1, random_projections, [last_axes, last_axes])<EOL>proj2 = tf.tensordot(samples2, random_projections, [last_axes, last_axes])<EOL>return assert_true_cdf_equal_by_dkwm_two_sample(<EOL>proj1, proj2, false_fail_rate=false_fail_rate)<EOL><DEDENT><DEDENT>

Asserts the given multivariate distributions are equal.

    The test is a 1-D DKWM-style test of equality in distribution along the given
    number of random projections.  This is of course imperfect, but can behave
    reasonably, especially if `num_projections` is significantly more than the
    dimensionality of the sample space.

    More precisely, the test
    (i) assumes the event shape is given by the trailing `event_ndims` dimensions
        in each `samples` Tensor;
    (ii) generates `num_projections` random projections from this space to scalar;
    (iii) fails if there is enough evidence (using the
        [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
        (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
        along any of these projections that `samples1` and `samples2` come from
        different true distributions.

    This test works as written even if the distribution in question has atoms
    (e.g., is discrete).

    Note that the top dimension of each `samples` is treated as iid, and the
    bottom `event_ndims` dimensions are projected to scalar.  The remaining
    dimensions, if any, are treated as batch dimensions, and `false_fail_rate` is
    a total false failure rate for all the assertions in the batch (and all
    projections).  As such, if the batch is nontrivial, the assertion will insist
    on stronger evidence to fail any one member.

    A note on experiment design: This test boils down to `num_projections`
    two-sample CDF equality tests.  As such, one can compute the number of samples
    to draw or the detectable discrepancy (along any projection) using
    `min_num_samples_for_dkwm_cdf_two_sample_test` and
    `min_discrepancy_of_true_cdfs_detectable_by_dkwm_two_sample` respectively,
    just being sure to divide the false failure rate by the number of projections
    requested (no need to adjust the false pass rate).

    Args:
      samples1: Tensor of shape [n] + B + E.  Samples from some (batch of)
        distribution(s) of interest.  Assumed IID across the 0 dimension.
      samples2: Tensor of shape [m] + B + E.  Samples from some (batch of)
        distribution(s) of interest.  Assumed IID across the 0 dimension.
      num_projections: Scalar integer Tensor.  Number of projections to use.  Each
        projection will be a random direction on the unit hypershpere of shape E.
      event_ndims: Number of trailing dimensions forming the event shape.
        `rank(E) == event_ndims`.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      seed: Optional PRNG seed to use for generating the random projections.
        Changing this from the default should generally not be necessary.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if the two samples do not match
        along any of the generated projections.


def _maximum_mean(samples, envelope, high, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [samples, envelope, high]):<EOL><INDENT>dtype = dtype_util.common_dtype([samples, envelope, high], tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>envelope = tf.convert_to_tensor(<EOL>value=envelope, name='<STR_LIT>', dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name='<STR_LIT>', dtype=dtype)<EOL>xmax = tf.reduce_max(input_tensor=samples, axis=[<NUM_LIT:0>])<EOL>msg = '<STR_LIT>'<EOL>check_op = tf.compat.v1.assert_less_equal(xmax, high, message=msg)<EOL>with tf.control_dependencies([check_op]):<EOL><INDENT>return tf.identity(_do_maximum_mean(samples, envelope, high))<EOL><DEDENT><DEDENT>

Returns a stochastic upper bound on the mean of a scalar distribution.

    The idea is that if the true CDF is within an `eps`-envelope of the
    empirical CDF of the samples, and the support is bounded above, then
    the mean is bounded above as well.  In symbols,

    ```none
    sup_x(|F_n(x) - F(x)|) < eps
    ```

    The 0th dimension of `samples` is interpreted as independent and
    identically distributed samples.  The remaining dimensions are
    broadcast together with `envelope` and `high`, and operated on
    separately.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `envelope` and `high`.
      envelope: Floating-point `Tensor` of sizes of admissible CDF
        envelopes (i.e., the `eps` above).
      high: Floating-point `Tensor` of upper bounds on the distributions'
        supports.  `samples <= high`.
      name: A name for this operation (optional).

    Returns:
      bound: Floating-point `Tensor` of upper bounds on the true means.

    Raises:
      InvalidArgumentError: If some `sample` is found to be larger than
        the corresponding `high`.


def empirical_cdfs(samples, positions, continuity='<STR_LIT:right>',<EOL>dtype=tf.float32, name=None):

if continuity not in ['<STR_LIT:left>', '<STR_LIT:right>']:<EOL><INDENT>msg = '<STR_LIT>'.format(<EOL>continuity)<EOL>raise ValueError(msg)<EOL><DEDENT>with tf.compat.v1.name_scope(name, '<STR_LIT>', [samples, positions]):<EOL><INDENT>n = tf.cast(tf.shape(input=samples)[-<NUM_LIT:1>], dtype=dtype)<EOL>indexes = tf.searchsorted(<EOL>sorted_sequence=samples, values=positions, side=continuity)<EOL>return tf.cast(indexes, dtype=dtype) / n<EOL><DEDENT>

Evaluates the empirical CDF of a batch of samples at a batch of positions.

    If elements of `positions` might be exactly equal to elements of `samples`
    (e.g., if the underlying distribution of interest is discrete), there is a
    difference between the conventional, right-continuous CDF (Pr[X <= x]) and a
    left-continuous variant (Pr[X < x]).  The latter can be accessed by setting
    `continuity='left'`.  The difference between the right-continuous and
    left-continuous CDFs is the empirical pmf at each point, i.e., how many times
    each element of `positions` occurs in its batch of `samples`.

    Note: Returns results parallel to `positions`, i.e., the values of the
    empirical CDF at those points.

    Note: The sample dimension is _last_, and the samples must be _sorted_ within
    each batch.

    Args:
      samples: Tensor of shape `batch + [num_samples]` of samples.  The samples
        must be in ascending order within each batch member.
      positions: Tensor of shape `batch + [m]` of positions where to evaluate the
        CDFs.  The positions need not be sorted.
      continuity: Whether to return a conventional, right-continuous CDF
        (`continuity = 'right'`, default) or a left-continuous CDF (`continuity =
        'left'`).  The value at each point `x` will be `F_n(X <= x)` or
        `F_n(X < x)`, respectively.
      dtype: dtype at which to evaluate the desired empirical CDFs.
      name: A name for this operation (optional).

    Returns:
      cdf: Tensor parallel to `positions`.  For each x in `positions`, gives the
        (right- or left-continuous, per the `continuity` argument) cdf at that
        position.  If `positions` contains duplicates, `cdf` will give each the
        same value.


def assert_true_mean_equal_by_dkwm(<EOL>samples, low, high, expected, false_fail_rate=<NUM_LIT>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[samples, low, high, expected, false_fail_rate]):<EOL><INDENT>return assert_true_mean_in_interval_by_dkwm(<EOL>samples, low, high, expected, expected, false_fail_rate)<EOL><DEDENT>

Asserts the mean of the given distribution is as expected.

    More precisely, fails if there is enough evidence (using the
    [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
    that the true mean of some distribution from which the given samples are
    drawn is _not_ the given expected mean with statistical significance
    `false_fail_rate` or stronger, otherwise passes.  If you also want to
    check that you are gathering enough evidence that a pass is not
    spurious, see `min_num_samples_for_dkwm_mean_test` and
    `min_discrepancy_of_true_means_detectable_by_dkwm`.

    Note that `false_fail_rate` is a total false failure rate for all
    the assertions in the batch.  As such, if the batch is nontrivial,
    the assertion will insist on stronger evidence to fail any one member.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `low` and `high`.
        The support is bounded: `low <= samples <= high`.
      low: Floating-point `Tensor` of lower bounds on the distributions'
        supports.
      high: Floating-point `Tensor` of upper bounds on the distributions'
        supports.
      expected: Floating-point `Tensor` of expected true means.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if any expected mean is
        outside the corresponding confidence interval.


def min_discrepancy_of_true_cdfs_detectable_by_dkwm(<EOL>n, false_fail_rate, false_pass_rate, name=None):

with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[n, false_fail_rate, false_pass_rate]):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[n, false_fail_rate, false_pass_rate], tf.float32)<EOL>n = tf.convert_to_tensor(value=n, name='<STR_LIT:n>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>false_pass_rate = tf.convert_to_tensor(<EOL>value=false_pass_rate, name='<STR_LIT>', dtype=dtype)<EOL>sampling_envelope = _dkwm_cdf_envelope(n, false_pass_rate)<EOL>itemwise_false_fail_rate = _itemwise_error_rate(<EOL>total_rate=false_fail_rate, param_tensors=[n])<EOL>analysis_envelope = _dkwm_cdf_envelope(n, itemwise_false_fail_rate)<EOL>return sampling_envelope + analysis_envelope<EOL><DEDENT>

Returns the minimum CDF discrepancy that a DKWM-based test can detect.

    DKWM is the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval).

    Note that `false_fail_rate` is a total false failure rate for all
    the tests in the batch.  As such, if the batch is nontrivial, each
    member will demand more samples.  The `false_pass_rate` is also
    interpreted as a total, but is treated asymmetrically: If each test
    in the batch detects its corresponding discrepancy with probability
    at least `1 - false_pass_rate`, then running all those tests and
    failing if any one fails will jointly detect all those discrepancies
    with the same `false_pass_rate`.

    Args:
      n: `Tensor` of numbers of samples to be drawn from the distributions
        of interest.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of false failures.
      false_pass_rate: *Scalar* floating-point `Tensor` admissible rate
        of false passes.
      name: A name for this operation (optional).

    Returns:
      discr: `Tensor` of lower bounds on the K-S distances between true
         CDFs detectable by a DKWM-based test.

    For each batch member `i`, of `K` total, drawing `n[i]` samples from some
    scalar distribution is enough to detect a K-S distance in CDFs of size
    `discr[i]` or more.  Specifically, we guarantee that (a) if the true CDF is
    the expected CDF, then `assert_true_cdf_equal_by_dkwm` will fail with
    probability at most `false_fail_rate / K` (which amounts to `false_fail_rate`
    if applied to the whole batch at once), and (b) if the true CDF differs from
    the expected CDF by at least `discr[i]`, `assert_true_cdf_equal_by_dkwm` will
    pass with probability at most `false_pass_rate`.

    The detectable discrepancy scales as

    - `O(1 / sqrt(n[i]))`,
    - `O(-log(false_fail_rate/K))`, and
    - `O(-log(false_pass_rate))`.


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

Construct HalfNormals with scale `scale`.

        Args:
          scale: Floating point tensor; the scales 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.


def _compute_non_batch_kl(self, loc_a, scale_a, loc_b, scale_b):

<EOL>covariance_a = np.dot(scale_a, scale_a.T)<EOL>covariance_b = np.dot(scale_b, scale_b.T)<EOL>covariance_b_inv = np.linalg.inv(covariance_b)<EOL>t = np.trace(covariance_b_inv.dot(covariance_a))<EOL>q = (loc_b - loc_a).dot(covariance_b_inv).dot(loc_b - loc_a)<EOL>k = loc_a.shape[<NUM_LIT:0>]<EOL>l = np.log(np.linalg.det(covariance_b) / np.linalg.det(covariance_a))<EOL>return <NUM_LIT:0.5> * (t + q - k + l)<EOL>

Non-batch KL for N(loc_a, scale_a), N(loc_b, scale_b).


@hps.composite<EOL>def broadcasting_shapes(draw, batch_shape, param_names):

n = len(param_names)<EOL>return dict(zip(draw(hps.permutations(param_names)),<EOL>draw(tfp_test_util.broadcasting_shapes(batch_shape, n))))<EOL>

Draws a set of parameter batch shapes that broadcast to `batch_shape`.

    For each parameter we need to choose its batch rank, and whether or not each
    axis i is 1 or batch_shape[i]. This function chooses a set of shapes that
    have possibly mismatched ranks, and possibly broadcasting axes, with the
    promise that the broadcast of the set of all shapes matches `batch_shape`.

    Args:
      draw: Hypothesis sampler.
      batch_shape: `tf.TensorShape`, the target (fully-defined) batch shape .
      param_names: Iterable of `str`, the parameters whose batch shapes need
        determination.

    Returns:
      param_batch_shapes: `dict` of `str->tf.TensorShape` where the set of
          shapes broadcast to `batch_shape`. The shapes are fully defined.


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

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>distance = tf.abs(a.loc - b.loc)<EOL>ratio = a.scale / b.scale<EOL>return (-tf.math.log(ratio) - <NUM_LIT:1> + distance / b.scale +<EOL>ratio * tf.exp(-distance / a.scale))<EOL><DEDENT>

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

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

    Returns:
      Batchwise KL(a || b)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._assertions = _maybe_validate_distributions(<EOL>distributions, dtype_override, validate_args)<EOL>if dtype_override is not None:<EOL><INDENT>dtype = dtype_override<EOL><DEDENT>else:<EOL><INDENT>dtype = set(<EOL>dtype_util.base_dtype(d.dtype)<EOL>for d in distributions<EOL>if d.dtype is not None)<EOL>if len(dtype) == <NUM_LIT:0>:  <EOL><INDENT>dtype = tf.float32<EOL><DEDENT>elif len(dtype) == <NUM_LIT:1>:<EOL><INDENT>dtype = dtype.pop()<EOL><DEDENT>else:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT><DEDENT>reparameterization_type = set(d.reparameterization_type<EOL>for d in distributions)<EOL>reparameterization_type = (reparameterization_type.pop()<EOL>if len(reparameterization_type) == <NUM_LIT:1><EOL>else reparameterization.NOT_REPARAMETERIZED)<EOL>self._distributions = distributions<EOL>super(Blockwise, self).__init__(<EOL>dtype=dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization_type,<EOL>parameters=parameters,<EOL>graph_parents=_flatten(d._graph_parents for d in distributions),  <EOL>name=name)<EOL><DEDENT>

Construct the `Blockwise` distribution.

        Args:
          distributions: Python `list` of `tfp.distributions.Distribution`
            instances. All distribution instances must have the same `batch_shape`
            and all must have `event_ndims==1`, i.e., be vector-variate
            distributions.
          dtype_override: samples of `distributions` will be cast to this `dtype`.
            If unspecified, all `distributions` must have the same `dtype`.
            Default value: `None` (i.e., do not cast).
          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.


@kullback_leibler.RegisterKL(Blockwise, Blockwise)<EOL>def _kl_blockwise_blockwise(b0, b1, name=None):

if len(b0.distributions) != len(b1.distributions):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>b0_event_sizes = [_event_size(d) for d in b0.distributions]<EOL>b1_event_sizes = [_event_size(d) for d in b1.distributions]<EOL>assertions = []<EOL>message = ('<STR_LIT>'<EOL>'<STR_LIT>')<EOL>if (all(isinstance(event_size, int) for event_size in b0_event_sizes) and<EOL>all(isinstance(event_size, int) for event_size in b1_event_sizes)):<EOL><INDENT>if b0_event_sizes != b1_event_sizes:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>if b0.validate_args or b1.validate_args:<EOL><INDENT>assertions.extend(<EOL>assert_util.assert_equal(  <EOL>e1, e2, message=message)<EOL>for e1, e2 in zip(b0_event_sizes, b1_event_sizes))<EOL><DEDENT><DEDENT>with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>return sum([<EOL>kullback_leibler.kl_divergence(d1, d2) for d1, d2 in zip(<EOL>b0.distributions, b1.distributions)])<EOL><DEDENT><DEDENT>

Calculate the batched KL divergence KL(b0 || b1) with b0 and b1 Blockwise distributions.

    Args:
      b0: instance of a Blockwise distribution object.
      b1: instance of a Blockwise distribution object.
      name: (optional) Name to use for created operations. Default is
        "kl_blockwise_blockwise".

    Returns:
      kl_blockwise_blockwise: `Tensor`. The batchwise KL(b0 || b1).


def kalman_transition(filtered_mean, filtered_cov,<EOL>transition_matrix, transition_noise):

predicted_mean = _propagate_mean(filtered_mean,<EOL>transition_matrix,<EOL>transition_noise)<EOL>predicted_cov = _propagate_cov(filtered_cov,<EOL>transition_matrix,<EOL>transition_noise)<EOL>return predicted_mean, predicted_cov<EOL>

Propagate a filtered distribution through a transition model.


def kalman_filter_step(state, elems_t):<EOL><INDENT>"""<STR_LIT>"""<EOL>if isinstance(elems_t, tuple):<EOL><INDENT>x_t, mask_t = elems_t<EOL><DEDENT>else:<EOL><INDENT>x_t = elems_t<EOL>mask_t = None<EOL><DEDENT>observation_matrix = get_observation_matrix_for_timestep(state.timestep)<EOL>observation_noise = get_observation_noise_for_timestep(state.timestep)<EOL>if mask_t is not None:<EOL><INDENT>x_expected = _propagate_mean(state.predicted_mean,<EOL>observation_matrix,<EOL>observation_noise) * tf.ones_like(x_t)<EOL>x_t = tf.where(<EOL>tf.broadcast_to(mask_t, tf.shape(input=x_expected)), x_expected,<EOL>tf.broadcast_to(x_t, tf.shape(input=x_expected)))<EOL><DEDENT>(filtered_mean,<EOL>filtered_cov,<EOL>observation_dist) = linear_gaussian_update(<EOL>state.predicted_mean, state.predicted_cov,<EOL>observation_matrix, observation_noise,<EOL>x_t)<EOL>log_marginal_likelihood = observation_dist.log_prob(x_t[..., <NUM_LIT:0>])<EOL>if mask_t is not None:<EOL><INDENT>filtered_mean = tf.where(<EOL>tf.broadcast_to(mask_t, tf.shape(input=filtered_mean)),<EOL>state.predicted_mean, filtered_mean)<EOL>filtered_cov = tf.where(<EOL>tf.broadcast_to(mask_t, tf.shape(input=filtered_cov)),<EOL>state.predicted_cov, filtered_cov)<EOL>log_marginal_likelihood = tf.where(<EOL>tf.broadcast_to(mask_t[..., <NUM_LIT:0>, <NUM_LIT:0>],<EOL>tf.shape(input=log_marginal_likelihood)),<EOL>tf.zeros_like(log_marginal_likelihood),<EOL>log_marginal_likelihood)<EOL><DEDENT>predicted_mean, predicted_cov = kalman_transition(<EOL>filtered_mean,<EOL>filtered_cov,<EOL>get_transition_matrix_for_timestep(state.timestep),<EOL>get_transition_noise_for_timestep(state.timestep))<EOL>return KalmanFilterState(<EOL>filtered_mean, filtered_cov,<EOL>predicted_mean, predicted_cov,<EOL>observation_dist.mean()[..., tf.newaxis],<EOL>observation_dist.covariance(),<EOL>log_marginal_likelihood,<EOL>state.timestep+<NUM_LIT:1>)<EOL><DEDENT>return kalman_filter_step<EOL>

Build a callable that performs one step of Kalman filtering.

    Args:
      get_transition_matrix_for_timestep: callable taking a timestep
        as an integer `Tensor` argument, and returning a `LinearOperator`
        of shape `[latent_size, latent_size]`.
      get_transition_noise_for_timestep: callable taking a timestep as
        an integer `Tensor` argument, and returning a
        `MultivariateNormalLinearOperator` of event shape
        `[latent_size]`.
      get_observation_matrix_for_timestep: callable taking a timestep
        as an integer `Tensor` argument, and returning a `LinearOperator`
        of shape `[observation_size, observation_size]`.
      get_observation_noise_for_timestep: callable taking a timestep as
        an integer `Tensor` argument, and returning a
        `MultivariateNormalLinearOperator` of event shape
        `[observation_size]`.

    Returns:
      kalman_filter_step: a callable that updates a KalmanFilterState
        from timestep `t-1` to `t`.


return self._scale<EOL>


return self._logits<EOL>


<EOL>ones = tf.ones_like(n - k)<EOL>k_eq_n = tf.equal(k, n)<EOL>safe_dn = tf.where(k_eq_n, ones, n - k)<EOL>dk = tf.math.betainc(a=safe_dn, b=k + <NUM_LIT:1>, x=<NUM_LIT:1> - p)<EOL>return tf.where(k_eq_n, ones, dk)<EOL>

The binomial cumulative distribution function.

    Args:
      k: floating point `Tensor`.
      n: floating point `Tensor`.
      p: floating point `Tensor`.

    Returns:
      `sum_{j=0}^k p^j (1 - p)^(n - j)`.


with tf.name_scope("<STR_LIT>"):<EOL><INDENT>return tf.math.log1p(-tf.exp(small - big)) + big<EOL><DEDENT>

Stable evaluation of `Log[exp{big} - exp{small}]`.

    To work correctly, we should have the pointwise relation:  `small <= big`.

    Args:
      big: Floating-point `Tensor`
      small: Floating-point `Tensor` with same `dtype` as `big` and broadcastable
        shape.

    Returns:
      `Tensor` of same `dtype` of `big` and broadcast shape.


<EOL>logsf_y = self.log_survival_function(y)<EOL>logsf_y_minus_1 = self.log_survival_function(y - <NUM_LIT:1>)<EOL>logcdf_y = self.log_cdf(y)<EOL>logcdf_y_minus_1 = self.log_cdf(y - <NUM_LIT:1>)<EOL>big = tf.where(logsf_y < logcdf_y, logsf_y_minus_1, logcdf_y)<EOL>small = tf.where(logsf_y < logcdf_y, logsf_y, logcdf_y_minus_1)<EOL>return _logsum_expbig_minus_expsmall(big, small)<EOL>

Compute log_prob(y) using log survival_function and cdf together.


return self._dist<EOL>


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>if (rate is None) == (log_rate is None):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>elif log_rate is None:<EOL><INDENT>rate = tf.convert_to_tensor(<EOL>value=rate,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([rate], preferred_dtype=tf.float32))<EOL>if not dtype_util.is_floating(rate.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>".format(<EOL>dtype_util.name(rate.dtype)))<EOL><DEDENT>with tf.control_dependencies(<EOL>[assert_util.assert_positive(rate)] if validate_args else []):<EOL><INDENT>self._rate = tf.identity(rate, name="<STR_LIT>")<EOL>self._log_rate = tf.math.log(rate, name="<STR_LIT>")<EOL><DEDENT><DEDENT>else:<EOL><INDENT>log_rate = tf.convert_to_tensor(<EOL>value=log_rate,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([log_rate], tf.float32))<EOL>if not dtype_util.is_floating(log_rate.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>".format(<EOL>dtype_util.name(log_rate.dtype)))<EOL><DEDENT>self._rate = tf.exp(log_rate, name="<STR_LIT>")<EOL>self._log_rate = tf.convert_to_tensor(value=log_rate, name="<STR_LIT>")<EOL><DEDENT><DEDENT>self._interpolate_nondiscrete = interpolate_nondiscrete<EOL>super(Poisson, self).__init__(<EOL>dtype=self._rate.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._rate],<EOL>name=name)<EOL>

Initialize a batch of Poisson distributions.

        Args:
          rate: Floating point tensor, the rate parameter. `rate` must be positive.
            Must specify exactly one of `rate` and `log_rate`.
          log_rate: Floating point tensor, the log of the rate parameter.
            Must specify exactly one of `rate` and `log_rate`.
          interpolate_nondiscrete: Python `bool`. When `False`,
            `log_prob` returns `-inf` (and `prob` returns `0`) for non-integer
            inputs. When `True`, `log_prob` evaluates the continuous function
            `k * log_rate - lgamma(k+1) - rate`, which matches the Poisson pmf
            at integer arguments `k` (note that this function is not itself
            a normalized probability log-density).
            Default value: `True`.
          validate_args: Python `bool`. 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`. 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.

        Raises:
          ValueError: if none or both of `rate`, `log_rate` are specified.
          TypeError: if `rate` is not a float-type.
          TypeError: if `log_rate` is not a float-type.


return self._rate<EOL>


def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None):

ds = []<EOL>values_out = []<EOL>seed = seed_stream.SeedStream('<STR_LIT>', seed)<EOL>gen = self._model()<EOL>index = <NUM_LIT:0><EOL>d = next(gen)<EOL>try:<EOL><INDENT>while True:<EOL><INDENT>actual_distribution = d.distribution if isinstance(d, self.Root) else d<EOL>ds.append(actual_distribution)<EOL>if (value is not None and len(value) > index and<EOL>value[index] is not None):<EOL><INDENT>seed()<EOL>next_value = value[index]<EOL><DEDENT>else:<EOL><INDENT>next_value = actual_distribution.sample(<EOL>sample_shape=sample_shape if isinstance(d, self.Root) else (),<EOL>seed=seed())<EOL><DEDENT>values_out.append(next_value)<EOL>index += <NUM_LIT:1><EOL>d = gen.send(next_value)<EOL><DEDENT><DEDENT>except StopIteration:<EOL><INDENT>pass<EOL><DEDENT>return ds, values_out<EOL>

Executes `model`, creating both samples and distributions.


def _verifyPdfWithNumpy(self, vmf, atol=<NUM_LIT>):

dim = tf.compat.dimension_value(vmf.event_shape[-<NUM_LIT:1>])<EOL>nsamples = <NUM_LIT:10><EOL>sample_shape = [nsamples] + tensorshape_util.as_list(<EOL>vmf.batch_shape) + [dim]<EOL>uniforms = np.random.randn(*sample_shape)<EOL>uniforms /= np.linalg.norm(uniforms, axis=-<NUM_LIT:1>, keepdims=True)<EOL>uniforms = uniforms.astype(dtype_util.as_numpy_dtype(vmf.dtype))<EOL>vmf_samples = vmf.sample(<EOL>sample_shape=[nsamples], seed=tfp_test_util.test_seed())<EOL>samples = tf.concat([uniforms, vmf_samples], axis=<NUM_LIT:0>)<EOL>samples = tf.debugging.check_numerics(samples, '<STR_LIT>')<EOL>samples = self.evaluate(samples)<EOL>log_prob = vmf.log_prob(samples)<EOL>log_prob = tf.debugging.check_numerics(log_prob, '<STR_LIT>')<EOL>try:<EOL><INDENT>from scipy.special import gammaln  <EOL>from scipy.special import ive  <EOL><DEDENT>except ImportError:<EOL><INDENT>tf.compat.v1.logging.warn('<STR_LIT>')<EOL>return<EOL><DEDENT>conc = self.evaluate(vmf.concentration)<EOL>mean_dir = self.evaluate(vmf.mean_direction)<EOL>log_true_sphere_surface_area = (<EOL>np.log(<NUM_LIT:2>) + (dim / <NUM_LIT:2>) * np.log(np.pi) - gammaln(dim / <NUM_LIT:2>))<EOL>expected = (<EOL>conc * np.sum(samples * mean_dir, axis=-<NUM_LIT:1>) +<EOL>np.where(conc > <NUM_LIT:0>,<EOL>(dim / <NUM_LIT:2> - <NUM_LIT:1>) * np.log(conc) -<EOL>(dim / <NUM_LIT:2>) * np.log(<NUM_LIT:2> * np.pi) -<EOL>np.log(ive(dim / <NUM_LIT:2> - <NUM_LIT:1>, conc)) -<EOL>np.abs(conc),<EOL>-log_true_sphere_surface_area))<EOL>self.assertAllClose(expected, self.evaluate(log_prob),<EOL>atol=atol)<EOL>

Verifies log_prob evaluations with numpy/scipy.

        Both uniform random points and sampled points are evaluated.

        Args:
          vmf: A `tfp.distributions.VonMisesFisher` instance.
          atol: Absolute difference tolerable.


return self.bijector.concentration0<EOL>

Concentration parameter associated with a `0` outcome.


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

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

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

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

    Returns:
      Batchwise KL(a || b)


def _broadcast_cat_event_and_params(event, params, base_dtype):

if dtype_util.is_integer(event.dtype):<EOL><INDENT>pass<EOL><DEDENT>elif dtype_util.is_floating(event.dtype):<EOL><INDENT>event = tf.cast(event, dtype=tf.int32)<EOL><DEDENT>else:<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(base_dtype))<EOL><DEDENT>shape_known_statically = (<EOL>tensorshape_util.rank(params.shape) is not None and<EOL>tensorshape_util.is_fully_defined(params.shape[:-<NUM_LIT:1>]) and<EOL>tensorshape_util.is_fully_defined(event.shape))<EOL>if not shape_known_statically or params.shape[:-<NUM_LIT:1>] != event.shape:<EOL><INDENT>params *= tf.ones_like(event[..., tf.newaxis],<EOL>dtype=params.dtype)<EOL>params_shape = tf.shape(input=params)[:-<NUM_LIT:1>]<EOL>event *= tf.ones(params_shape, dtype=event.dtype)<EOL>if tensorshape_util.rank(params.shape) is not None:<EOL><INDENT>tensorshape_util.set_shape(event, params.shape[:-<NUM_LIT:1>])<EOL><DEDENT><DEDENT>return event, params<EOL>

Broadcasts the event or distribution parameters.


if not isinstance(model, collections.Sequence):<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>type(model).__name__))<EOL><DEDENT>self._dist_fn = model<EOL>self._dist_fn_wrapped, self._dist_fn_args = zip(*[<EOL>_unify_call_signature(i, dist_fn)<EOL>for i, dist_fn in enumerate(model)])<EOL>

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


def _resolve_distribution_names(dist_fn_args, dist_names, leaf_name):

if dist_names is None:<EOL><INDENT>dist_names = []<EOL><DEDENT>else:<EOL><INDENT>dist_names = dist_names.copy()<EOL><DEDENT>n = len(dist_fn_args)<EOL>dist_names.extend([None]*(n - len(dist_names)))<EOL>for i_, args in enumerate(reversed(dist_fn_args)):<EOL><INDENT>if not args:<EOL><INDENT>continue  <EOL><DEDENT>i = n - i_ - <NUM_LIT:1><EOL>for j, arg_name in enumerate(args):<EOL><INDENT>dist_names[i - j - <NUM_LIT:1>] = arg_name<EOL><DEDENT><DEDENT>j = <NUM_LIT:0><EOL>for i_ in range(len(dist_names)):<EOL><INDENT>i = n - i_ - <NUM_LIT:1><EOL>if dist_names[i] is None:<EOL><INDENT>dist_names[i] = leaf_name if j == <NUM_LIT:0> else leaf_name + str(j)<EOL>j += <NUM_LIT:1><EOL><DEDENT><DEDENT>return tuple(dist_names)<EOL>

Uses arg names to resolve distribution names.


if len(d0._dist_fn_wrapped) != len(d1._dist_fn_wrapped):  <EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if (not all(a is None for a in d0._dist_fn_args) or  <EOL>not all(a is None for a in d1._dist_fn_args)):  <EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>return sum(kullback_leibler.kl_divergence(d0_(), d1_())<EOL>for d0_, d1_ in zip(d0._dist_fn_wrapped, d1._dist_fn_wrapped))<EOL><DEDENT>

Calculate the KL divergence between two `JointDistributionSequential`s.

    Args:
      d0: instance of a `JointDistributionSequential` object.
      d1: instance of a `JointDistributionSequential` object.
      name: (optional) Name to use for created operations.
        Default value: `"kl_joint_joint"`.

    Returns:
      kl_joint_joint: `Tensor` The sum of KL divergences between elemental
        distributions of two joint distributions.

    Raises:
      ValueError: when joint distributions have a different number of elemental
        distributions.
      ValueError: when either joint distribution has a distribution with dynamic
        dependency, i.e., when either joint distribution is not a collection of
        independent distributions.


argspec = tf_inspect.getfullargspec(fn)<EOL>args = argspec.args<EOL>if tf_inspect.isclass(fn):<EOL><INDENT>args = args[<NUM_LIT:1>:]  <EOL><DEDENT>if argspec.defaults:<EOL><INDENT>args = args[:-len(argspec.defaults)]<EOL><DEDENT>return tuple(args)<EOL>

Returns the distribution's required args.


def _resolve_graph(self, distribution_names=None, leaf_name='<STR_LIT:x>'):

<EOL>if distribution_names is None or any(self._dist_fn_args):<EOL><INDENT>distribution_names = _resolve_distribution_names(<EOL>self._dist_fn_args, distribution_names, leaf_name)<EOL><DEDENT>if len(set(distribution_names)) != len(distribution_names):<EOL><INDENT>raise ValueError('<STR_LIT>'.format(<EOL>distribution_names))<EOL><DEDENT>if len(distribution_names) != len(self._dist_fn_wrapped):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>return tuple(zip(distribution_names,<EOL>tuple(() if a is None else a for a in self._dist_fn_args)))<EOL>

Creates a `tuple` of `tuple`s of dependencies.

        This function is **experimental**. That said, we encourage its use
        and ask that you report problems to `tfprobability@tensorflow.org`.

        Args:
          distribution_names: `list` of `str` or `None` names corresponding to each
            of `model` elements. (`None`s are expanding into the
            appropriate `str`.)
          leaf_name: `str` used when no maker depends on a particular
            `model` element.

        Returns:
          graph: `tuple` of `(str tuple)` pairs representing the name of each
            distribution (maker) and the names of its dependencies.

        #### Example

        ```python
        d = tfd.JointDistributionSequential([
                         tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
            lambda    e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
                         tfd.Normal(loc=0, scale=2.),
            lambda n, g: tfd.Normal(loc=n, scale=g),
        ])
        d._resolve_graph()
        # ==> (
        #       ('e', ()),
        #       ('g', ('e',)),
        #       ('n', ()),
        #       ('x', ('n', 'g')),
        #     )
        ```


if distribution_util.is_distribution_instance(dist_fn):<EOL><INDENT>return (lambda *_: dist_fn), None<EOL><DEDENT>if not callable(dist_fn):<EOL><INDENT>raise TypeError('<STR_LIT>'<EOL>'<STR_LIT>'.format(dist_fn))<EOL><DEDENT>args = _get_required_args(dist_fn)<EOL>if not args:<EOL><INDENT>return (lambda *_: dist_fn()), ()<EOL><DEDENT>@functools.wraps(dist_fn)<EOL>def dist_fn_wrapped(*xs):<EOL><INDENT>"""<STR_LIT>"""<EOL>if i != len(xs):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>i, dist_fn, i, len(xs)))<EOL><DEDENT>if len(xs) < len(args):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>i, dist_fn, len(args), len(xs)))<EOL><DEDENT>return dist_fn(*reversed(xs[-len(args):]))<EOL><DEDENT>return dist_fn_wrapped, args<EOL>

Creates `dist_fn_wrapped` which calls `dist_fn` with all prev nodes.

    Args:
      i: Python `int` corresponding to position in topologically sorted DAG.
      dist_fn: Python `callable` which takes a subset of previously constructed
        distributions (in reverse order) and produces a new distribution instance.

    Returns:
      dist_fn_wrapped: Python `callable` which takes all previous distributions
        (in non reverse order) and produces a  new distribution instance.
      args: `tuple` of `str` representing the arg names of `dist_fn` (and in non
        wrapped, "natural" order). `None` is returned only if the input is not a
        `callable`.


def _fn(self):<EOL><INDENT>if any(self._dist_fn_args):  <EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>' + attr + '<STR_LIT>'<EOL>'<STR_LIT>'.format(self.model))<EOL><DEDENT>return self._unflatten(getattr(d(), attr)() for d in self._dist_fn_wrapped)  <EOL><DEDENT>return _fn<EOL>

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


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([total_count, logits, probs], tf.float32)<EOL>self._total_count = tf.convert_to_tensor(<EOL>value=total_count, name="<STR_LIT>", dtype=dtype)<EOL>if validate_args:<EOL><INDENT>self._total_count = (<EOL>distribution_util.embed_check_nonnegative_integer_form(<EOL>self._total_count))<EOL><DEDENT>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>logits=logits,<EOL>probs=probs,<EOL>multidimensional=True,<EOL>validate_args=validate_args,<EOL>name=name,<EOL>dtype=dtype)<EOL>self._mean_val = self._total_count[..., tf.newaxis] * self._probs<EOL><DEDENT>super(Multinomial, self).__init__(<EOL>dtype=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._logits, self._probs],<EOL>name=name)<EOL>

Initialize a batch of Multinomial distributions.

        Args:
          total_count: Non-negative floating point tensor with shape broadcastable
            to `[N1,..., Nm]` with `m >= 0`. Defines this as a batch of
            `N1 x ... x Nm` different Multinomial distributions. Its components
            should be equal to integer values.
          logits: Floating point tensor representing unnormalized log-probabilities
            of a positive event with shape broadcastable to
            `[N1,..., Nm, K]` `m >= 0`, and the same dtype as `total_count`. Defines
            this as a batch of `N1 x ... x Nm` different `K` class Multinomial
            distributions. Only one of `logits` or `probs` should be passed in.
          probs: Positive floating point tensor with shape broadcastable to
            `[N1,..., Nm, K]` `m >= 0` and same dtype as `total_count`. Defines
            this as a batch of `N1 x ... x Nm` different `K` class Multinomial
            distributions. `probs`'s components in the last portion of its shape
            should sum to `1`. Only one of `logits` or `probs` should be passed in.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


parameters = dict(locals())<EOL>name = name or "<STR_LIT>" + distribution.name<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._batch_shape_unexpanded = tf.convert_to_tensor(<EOL>value=batch_shape, dtype=tf.int32, name="<STR_LIT>")<EOL>validate_init_args_statically(distribution, self._batch_shape_unexpanded)<EOL>batch_shape, batch_shape_static, runtime_assertions = calculate_reshape(<EOL>distribution.batch_shape_tensor(), self._batch_shape_unexpanded,<EOL>validate_args)<EOL>self._distribution = distribution<EOL>self._batch_shape_ = batch_shape<EOL>self._batch_shape_static = batch_shape_static<EOL>self._runtime_assertions = runtime_assertions<EOL>super(BatchReshape, self).__init__(<EOL>dtype=distribution.dtype,<EOL>reparameterization_type=distribution.reparameterization_type,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=(<EOL>[self._batch_shape_unexpanded] + distribution._graph_parents),  <EOL>name=name)<EOL><DEDENT>

Construct BatchReshape distribution.

        Args:
          distribution: The base distribution instance to reshape. Typically an
            instance of `Distribution`.
          batch_shape: Positive `int`-like vector-shaped `Tensor` representing
            the new shape of the batch dimensions. Up to one dimension may contain
            `-1`, meaning the remainder of the batch 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: The name to give Ops created by the initializer.
            Default value: `"BatchReshape" + distribution.name`.

        Raises:
          ValueError: if `batch_shape` is not a vector.
          ValueError: if `batch_shape` has non-positive elements.
          ValueError: if `batch_shape` size is not the same as a
            `distribution.batch_shape` size.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>power = tf.convert_to_tensor(<EOL>value=power,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([power], preferred_dtype=tf.float32))<EOL>if (not dtype_util.is_floating(power.dtype) or<EOL>dtype_util.base_equal(power.dtype, tf.float16)):<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>".format(<EOL>dtype_util.name(power.dtype)))<EOL><DEDENT>runtime_assertions = []<EOL>if validate_args:<EOL><INDENT>runtime_assertions.append(assert_util.assert_greater(<EOL>power, np.ones([], power.dtype.as_numpy_dtype)))<EOL><DEDENT>with tf.control_dependencies(runtime_assertions):<EOL><INDENT>self._power = tf.identity(power, name="<STR_LIT>")<EOL><DEDENT><DEDENT>self._interpolate_nondiscrete = interpolate_nondiscrete<EOL>self._sample_maximum_iterations = sample_maximum_iterations<EOL>super(Zipf, self).__init__(<EOL>dtype=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._power],<EOL>name=name)<EOL>

Initialize a batch of Zipf distributions.

        Args:
          power: `Float` like `Tensor` representing the power parameter. Must be
            strictly greater than `1`.
          dtype: The `dtype` of `Tensor` returned by `sample`.
            Default value: `tf.int32`.
          interpolate_nondiscrete: Python `bool`. When `False`, `log_prob` returns
            `-inf` (and `prob` returns `0`) for non-integer inputs. When `True`,
            `log_prob` evaluates the continuous function `-power log(k) -
            log(zeta(power))` , which matches the Zipf pmf at integer arguments `k`
            (note that this function is not itself a normalized probability
            log-density).
            Default value: `True`.
          sample_maximum_iterations: Maximum number of iterations of allowable
            iterations in `sample`. When `validate_args=True`, samples which fail to
            reach convergence (subject to this cap) are masked out with
            `self.dtype.min` or `nan` depending on `self.dtype.is_integer`.
            Default value: `100`.
          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: `'Zipf'`.

        Raises:
          TypeError: if `power` is not `float` like.


return self._power<EOL>


parameters = dict(locals())<EOL>with tf.name_scope(name):<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[loc, scale_diag, scale_identity_multiplier], tf.float32)<EOL>scale = distribution_util.make_diag_scale(<EOL>loc=loc,<EOL>scale_diag=scale_diag,<EOL>scale_identity_multiplier=scale_identity_multiplier,<EOL>validate_args=False,<EOL>assert_positive=False,<EOL>dtype=dtype)<EOL><DEDENT><DEDENT>super(VectorLaplaceDiag, self).__init__(<EOL>loc=loc,<EOL>scale=scale,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL>self._parameters = parameters<EOL>

Construct Vector Laplace distribution on `R^k`.

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

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

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

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

        where:

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

        Additional leading dimensions (if any) will index batches.

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

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

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


with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>if observation_index_points is None:<EOL><INDENT>observation_index_points = self._index_points<EOL><DEDENT>observation_index_points = tf.convert_to_tensor(<EOL>value=observation_index_points, dtype=self._dtype,<EOL>name='<STR_LIT>')<EOL>observations = tf.convert_to_tensor(<EOL>value=observations, dtype=self._dtype, name='<STR_LIT>')<EOL>kl_weight = tf.convert_to_tensor(<EOL>value=kl_weight, dtype=self._dtype,<EOL>name='<STR_LIT>')<EOL>kzx = self.kernel.matrix(self._inducing_index_points,<EOL>observation_index_points)<EOL>kzx_linop = tf.linalg.LinearOperatorFullMatrix(kzx)<EOL>loc = (self._mean_fn(observation_index_points) +<EOL>kzx_linop.matvec(self._kzz_inv_varloc, adjoint=True))<EOL>likelihood = independent.Independent(<EOL>normal.Normal(<EOL>loc=loc,<EOL>scale=tf.sqrt(self._observation_noise_variance + self._jitter),<EOL>name='<STR_LIT>'),<EOL>reinterpreted_batch_ndims=<NUM_LIT:1>)<EOL>obs_ll = likelihood.log_prob(observations)<EOL>chol_kzz_linop = tf.linalg.LinearOperatorLowerTriangular(self._chol_kzz)<EOL>chol_kzz_inv_kzx = chol_kzz_linop.solve(kzx)<EOL>kzz_inv_kzx = chol_kzz_linop.solve(chol_kzz_inv_kzx, adjoint=True)<EOL>kxx_diag = tf.linalg.diag_part(<EOL>self.kernel.matrix(<EOL>observation_index_points, observation_index_points))<EOL>ktilde_trace_term = (<EOL>tf.reduce_sum(input_tensor=kxx_diag, axis=-<NUM_LIT:1>) -<EOL>tf.reduce_sum(input_tensor=chol_kzz_inv_kzx ** <NUM_LIT:2>, axis=[-<NUM_LIT:2>, -<NUM_LIT:1>]))<EOL>other_trace_term = tf.reduce_sum(<EOL>input_tensor=(<EOL>self._variational_inducing_observations_posterior.scale.matmul(<EOL>kzz_inv_kzx) ** <NUM_LIT:2>),<EOL>axis=[-<NUM_LIT:2>, -<NUM_LIT:1>])<EOL>trace_term = (<NUM_LIT> * (ktilde_trace_term + other_trace_term) /<EOL>self._observation_noise_variance)<EOL>inducing_prior = gaussian_process.GaussianProcess(<EOL>kernel=self._kernel,<EOL>mean_fn=self._mean_fn,<EOL>index_points=self._inducing_index_points,<EOL>observation_noise_variance=self._observation_noise_variance)<EOL>kl_term = kl_weight * kullback_leibler.kl_divergence(<EOL>self._variational_inducing_observations_posterior,<EOL>inducing_prior)<EOL>lower_bound = (obs_ll - trace_term - kl_term)<EOL>return -tf.reduce_mean(input_tensor=lower_bound)<EOL><DEDENT>

Variational loss for the VGP.

        Given `observations` and `observation_index_points`, compute the
        negative variational lower bound as specified in [Hensman, 2013][1].

        Args:
          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`, `observation_index_points`, etc.).
          observation_index_points: `float` `Tensor` representing finite (batch of)
            vector(s) of points where observations are defined. Shape has the
            form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
            dimensions and must equal `kernel.feature_ndims` and `e1` is the number
            (size) of index points in each batch (we denote it `e1` to distinguish
            it from the numer of inducing index points, denoted `e2` below). If
            set to `None` uses `index_points` as the origin for observations.
            Default value: None.
          kl_weight: Amount by which to scale the KL divergence loss between prior
            and posterior.
            Default value: 1.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: "GaussianProcess".
        Returns:
          loss: Scalar tensor representing the negative variational lower bound.
            Can be directly used in a `tf.Optimizer`.
        Raises:
          ValueError: if `mean_fn` is not `None` and is not callable.

        #### References

        [1]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
             https://arxiv.org/abs/1309.6835


return self.bijector.scale<EOL>

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


return self._total_concentration<EOL>

Sum of last dim of concentration parameter.