with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>logx = tf.convert_to_tensor(value=logx, name="<STR_LIT>")<EOL>if w is None:<EOL><INDENT>lswe = tf.reduce_logsumexp(<EOL>input_tensor=logx, axis=axis, keepdims=keep_dims)<EOL>if return_sign:<EOL><INDENT>sgn = tf.ones_like(lswe)<EOL>return lswe, sgn<EOL><DEDENT>return lswe<EOL><DEDENT>w = tf.convert_to_tensor(value=w, dtype=logx.dtype, name="<STR_LIT:w>")<EOL>log_absw_x = logx + tf.math.log(tf.abs(w))<EOL>max_log_absw_x = tf.reduce_max(<EOL>input_tensor=log_absw_x, axis=axis, keepdims=True)<EOL>max_log_absw_x = tf.where(<EOL>tf.math.is_inf(max_log_absw_x), tf.zeros_like(max_log_absw_x),<EOL>max_log_absw_x)<EOL>wx_over_max_absw_x = (tf.sign(w) * tf.exp(log_absw_x - max_log_absw_x))<EOL>sum_wx_over_max_absw_x = tf.reduce_sum(<EOL>input_tensor=wx_over_max_absw_x, axis=axis, keepdims=keep_dims)<EOL>if not keep_dims:<EOL><INDENT>max_log_absw_x = tf.squeeze(max_log_absw_x, axis)<EOL><DEDENT>sgn = tf.sign(sum_wx_over_max_absw_x)<EOL>lswe = max_log_absw_x + tf.math.log(sgn * sum_wx_over_max_absw_x)<EOL>if return_sign:<EOL><INDENT>return lswe, sgn<EOL><DEDENT>return lswe<EOL><DEDENT>

Computes `log(abs(sum(weight * exp(elements across tensor dimensions))))`.

    If all weights `w` are known to be positive, it is more efficient to directly
    use `reduce_logsumexp`, i.e., `tf.reduce_logsumexp(logx + tf.log(w))` is more
    efficient than `du.reduce_weighted_logsumexp(logx, w)`.

    Reduces `input_tensor` along the dimensions given in `axis`.
    Unless `keep_dims` is true, the rank of the tensor is reduced by 1 for each
    entry in `axis`. If `keep_dims` is true, the reduced dimensions
    are retained with length 1.

    If `axis` has no entries, all dimensions are reduced, and a
    tensor with a single element is returned.

    This function is more numerically stable than log(sum(w * exp(input))). It
    avoids overflows caused by taking the exp of large inputs and underflows
    caused by taking the log of small inputs.

    For example:

    ```python
    x = tf.constant([[0., 0, 0],
                     [0, 0, 0]])

    w = tf.constant([[-1., 1, 1],
                     [1, 1, 1]])

    du.reduce_weighted_logsumexp(x, w)
    # ==> log(-1*1 + 1*1 + 1*1 + 1*1 + 1*1 + 1*1) = log(4)

    du.reduce_weighted_logsumexp(x, w, axis=0)
    # ==> [log(-1+1), log(1+1), log(1+1)]

    du.reduce_weighted_logsumexp(x, w, axis=1)
    # ==> [log(-1+1+1), log(1+1+1)]

    du.reduce_weighted_logsumexp(x, w, axis=1, keep_dims=True)
    # ==> [[log(-1+1+1)], [log(1+1+1)]]

    du.reduce_weighted_logsumexp(x, w, axis=[0, 1])
    # ==> log(-1+5)
    ```

    Args:
      logx: The tensor to reduce. Should have numeric type.
      w: The weight tensor. Should have numeric type identical to `logx`.
      axis: The dimensions to reduce. If `None` (the default), reduces all
        dimensions. Must be in the range `[-rank(input_tensor),
        rank(input_tensor))`.
      keep_dims: If true, retains reduced dimensions with length 1.
      return_sign: If `True`, returns the sign of the result.
      name: A name for the operation (optional).

    Returns:
      lswe: The `log(abs(sum(weight * exp(x))))` reduced tensor.
      sign: (Optional) The sign of `sum(weight * exp(x))`.


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


<EOL>s = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(x.shape, np.abs(axis))[axis])<EOL>if s is not None:<EOL><INDENT>return s<EOL><DEDENT>return tf.shape(input=x)[axis]<EOL>

Returns the size of a specific dimension.


with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>if quadrature_grid_and_probs is None:<EOL><INDENT>grid, probs = np.polynomial.hermite.hermgauss(deg=<NUM_LIT:8>)<EOL>grid = grid.astype(dtype_util.as_numpy_dtype(dtype))<EOL>probs = probs.astype(dtype_util.as_numpy_dtype(dtype))<EOL>probs /= np.linalg.norm(probs, ord=<NUM_LIT:1>, keepdims=True)<EOL>grid = tf.convert_to_tensor(value=grid, name="<STR_LIT>", dtype=dtype)<EOL>probs = tf.convert_to_tensor(value=probs, name="<STR_LIT>", dtype=dtype)<EOL>return grid, probs<EOL><DEDENT>grid, probs = tuple(quadrature_grid_and_probs)<EOL>grid = tf.convert_to_tensor(value=grid, name="<STR_LIT>", dtype=dtype)<EOL>probs = tf.convert_to_tensor(<EOL>value=probs, name="<STR_LIT>", dtype=dtype)<EOL>probs /= tf.norm(tensor=probs, ord=<NUM_LIT:1>, axis=-<NUM_LIT:1>, keepdims=True, name="<STR_LIT>")<EOL>def _static_event_size(x):<EOL><INDENT>"""<STR_LIT>"""<EOL>return tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(x.shape, <NUM_LIT:1>)[-<NUM_LIT:1>])<EOL><DEDENT>m, n = _static_event_size(probs), _static_event_size(grid)<EOL>if m is not None and n is not None:<EOL><INDENT>if m != n:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(m, n))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions = [<EOL>assert_util.assert_equal(<EOL>dimension_size(probs, axis=-<NUM_LIT:1>),<EOL>dimension_size(grid, axis=-<NUM_LIT:1>),<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>")),<EOL>]<EOL>with tf.control_dependencies(assertions):<EOL><INDENT>grid = tf.identity(grid)<EOL>probs = tf.identity(probs)<EOL><DEDENT><DEDENT>return grid, probs<EOL><DEDENT>

Validates quadrature grid, probs or computes them as necessary.

    Args:
      quadrature_grid_and_probs: Python pair of `float`-like `Tensor`s
        representing the sample points and the corresponding (possibly
        normalized) weight.  When `None`, defaults to:
          `np.polynomial.hermite.hermgauss(deg=8)`.
      dtype: The expected `dtype` of `grid` and `probs`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      name: Python `str` name prefixed to Ops created by this class.

    Returns:
       quadrature_grid_and_probs: Python pair of `float`-like `Tensor`s
        representing the sample points and the corresponding (possibly
        normalized) weight.

    Raises:
      ValueError: if `quadrature_grid_and_probs is not None` and
        `len(quadrature_grid_and_probs[0]) != len(quadrature_grid_and_probs[1])`


def pad(x, axis, front=False, back=False, value=<NUM_LIT:0>, count=<NUM_LIT:1>, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>value = tf.convert_to_tensor(value=value, dtype=x.dtype, name="<STR_LIT:value>")<EOL>count = tf.convert_to_tensor(value=count, name="<STR_LIT:count>")<EOL>if not dtype_util.is_integer(count.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>".format(<EOL>dtype_util.name(count.dtype)))<EOL><DEDENT>if not front and not back:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>ndims = (<EOL>tensorshape_util.rank(x.shape)<EOL>if tensorshape_util.rank(x.shape) is not None else tf.rank(<EOL>x, name="<STR_LIT>"))<EOL>axis = tf.convert_to_tensor(value=axis, name="<STR_LIT>")<EOL>axis_ = tf.get_static_value(axis)<EOL>if axis_ is not None:<EOL><INDENT>axis = axis_<EOL>if axis < <NUM_LIT:0>:<EOL><INDENT>axis = ndims + axis<EOL><DEDENT>count_ = tf.get_static_value(count)<EOL>if axis_ >= <NUM_LIT:0> or tensorshape_util.rank(x.shape) is not None:<EOL><INDENT>head = x.shape[:axis]<EOL>mid_dim_value = tf.compat.dimension_value(x.shape[axis])<EOL>if count_ is None or mid_dim_value is None:<EOL><INDENT>middle = tf.TensorShape(None)<EOL><DEDENT>else:<EOL><INDENT>middle = tf.TensorShape(mid_dim_value + count_ * (front + back))<EOL><DEDENT>tail = x.shape[axis + <NUM_LIT:1>:]<EOL>final_shape = head.concatenate(middle.concatenate(tail))<EOL><DEDENT>else:<EOL><INDENT>final_shape = None<EOL><DEDENT><DEDENT>else:<EOL><INDENT>axis = tf.where(axis < <NUM_LIT:0>, ndims + axis, axis)<EOL>final_shape = None<EOL><DEDENT>x = tf.pad(<EOL>tensor=x,<EOL>paddings=tf.one_hot(<EOL>indices=tf.stack([axis if front else -<NUM_LIT:1>, axis if back else -<NUM_LIT:1>]),<EOL>depth=ndims,<EOL>axis=<NUM_LIT:0>,<EOL>on_value=count,<EOL>dtype=tf.int32),<EOL>constant_values=value)<EOL>if final_shape is not None:<EOL><INDENT>tensorshape_util.set_shape(x, final_shape)<EOL><DEDENT>return x<EOL><DEDENT>

Pads `value` to the front and/or back of a `Tensor` dim, `count` times.

    Args:
      x: `Tensor` input.
      axis: Scalar `int`-like `Tensor` representing the single dimension to pad.
        (Negative indexing is supported.)
      front: Python `bool`; if `True` the beginning of the `axis` dimension is
        padded with `value`, `count` times. If `False` no front padding is made.
      back: Python `bool`; if `True` the end of the `axis` dimension is padded
        with `value`, `count` times. If `False` no end padding is made.
      value: Scalar `int`-like `Tensor` representing the actual value added to the
        front and/or back of the `axis` dimension of `x`.
      count: Scalar `int`-like `Tensor` representing number of elements added to
        the front and/or back of the `axis` dimension of `x`. E.g., if `front =
        back = True` then `2 * count` elements are added.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      pad: The padded version of input `x`.

    Raises:
      ValueError: if both `front` and `back` are `False`.
      TypeError: if `count` is not `int`-like.


<EOL>arg_names, variable_arg_name, keyword_arg_name, local_vars = (<EOL>tf_inspect._inspect.getargvalues(  <EOL>tf_inspect._inspect.stack()[<NUM_LIT:1>][<NUM_LIT:0>]))  <EOL>local_vars.pop(variable_arg_name, {})<EOL>keyword_args = local_vars.pop(keyword_arg_name, {})<EOL>final_args = {}<EOL>for arg_name in arg_names:<EOL><INDENT>final_args[arg_name] = local_vars.pop(arg_name)<EOL><DEDENT>final_args.update(keyword_args)<EOL>return final_args<EOL>

Returns parent frame arguments.

    When called inside a function, returns a dictionary with the caller's function
    arguments. These are positional arguments and keyword arguments (**kwargs),
    while variable arguments (*varargs) are excluded.

    When called at global scope, this will return an empty dictionary, since there
    are no arguments.

    WARNING: If caller function argument names are overloaded before invoking
    this method, then values will reflect the overloaded value. For this reason,
    we recommend calling `parent_frame_arguments` at the beginning of the
    function.


def expand_to_vector(x, tensor_name=None, op_name=None, validate_args=False):

with tf.name_scope(op_name or "<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>ndims = tensorshape_util.rank(x.shape)<EOL>if ndims is None:<EOL><INDENT>if validate_args:<EOL><INDENT>x = with_dependencies([<EOL>assert_util.assert_rank_at_most(<EOL>x, <NUM_LIT:1>, message="<STR_LIT>")<EOL>], x)<EOL><DEDENT>ndims = tf.rank(x)<EOL>expanded_shape = pick_vector(<EOL>tf.equal(ndims, <NUM_LIT:0>), np.array([<NUM_LIT:1>], dtype=np.int32), tf.shape(input=x))<EOL>return tf.reshape(x, expanded_shape)<EOL><DEDENT>elif ndims == <NUM_LIT:0>:<EOL><INDENT>x_const = tf.get_static_value(x)<EOL>if x_const is not None:<EOL><INDENT>return tf.convert_to_tensor(<EOL>value=dtype_util.as_numpy_dtype(x.dtype)([x_const]),<EOL>name=tensor_name)<EOL><DEDENT>else:<EOL><INDENT>return tf.reshape(x, [<NUM_LIT:1>])<EOL><DEDENT><DEDENT>elif ndims != <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>return x<EOL><DEDENT>

Transform a 0-D or 1-D `Tensor` to be 1-D.

    For user convenience, many parts of the TensorFlow Probability API accept
    inputs of rank 0 or 1 -- i.e., allowing an `event_shape` of `[5]` to be passed
    to the API as either `5` or `[5]`.  This function can be used to transform
    such an argument to always be 1-D.

    NOTE: Python or NumPy values will be converted to `Tensor`s with standard type
    inference/conversion.  In particular, an empty list or tuple will become an
    empty `Tensor` with dtype `float32`.  Callers should convert values to
    `Tensor`s before calling this function if different behavior is desired
    (e.g. converting empty lists / other values to `Tensor`s with dtype `int32`).

    Args:
      x: A 0-D or 1-D `Tensor`.
      tensor_name: Python `str` name for `Tensor`s created by this function.
      op_name: Python `str` name for `Op`s created by this function.
      validate_args: Python `bool, default `False`.  When `True`, arguments may be
        checked for validity at execution time, possibly degrading runtime
        performance.  When `False`, invalid inputs may silently render incorrect
          outputs.
    Returns:
      vector: a 1-D `Tensor`.


def with_dependencies(dependencies, output_tensor, name=None):

if tf.executing_eagerly():<EOL><INDENT>return output_tensor<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>") as name:<EOL><INDENT>with tf.control_dependencies(d for d in dependencies if d is not None):<EOL><INDENT>output_tensor = tf.convert_to_tensor(value=output_tensor)<EOL>if isinstance(output_tensor, tf.Tensor):<EOL><INDENT>return tf.identity(output_tensor, name=name)<EOL><DEDENT>else:<EOL><INDENT>return tf.IndexedSlices(<EOL>tf.identity(output_tensor.values, name=name),<EOL>output_tensor.indices,<EOL>output_tensor.dense_shape)<EOL><DEDENT><DEDENT><DEDENT>

Produces the content of `output_tensor` only after `dependencies`.

    In some cases, a user may want the output of an operation to be consumed
    externally only after some other dependencies have run first. This function
    returns `output_tensor`, but only after all operations in `dependencies` have
    run. Note that this means that there is no guarantee that `output_tensor` will
    be evaluated after any `dependencies` have run.

    See also `tf.tuple` and `tf.group`.

    Args:
      dependencies: Iterable of operations to run before this op finishes.
      output_tensor: A `Tensor` or `IndexedSlices` that will be returned.
      name: (Optional) A name for this operation.

    Returns:
      output_with_deps: Same as `output_tensor` but with embedded dependencies.

    Raises:
      TypeError: if `output_tensor` is not a `Tensor` or `IndexedSlices`.


return (not tf_inspect.isclass(d) and<EOL>hasattr(d, "<STR_LIT>") and<EOL>hasattr(d, "<STR_LIT>"))<EOL>

Standardizes our definition of being a `tfd.Distribution`.


def __init__(self, additional_note="<STR_LIT>", kwargs_dict=None):

self._additional_note = additional_note<EOL>if kwargs_dict:<EOL><INDENT>bullets = []<EOL>for key in sorted(kwargs_dict.keys()):<EOL><INDENT>value = kwargs_dict[key]<EOL>if any(x.isspace() for x in key):<EOL><INDENT>raise ValueError("<STR_LIT>" % key)<EOL><DEDENT>value = value.lstrip()<EOL>if "<STR_LIT:\n>" in value:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>" % key)<EOL><DEDENT>bullets.append("<STR_LIT>" % (key, value))<EOL><DEDENT>self._additional_note += ("<STR_LIT>" + "<STR_LIT:\n>".join(bullets))<EOL><DEDENT>

Initializes the AppendDocstring object.

        Args:
          additional_note: Python string added as additional docstring to public
            version of function.
          kwargs_dict: Python string/string dictionary representing specific kwargs
            expanded from the **kwargs input.

        Raises:
          ValueError: if kwargs_dict.key contains whitespace.
          ValueError: if kwargs_dict.value contains newlines.


def _compute_rank_and_fullsize_reqd(draw, target_shape, current_shape, is_last):

target_rank = target_shape.ndims<EOL>if is_last:<EOL><INDENT>full_rank_current = tf.broadcast_static_shape(<EOL>current_shape, tf.TensorShape([<NUM_LIT:1>] * target_rank))<EOL>axis_is_mismatched = [<EOL>full_rank_current[i] != target_shape[i] for i in range(target_rank)<EOL>]<EOL>min_rank = target_rank<EOL>if current_shape.ndims == target_rank:<EOL><INDENT>min_rank -= (axis_is_mismatched + [True]).index(True)<EOL><DEDENT>next_rank = draw(<EOL>hps.integers(min_value=min_rank, max_value=target_rank))<EOL>force_fullsize_dim = axis_is_mismatched[target_rank - next_rank:]<EOL><DEDENT>else:<EOL><INDENT>next_rank = draw(hps.integers(min_value=<NUM_LIT:0>, max_value=target_rank))<EOL>force_fullsize_dim = [False] * next_rank<EOL><DEDENT>return next_rank, force_fullsize_dim<EOL>

Returns a param rank and a list of bools for full-size-required by axis.

    Args:
      draw: Hypothesis data sampler.
      target_shape: `tf.TensorShape`, the target broadcasted shape.
      current_shape: `tf.TensorShape`, the broadcasted shape of the shapes
        selected thus far. This is ignored for non-last shapes.
      is_last: bool indicator of whether this is the last shape (in which case, we
        must achieve the target shape).

    Returns:
      next_rank: Sampled rank for the next shape.
      force_fullsize_dim: `next_rank`-sized list of bool indicating whether the
        corresponding axis of the shape must be full-sized (True) or is allowed to
        be 1 (i.e., broadcast) (False).


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


with tf.compat.v2.name_scope(name or '<STR_LIT>'):<EOL><INDENT>return x[..., :, tf.newaxis] * x[..., tf.newaxis, :]<EOL><DEDENT>

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


def histogram(self, x, value_range=None, nbins=None, name=None):

with tf.compat.v2.name_scope(name or '<STR_LIT>'):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if value_range is None:<EOL><INDENT>value_range = [<EOL>tf.reduce_min(input_tensor=x), <NUM_LIT:1> + tf.reduce_max(input_tensor=x)<EOL>]<EOL><DEDENT>value_range = tf.convert_to_tensor(value=value_range, name='<STR_LIT>')<EOL>lo = value_range[<NUM_LIT:0>]<EOL>hi = value_range[<NUM_LIT:1>]<EOL>if nbins is None:<EOL><INDENT>nbins = tf.cast(hi - lo, dtype=tf.int32)<EOL><DEDENT>delta = (hi - lo) / tf.cast(<EOL>nbins, dtype=dtype_util.base_dtype(value_range.dtype))<EOL>edges = tf.range(<EOL>start=lo, limit=hi, delta=delta, dtype=dtype_util.base_dtype(x.dtype))<EOL>counts = tf.histogram_fixed_width(x, value_range=value_range, nbins=nbins)<EOL>return counts, edges<EOL><DEDENT>

Return histogram of values.

        Given the tensor `values`, this operation returns a rank 1 histogram
        counting the number of entries in `values` that fell into every bin. The
        bins are equal width and determined by the arguments `value_range` and
        `nbins`.

        Args:
          x: 1D numeric `Tensor` of items to count.
          value_range:  Shape [2] `Tensor`. `new_values <= value_range[0]` will be
            mapped to `hist[0]`, `values >= value_range[1]` will be mapped to
            `hist[-1]`. Must be same dtype as `x`.
          nbins:  Scalar `int32 Tensor`.  Number of histogram bins.
          name: Python `str` name prefixed to Ops created by this class.

        Returns:
          counts: 1D `Tensor` of counts, i.e.,
            `counts[i] = sum{ edges[i-1] <= values[j] < edges[i] : j }`.
          edges: 1D `Tensor` characterizing intervals used for counting.


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

Returns a `list` of integers or `None` for each dimension.

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

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

    Returns:
      shape_as_list: list of `int` or `None` values representing each dimensions
        size if known.

    Raises:
      ValueError: If `x` has unknown rank.


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

Raises an exception if `x` is not compatible with the given `rank`.

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

    Args:
      x: object representing a shape; convertible to `tf.TensorShape`.
      rank: an `int` representing the minimum required rank of `x`.

    Returns:
      None

    Raises:
      ValueError: If `x` does not represent a shape with the given `rank`.


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.


return type(x)(tf.TensorShape(x).concatenate(other))<EOL>

Returns the concatenation of the dimension in `x` and `other`.

    *Note:* If either `x` or `other` is completely unknown, concatenation will
    discard information about the other shape. In future, we might support
    concatenation that preserves this information for use with slicing.

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

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

    Returns:
      new_shape: an object like `x` whose elements are the concatenation of the
        dimensions in `x` and `other`.


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.


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.


return tf.TensorShape(x).is_compatible_with(other)<EOL>

Returns `True` iff `x` is compatible with `other`.

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

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

    Returns:
      is_compatible: `bool` indicating of the shapes are compatible.


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

Returns True iff `x` is fully defined in every dimension.

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

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

    Returns:
      is_fully_defined: `bool` indicating that the shape is fully known.


return type(x)(tf.TensorShape(x).merge_with(other))<EOL>

Returns a shape combining the information in `x` and `other`.

    The dimensions in `x` and `other` are merged elementwise, according to the
    rules defined for `tf.Dimension.merge_with()`.

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

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

    Returns:
      merged_shape: shape having `type(x)` containing the combined information of
        `x` and `other`.

    Raises:
      ValueError: If `x` and `other` are not compatible.


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

Returns the total number of elements, or `None` for incomplete shapes.

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

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

    Returns:
      num_elements: `int` representing the total number of elements implied by
        shape `x`.


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.


if hasattr(tensor, '<STR_LIT>'):<EOL><INDENT>tensor.set_shape(shape)<EOL><DEDENT>

Updates the shape of this tensor.

    This method can be called multiple times, and will merge the given
    `shape` with the current shape of this tensor. It can be used to
    provide additional information about the shape of this tensor that
    cannot be inferred from the graph alone. For example, this can be used
    to provide additional information about the shapes of images:

    ```python
    _, image_data = tf.TFRecordReader(...).read(...)
    image = tf.image.decode_png(image_data, channels=3)

    # The height and width dimensions of `image` are data dependent, and
    # cannot be computed without executing the op.
    print(image.shape)
    ==> TensorShape([Dimension(None), Dimension(None), Dimension(3)])

    # We know that each image in this dataset is 28 x 28 pixels.
    image.set_shape([28, 28, 3])
    print(image.shape)
    ==> TensorShape([Dimension(28), Dimension(28), Dimension(3)])
    ```

    NOTE: This shape is not enforced at runtime. Setting incorrect shapes can
    result in inconsistencies between the statically-known graph and the runtime
    value of tensors. For runtime validation of the shape, use `tf.ensure_shape`
    instead.

    Args:
      tensor: `Tensor` which will have its static shape set.
      shape: A `TensorShape` representing the shape of this tensor, a
      `TensorShapeProto`, a list, a tuple, or None.

    Raises:
      ValueError: If `shape` is not compatible with the current shape of
        this tensor.


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

Returns a shape based on `x` with at least the given `rank`.

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

    Args:
      x: object representing a shape; convertible to `tf.TensorShape`.
      rank: An `int` representing the minimum rank of `x` or else an assertion is
        raised.

    Returns:
      shape: a shape having `type(x)` but guaranteed to have at least the given
        rank (or else an assertion was raised).

    Raises:
      ValueError: If `x` does not represent a shape with at least the given
        `rank`.


orig_shape = d.shape<EOL>d = np.reshape(d, (int(np.prod(d.shape[:-<NUM_LIT:1>])), d.shape[-<NUM_LIT:1>]))<EOL>diag_list = []<EOL>for i in range(d.shape[<NUM_LIT:0>]):<EOL><INDENT>diag_list.append(np.diag(d[i, ...]))<EOL><DEDENT>return np.reshape(diag_list, orig_shape + (d.shape[-<NUM_LIT:1>],))<EOL>


def _fill_triangular(self, x, upper=False):

x = np.asarray(x)<EOL>m = np.int32(x.shape[-<NUM_LIT:1>])<EOL>n = np.sqrt(<NUM_LIT> + <NUM_LIT> * m) - <NUM_LIT:0.5><EOL>if n != np.floor(n):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>n = np.int32(n)<EOL>x_tail = x[..., (m - (n * n - m)):]<EOL>y = np.concatenate(<EOL>[x, x_tail[..., ::-<NUM_LIT:1>]] if upper else [x_tail, x[..., ::-<NUM_LIT:1>]],<EOL>axis=-<NUM_LIT:1>)<EOL>y = y.reshape(np.concatenate([<EOL>np.int32(x.shape[:-<NUM_LIT:1>]),<EOL>np.int32([n, n]),<EOL>], axis=<NUM_LIT:0>))<EOL>return np.triu(y) if upper else np.tril(y)<EOL>

Numpy implementation of `fill_triangular`.


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

Assert that all entries in a `Tensor` are finite.

        Args:
          a: A `Tensor` whose entries are checked for finiteness.


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.


dtype = tf.as_dtype(dtype)<EOL>if hasattr(dtype, '<STR_LIT>'):<EOL><INDENT>return dtype.as_numpy_dtype<EOL><DEDENT>return dtype<EOL>

Returns a `np.dtype` based on this `dtype`.


dtype = tf.as_dtype(dtype)<EOL>if hasattr(dtype, '<STR_LIT>'):<EOL><INDENT>return dtype.base_dtype<EOL><DEDENT>return dtype<EOL>

Returns a non-reference `dtype` based on this `dtype`.


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

Returns `True` if base dtypes are identical.


dtype = None<EOL>for a in tf.nest.flatten(args_list):<EOL><INDENT>if hasattr(a, '<STR_LIT>'):<EOL><INDENT>dt = as_numpy_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>return preferred_dtype if dtype is None else tf.as_dtype(dtype)<EOL>

Returns explict dtype from `args_list` if there is one.


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.


dtype = tf.as_dtype(dtype)<EOL>if hasattr(dtype, '<STR_LIT>'):<EOL><INDENT>return dtype.is_complex<EOL><DEDENT>return np.issubdtype(np.dtype(dtype), np.complex)<EOL>

Returns whether this is a complex floating point type.


dtype = tf.as_dtype(dtype)<EOL>if hasattr(dtype, '<STR_LIT>'):<EOL><INDENT>return dtype.is_floating<EOL><DEDENT>return np.issubdtype(np.dtype(dtype), np.float)<EOL>

Returns whether this is a (non-quantized, real) floating point type.


dtype = tf.as_dtype(dtype)<EOL>if hasattr(dtype, '<STR_LIT>'):<EOL><INDENT>return dtype.is_integer<EOL><DEDENT>return np.issubdtype(np.dtype(dtype), np.integer)<EOL>

Returns whether this is a (non-quantized) integer type.


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.min<EOL><DEDENT>use_finfo = is_floating(dtype) or is_complex(dtype)<EOL>return np.finfo(dtype).min if use_finfo else np.iinfo(dtype).min<EOL>

Returns the minimum representable value in this data type.


dtype = tf.as_dtype(dtype)<EOL>if hasattr(dtype, '<STR_LIT:name>'):<EOL><INDENT>return dtype.name<EOL><DEDENT>if hasattr(dtype, '<STR_LIT>'):<EOL><INDENT>return dtype.__name__<EOL><DEDENT>return str(dtype)<EOL>

Returns the string name for this `dtype`.


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

Returns the number of bytes to represent this `dtype`.


original_expected_type = expected_type<EOL>mismatch = False<EOL>for item in items:<EOL><INDENT>if item is not None:<EOL><INDENT>item_type = base_dtype(item.dtype)<EOL>if not expected_type:<EOL><INDENT>expected_type = item_type<EOL><DEDENT>elif expected_type != item_type:<EOL><INDENT>mismatch = True<EOL>break<EOL><DEDENT><DEDENT><DEDENT>if mismatch:<EOL><INDENT>expected_type = original_expected_type<EOL>original_item_str = None<EOL>get_name = lambda x: x.name if hasattr(x, '<STR_LIT:name>') else str(x)<EOL>for item in items:<EOL><INDENT>if item is not None:<EOL><INDENT>item_type = base_dtype(item.dtype)<EOL>if not expected_type:<EOL><INDENT>expected_type = item_type<EOL>original_item_str = get_name(item)<EOL><DEDENT>elif expected_type != item_type:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(<EOL>get_name(item),<EOL>item_type,<EOL>expected_type,<EOL>(('<STR_LIT>'.format(original_item_str))<EOL>if original_item_str else '<STR_LIT>')))<EOL><DEDENT><DEDENT><DEDENT>return expected_type  <EOL><DEDENT>else:<EOL><INDENT>return expected_type<EOL><DEDENT>

r"""Asserts all items are of the same base type.

    Args:
      items: List of graph items (e.g., `Variable`, `Tensor`, `SparseTensor`,
          `Operation`, or `IndexedSlices`). Can include `None` elements, which
          will be ignored.
      expected_type: Expected type. If not specified, assert all items are
          of the same base type.

    Returns:
      Validated type, or none if neither expected_type nor items provided.

    Raises:
      ValueError: If any types do not match.


if tensors:<EOL><INDENT>dtype = _assert_same_base_type(tensors, dtype)<EOL><DEDENT>if not dtype:<EOL><INDENT>dtype = tf.float32<EOL><DEDENT>elif not is_floating(dtype):<EOL><INDENT>raise ValueError('<STR_LIT>'.format(dtype))<EOL><DEDENT>return dtype<EOL>

Validate and return float type based on `tensors` and `dtype`.

    For ops such as matrix multiplication, inputs and weights must be of the
    same float type. This function validates that all `tensors` are the same type,
    validates that type is `dtype` (if supplied), and returns the type. Type must
    be a floating point type. If neither `tensors` nor `dtype` is supplied,
    the function will return `dtypes.float32`.

    Args:
      tensors: Tensors of input values. Can include `None` elements, which will
        be ignored.
      dtype: Expected type.

    Returns:
      Validated type.

    Raises:
      ValueError: if neither `tensors` nor `dtype` is supplied, or result is not
        float, or the common type of the inputs is not a floating point 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(<EOL>"<STR_LIT>"<EOL>% x.dtype)<EOL><DEDENT>return _ndtr(x)<EOL><DEDENT>

Normal distribution function.

    Returns the area under the Gaussian probability density function, integrated
    from minus infinity to x:

    ```
                      1       / x
       ndtr(x)  = ----------  |    exp(-0.5 t**2) dt
                  sqrt(2 pi)  /-inf

                = 0.5 (1 + erf(x / sqrt(2)))
                = 0.5 erfc(x / sqrt(2))
    ```

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

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

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


half_sqrt_2 = tf.constant(<EOL><NUM_LIT:0.5> * np.sqrt(<NUM_LIT>), dtype=x.dtype, name="<STR_LIT>")<EOL>w = x * half_sqrt_2<EOL>z = tf.abs(w)<EOL>y = tf.where(<EOL>tf.less(z, half_sqrt_2), <NUM_LIT:1.> + tf.math.erf(w),<EOL>tf.where(tf.greater(w, <NUM_LIT:0.>), <NUM_LIT> - tf.math.erfc(z), tf.math.erfc(z)))<EOL>return <NUM_LIT:0.5> * y<EOL>


with tf.name_scope(name):<EOL><INDENT>p = tf.convert_to_tensor(value=p, name="<STR_LIT:p>")<EOL>if dtype_util.as_numpy_dtype(p.dtype) not in [np.float32, np.float64]:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>"<EOL>% p.dtype)<EOL><DEDENT>return _ndtri(p)<EOL><DEDENT>

The inverse of the CDF of the Normal distribution function.

    Returns x such that the area under the pdf from minus infinity to x is equal
    to p.

    A piece-wise rational approximation is done for the function.
    This is a port of the implementation in netlib.

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

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

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


<EOL>p0 = list(reversed([-<NUM_LIT>,<EOL><NUM_LIT>,<EOL>-<NUM_LIT>,<EOL><NUM_LIT>,<EOL>-<NUM_LIT>]))<EOL>q0 = list(reversed([<NUM_LIT:1.0>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL>-<NUM_LIT>,<EOL><NUM_LIT>,<EOL>-<NUM_LIT>,<EOL><NUM_LIT>,<EOL>-<NUM_LIT>]))<EOL>p1 = list(reversed([<NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL>-<NUM_LIT>,<EOL>-<NUM_LIT>,<EOL>-<NUM_LIT>]))<EOL>q1 = list(reversed([<NUM_LIT:1.0>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL>-<NUM_LIT>,<EOL>-<NUM_LIT>,<EOL>-<NUM_LIT>]))<EOL>p2 = list(reversed([<NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>]))<EOL>q2 = list(reversed([<NUM_LIT:1.0>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>,<EOL><NUM_LIT>]))<EOL>def _create_polynomial(var, coeffs):<EOL><INDENT>"""<STR_LIT>"""<EOL>coeffs = np.array(coeffs, dtype_util.as_numpy_dtype(var.dtype))<EOL>if not coeffs.size:<EOL><INDENT>return tf.zeros_like(var)<EOL><DEDENT>return coeffs[<NUM_LIT:0>] + _create_polynomial(var, coeffs[<NUM_LIT:1>:]) * var<EOL><DEDENT>maybe_complement_p = tf.where(p > -np.expm1(-<NUM_LIT>), <NUM_LIT:1.> - p, p)<EOL>sanitized_mcp = tf.where(<EOL>maybe_complement_p <= <NUM_LIT:0.>,<EOL>tf.fill(tf.shape(input=p),<EOL>dtype_util.as_numpy_dtype(p.dtype)(<NUM_LIT:0.5>)), maybe_complement_p)<EOL>w = sanitized_mcp - <NUM_LIT:0.5><EOL>ww = w ** <NUM_LIT:2><EOL>x_for_big_p = w + w * ww * (_create_polynomial(ww, p0)<EOL>/ _create_polynomial(ww, q0))<EOL>x_for_big_p *= -np.sqrt(<NUM_LIT> * np.pi)<EOL>z = tf.sqrt(-<NUM_LIT> * tf.math.log(sanitized_mcp))<EOL>first_term = z - tf.math.log(z) / z<EOL>second_term_small_p = (<EOL>_create_polynomial(<NUM_LIT:1.> / z, p2) /<EOL>_create_polynomial(<NUM_LIT:1.> / z, q2) / z)<EOL>second_term_otherwise = (<EOL>_create_polynomial(<NUM_LIT:1.> / z, p1) /<EOL>_create_polynomial(<NUM_LIT:1.> / z, q1) / z)<EOL>x_for_small_p = first_term - second_term_small_p<EOL>x_otherwise = first_term - second_term_otherwise<EOL>x = tf.where(sanitized_mcp > np.exp(-<NUM_LIT>), x_for_big_p,<EOL>tf.where(z >= <NUM_LIT>, x_for_small_p, x_otherwise))<EOL>x = tf.where(p > <NUM_LIT:1.> - np.exp(-<NUM_LIT>), x, -x)<EOL>infinity_scalar = tf.constant(np.inf, dtype=p.dtype)<EOL>infinity = tf.fill(tf.shape(input=p), infinity_scalar)<EOL>x_nan_replaced = tf.where(<EOL>p <= <NUM_LIT:0.0>, -infinity, tf.where(p >= <NUM_LIT:1.0>, infinity, x))<EOL>return x_nan_replaced<EOL>


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


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


npdt = dtype_util.as_numpy_dtype(x.dtype)<EOL>if series_order <= <NUM_LIT:0>:<EOL><INDENT>return npdt(<NUM_LIT:1>)<EOL><DEDENT>x_2 = tf.square(x)<EOL>even_sum = tf.zeros_like(x)<EOL>odd_sum = tf.zeros_like(x)<EOL>x_2n = x_2  <EOL>for n in range(<NUM_LIT:1>, series_order + <NUM_LIT:1>):<EOL><INDENT>y = npdt(_double_factorial(<NUM_LIT:2> * n - <NUM_LIT:1>)) / x_2n<EOL>if n % <NUM_LIT:2>:<EOL><INDENT>odd_sum += y<EOL><DEDENT>else:<EOL><INDENT>even_sum += y<EOL><DEDENT>x_2n *= x_2<EOL><DEDENT>return <NUM_LIT:1.> + even_sum - odd_sum<EOL>

Calculates the asymptotic series used in log_ndtr.


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.


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

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


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.


q = sampling_dist_q<EOL>with tf.name_scope(name):<EOL><INDENT>z = _get_samples(q, z, n, seed)<EOL>log_p_z = log_p(z)<EOL>q_log_prob_z = q.log_prob(z)<EOL>def _importance_sampler_positive_f(log_f_z):<EOL><INDENT>log_values = log_f_z + log_p_z - q_log_prob_z<EOL>return _logspace_mean(log_values)<EOL><DEDENT>f_z = f(z)<EOL>log_f_plus_z = tf.math.log1p(tf.nn.relu(f_z))<EOL>log_f_minus_z = tf.math.log1p(tf.nn.relu(-<NUM_LIT:1.> * f_z))<EOL>log_f_plus_integral = _importance_sampler_positive_f(log_f_plus_z)<EOL>log_f_minus_integral = _importance_sampler_positive_f(log_f_minus_z)<EOL><DEDENT>return tf.math.exp(log_f_plus_integral) - tf.math.exp(log_f_minus_integral)<EOL>

r"""Monte Carlo estimate of \\(E_p[f(Z)] = E_q[f(Z) p(Z) / q(Z)]\\).

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

    \\(n^{-1} sum_{i=1}^n [ f(z_i) p(z_i) / q(z_i) ],  z_i ~ q,\\)
    \\(\approx E_q[ f(Z) p(Z) / q(Z) ]\\)
    \\(=       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.

    If `f >= 0`, it is up to 2x more efficient to exponentiate the result of
    `expectation_importance_sampler_logspace` applied to `Log[f]`.

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

    Args:
      f: Callable mapping samples from `sampling_dist_q` to `Tensors` with shape
        broadcastable to `q.batch_shape`.
        For example, `f` works "just like" `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" `sampling_dist_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:
      The importance sampling estimate.  `Tensor` with `shape` equal
        to batch shape of `q`, and `dtype` = `q.dtype`.


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


<EOL>center = tf.stop_gradient(_sample_max(log_values))<EOL>centered_values = tf.math.exp(log_values - center)<EOL>log_mean_of_values = tf.math.log(_sample_mean(centered_values)) + center<EOL>return log_mean_of_values<EOL>

Evaluate `Log[E[values]]` in a stable manner.

    Args:
      log_values:  `Tensor` holding `Log[values]`.

    Returns:
      `Tensor` of same `dtype` as `log_values`, reduced across dim 0.
        `Log[Mean[values]]`.


def expectation(f, samples, log_prob=None, use_reparametrization=True,<EOL>axis=<NUM_LIT:0>, keep_dims=False, name=None):

with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>if not callable(f):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>if use_reparametrization:<EOL><INDENT>return tf.math.reduce_mean(<EOL>input_tensor=f(samples), axis=axis, keepdims=keep_dims)<EOL><DEDENT>else:<EOL><INDENT>if not callable(log_prob):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>stop = tf.stop_gradient  <EOL>x = stop(samples)<EOL>logpx = log_prob(x)<EOL>fx = f(x)  <EOL>fx += stop(fx) * (logpx - stop(logpx))  <EOL>return tf.math.reduce_mean(input_tensor=fx, axis=axis, keepdims=keep_dims)<EOL><DEDENT><DEDENT>

r"""Computes the Monte-Carlo approximation of \\(E_p[f(X)]\\).

    This function computes the Monte-Carlo approximation of an expectation, i.e.,

    \\(E_p[f(X)] \approx= m^{-1} sum_i^m f(x_j),  x_j\  ~iid\ p(X)\\)

    where:

    - `x_j = samples[j, ...]`,
    - `log(p(samples)) = log_prob(samples)` and
    - `m = prod(shape(samples)[axis])`.

    Tricks: Reparameterization and Score-Gradient

    When p is "reparameterized", i.e., a diffeomorphic transformation of a
    parameterless distribution (e.g.,
    `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and
    expectation, i.e.,
    grad[ Avg{ \\(s_i : i=1...n\\) } ] = Avg{ grad[\\(s_i\\)] : i=1...n } where
    S_n = Avg{\\(s_i\\)}` and `\\(s_i = f(x_i), x_i ~ p\\).

    However, if p is not reparameterized, TensorFlow's gradient will be incorrect
    since the chain-rule stops at samples of non-reparameterized distributions.
    (The non-differentiated result, `approx_expectation`, is the same regardless
    of `use_reparametrization`.) In this circumstance using the Score-Gradient
    trick results in an unbiased gradient, i.e.,

    ```none
    grad[ E_p[f(X)] ]
    = grad[ int dx p(x) f(x) ]
    = int dx grad[ p(x) f(x) ]
    = int dx [ p'(x) f(x) + p(x) f'(x) ]
    = int dx p(x) [p'(x) / p(x) f(x) + f'(x) ]
    = int dx p(x) grad[ f(x) p(x) / stop_grad[p(x)] ]
    = E_p[ grad[ f(x) p(x) / stop_grad[p(x)] ] ]
    ```

    Unless p is not reparametrized, it is usually preferable to
    `use_reparametrization = True`.

    Warning: users are responsible for verifying `p` is a "reparameterized"
    distribution.

    Example Use:

    ```python
    import tensorflow_probability as tfp
    tfd = tfp.distributions

    # Monte-Carlo approximation of a reparameterized distribution, e.g., Normal.

    num_draws = int(1e5)
    p = tfd.Normal(loc=0., scale=1.)
    q = tfd.Normal(loc=1., scale=2.)
    exact_kl_normal_normal = tfd.kl_divergence(p, q)
    # ==> 0.44314718
    approx_kl_normal_normal = tfp.monte_carlo.expectation(
        f=lambda x: p.log_prob(x) - q.log_prob(x),
        samples=p.sample(num_draws, seed=42),
        log_prob=p.log_prob,
        use_reparametrization=(p.reparameterization_type
                               == distribution.FULLY_REPARAMETERIZED))
    # ==> 0.44632751
    # Relative Error: <1%

    # Monte-Carlo approximation of non-reparameterized distribution, e.g., Gamma.

    num_draws = int(1e5)
    p = ds.Gamma(concentration=1., rate=1.)
    q = ds.Gamma(concentration=2., rate=3.)
    exact_kl_gamma_gamma = tfd.kl_divergence(p, q)
    # ==> 0.37999129
    approx_kl_gamma_gamma = tfp.monte_carlo.expectation(
        f=lambda x: p.log_prob(x) - q.log_prob(x),
        samples=p.sample(num_draws, seed=42),
        log_prob=p.log_prob,
        use_reparametrization=(p.reparameterization_type
                               == distribution.FULLY_REPARAMETERIZED))
    # ==> 0.37696719
    # Relative Error: <1%

    # For comparing the gradients, see `monte_carlo_test.py`.
    ```

    Note: The above example is for illustration only. To compute approximate
    KL-divergence, the following is preferred:

    ```python
    approx_kl_p_q = tfp.vi.monte_carlo_csiszar_f_divergence(
        f=bf.kl_reverse,
        p_log_prob=q.log_prob,
        q=p,
        num_draws=num_draws)
    ```

    Args:
      f: Python callable which can return `f(samples)`.
      samples: `Tensor` of samples used to form the Monte-Carlo approximation of
        \\(E_p[f(X)]\\).  A batch of samples should be indexed by `axis`
        dimensions.
      log_prob: Python callable which can return `log_prob(samples)`. Must
        correspond to the natural-logarithm of the pdf/pmf of each sample. Only
        required/used if `use_reparametrization=False`.
        Default value: `None`.
      use_reparametrization: Python `bool` indicating that the approximation
        should use the fact that the gradient of samples is unbiased. Whether
        `True` or `False`, this arg only affects the gradient of the resulting
        `approx_expectation`.
        Default value: `True`.
      axis: The dimensions to average. If `None`, averages all
        dimensions.
        Default value: `0` (the left-most dimension).
      keep_dims: If True, retains averaged dimensions using size `1`.
        Default value: `False`.
      name: A `name_scope` for operations created by this function.
        Default value: `None` (which implies "expectation").

    Returns:
      approx_expectation: `Tensor` corresponding to the Monte-Carlo approximation
        of \\(E_p[f(X)]\\).

    Raises:
      ValueError: if `f` is not a Python `callable`.
      ValueError: if `use_reparametrization=False` and `log_prob` is not a Python
        `callable`.


return tf.math.reduce_mean(input_tensor=values, axis=[<NUM_LIT:0>])<EOL>

Mean over sample indices.  In this module this is always [0].


return tf.math.reduce_max(input_tensor=values, axis=[<NUM_LIT:0>])<EOL>

Max over sample indices.  In this module this is always [0].


with tf.name_scope('<STR_LIT>'):<EOL><INDENT>if (n is None) == (z is None):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>' % (n, z))<EOL><DEDENT>if n is not None:<EOL><INDENT>return dist.sample(n, seed=seed)<EOL><DEDENT>else:<EOL><INDENT>return tf.convert_to_tensor(value=z, name='<STR_LIT:z>')<EOL><DEDENT><DEDENT>


return self is other<EOL>

Determine if this `ReparameterizationType` is equal to another.

        Since RepaparameterizationType instances are constant static global
        instances, equality checks if two instances' id() values are equal.

        Args:
          other: Object to compare against.

        Returns:
          `self is other`.


def _matmul(a, b,<EOL>transpose_a=False, transpose_b=False,<EOL>adjoint_a=False, adjoint_b=False,<EOL>a_is_sparse=False, b_is_sparse=False,<EOL>name=None):

if a_is_sparse or b_is_sparse:<EOL><INDENT>raise NotImplementedError('<STR_LIT>')<EOL><DEDENT>if transpose_a or adjoint_a:<EOL><INDENT>a = _matrix_transpose(a, conjugate=adjoint_a)<EOL><DEDENT>if transpose_b or adjoint_b:<EOL><INDENT>b = _matrix_transpose(b, conjugate=adjoint_b)<EOL><DEDENT>return np.matmul(a, b)<EOL>


def _argsort(values, axis=-<NUM_LIT:1>, direction='<STR_LIT>', stable=False, name=None):

if direction == '<STR_LIT>':<EOL><INDENT>pass<EOL><DEDENT>elif direction == '<STR_LIT>':<EOL><INDENT>values = np.negative(values)<EOL><DEDENT>else:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(direction))<EOL><DEDENT>return np.argsort(values, axis, kind='<STR_LIT>' if stable else '<STR_LIT>')<EOL>


def _sort(values, axis=-<NUM_LIT:1>, direction='<STR_LIT>', stable=False, name=None):

if direction == '<STR_LIT>':<EOL><INDENT>pass<EOL><DEDENT>elif direction == '<STR_LIT>':<EOL><INDENT>values = np.negative(values)<EOL><DEDENT>else:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(direction))<EOL><DEDENT>result = np.sort(values, axis, kind='<STR_LIT>' if stable else '<STR_LIT>')<EOL>if direction == '<STR_LIT>':<EOL><INDENT>return np.negative(result)<EOL><DEDENT>return result<EOL>


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 _max_mask_non_finite(x, axis=-<NUM_LIT:1>, keepdims=False, mask=<NUM_LIT:0>):

m = np.max(x, axis=_astuple(axis), keepdims=keepdims)<EOL>needs_masking = ~np.isfinite(m)<EOL>if needs_masking.ndim > <NUM_LIT:0>:<EOL><INDENT>m[needs_masking] = mask<EOL><DEDENT>elif needs_masking:<EOL><INDENT>m = mask<EOL><DEDENT>return m<EOL>

Returns `max` or `mask` if `max` is not finite.


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>static_spec = tf_inspect.getfullargspec(static_fn)<EOL>if original_spec != static_spec:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(<EOL>original_spec, static_spec, original_fn))<EOL><DEDENT>@decorator.decorator<EOL>def wrap(wrapped_fn, *args, **kwargs):<EOL><INDENT>del wrapped_fn<EOL>[args_, kwargs_], all_static = _maybe_get_static_args([args, kwargs])<EOL>if all_static:<EOL><INDENT>return static_fn(*args_, **kwargs_)<EOL><DEDENT>return original_fn(*args, **kwargs)<EOL><DEDENT>return wrap(original_fn)<EOL>

Wraps original_fn, preferring to call static_fn when inputs are static.


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.


if pred in {<NUM_LIT:0>, <NUM_LIT:1>}:  <EOL><INDENT>pred_value = bool(pred)<EOL><DEDENT>elif isinstance(pred, bool):<EOL><INDENT>pred_value = pred<EOL><DEDENT>elif isinstance(pred, tf.Tensor):<EOL><INDENT>pred_value = tf.get_static_value(pred)<EOL>if pred_value is None:<EOL><INDENT>pred_value = c_api.TF_TryEvaluateConstant_wrapper(pred.graph._c_graph,<EOL>pred._as_tf_output())<EOL><DEDENT><DEDENT>else:<EOL><INDENT>raise TypeError('<STR_LIT>'<EOL>'<STR_LIT>'.format(pred))<EOL><DEDENT>return pred_value<EOL>

Helper function for statically evaluating predicates in `cond`.


if tensorshape is None:<EOL><INDENT>shape_tensor = (shape_tensor_fn() if callable(shape_tensor_fn)<EOL>else shape_tensor_fn)<EOL>if (hasattr(shape_tensor, '<STR_LIT>') and<EOL>hasattr(shape_tensor.shape, '<STR_LIT>')):<EOL><INDENT>ndims_ = tensorshape_util.num_elements(shape_tensor.shape)<EOL><DEDENT>else:<EOL><INDENT>ndims_ = len(shape_tensor)<EOL><DEDENT>ndims_fn = lambda: tf.size(input=shape_tensor)<EOL><DEDENT>else:<EOL><INDENT>ndims_ = tensorshape_util.rank(tensorshape)<EOL>ndims_fn = lambda: tf.size(input=shape_tensor_fn()  <EOL>if callable(shape_tensor_fn)<EOL>else shape_tensor_fn)<EOL><DEDENT>return ndims_fn() if ndims_ is None else ndims_<EOL>

Computes `rank` given a `Tensor`'s `shape`.


def cond(pred, true_fn=None, false_fn=None, name=None):

if not callable(true_fn):<EOL><INDENT>raise TypeError('<STR_LIT>')<EOL><DEDENT>if not callable(false_fn):<EOL><INDENT>raise TypeError('<STR_LIT>')<EOL><DEDENT>pred_value = _get_static_predicate(pred)<EOL>if pred_value is not None:<EOL><INDENT>if pred_value:<EOL><INDENT>return true_fn()<EOL><DEDENT>else:<EOL><INDENT>return false_fn()<EOL><DEDENT><DEDENT>else:<EOL><INDENT>return tf.cond(pred=pred, true_fn=true_fn, false_fn=false_fn, name=name)<EOL><DEDENT>

Return either `true_fn()` if predicate `pred` is true else `false_fn()`.

    If `pred` is a bool or has a constant value, we return either `true_fn()`
    or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both.

    Arguments:
      pred: A scalar determining whether to return the result of `true_fn` or
        `false_fn`.
      true_fn: The callable to be performed if pred is true.
      false_fn: The callable to be performed if pred is false.
      name: Optional name prefix when using `tf.cond`.

    Returns:
      Tensors returned by the call to either `true_fn` or `false_fn`.

    Raises:
      TypeError: If `true_fn` or `false_fn` is not callable.


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 assert_finite(x, data=None, summarize=None, message=None, name=None):

with tf.compat.v2.name_scope(name or '<STR_LIT>'):<EOL><INDENT>x_ = tf.get_static_value(x)<EOL>if x_ is not None:<EOL><INDENT>if ~np.all(np.isfinite(x_)):<EOL><INDENT>raise ValueError(message)<EOL><DEDENT>return x<EOL><DEDENT>assertion = tf.compat.v1.assert_equal(<EOL>tf.math.is_finite(x), tf.ones_like(x, tf.bool),<EOL>data=data, summarize=summarize, message=message)<EOL>with tf.control_dependencies([assertion]):<EOL><INDENT>return tf.identity(x)<EOL><DEDENT><DEDENT>

Assert all elements of `x` are finite.

    Args:
      x:  Numeric `Tensor`.
      data:  The tensors to print out if the condition is False.  Defaults to
        error message and first few entries of `x`.
      summarize: Print this many entries of each tensor.
      message: A string to prefix to the default message.
      name: A name for this operation (optional).
        Defaults to "assert_finite".

    Returns:
      Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
      If static checks determine `x` has correct rank, a `no_op` is returned.

    Raises:
      ValueError:  If static checks determine `x` has wrong rank.


def assert_rank_at_most(x, rank, data=None, summarize=None, message=None,<EOL>name=None):

with tf.compat.v2.name_scope(name or '<STR_LIT>'):<EOL><INDENT>return tf.compat.v1.assert_less_equal(<EOL>tf.rank(x), rank, data=data, summarize=summarize, message=message)<EOL><DEDENT>

Assert `x` has rank equal to `rank` or smaller.

    Example of adding a dependency to an operation:

    ```python
    with tf.control_dependencies([tf.assert_rank_at_most(x, 2)]):
      output = tf.reduce_sum(x)
    ```

    Args:
      x:  Numeric `Tensor`.
      rank:  Scalar `Tensor`.
      data:  The tensors to print out if the condition is False.  Defaults to
        error message and first few entries of `x`.
      summarize: Print this many entries of each tensor.
      message: A string to prefix to the default message.
      name: A name for this operation (optional).
        Defaults to "assert_rank_at_most".

    Returns:
      Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
      If static checks determine `x` has correct rank, a `no_op` is returned.

    Raises:
      ValueError:  If static checks determine `x` has wrong rank.


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 prefer_static.equal(prefer_static.rank_from_shape(shape), <NUM_LIT:0>)<EOL>

Returns `True` `Tensor` if `Tensor` shape implies a scalar.


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


parameters = dict(locals()) if parameters is None else parameters<EOL>name = name or (("<STR_LIT>" if bijector is None else bijector.name) +<EOL>(distribution.name or "<STR_LIT>"))<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._kwargs_split_fn = (_default_kwargs_split_fn<EOL>if kwargs_split_fn is None<EOL>else kwargs_split_fn)<EOL>self._zero = tf.constant(<NUM_LIT:0>, dtype=tf.int32, name="<STR_LIT>")<EOL>self._empty = tf.constant([], dtype=tf.int32, name="<STR_LIT>")<EOL>self._override_batch_shape = self._maybe_validate_shape_override(<EOL>batch_shape, distribution.is_scalar_batch(), validate_args,<EOL>"<STR_LIT>")<EOL>self._is_batch_override = prefer_static.logical_not(<EOL>prefer_static.equal(<EOL>prefer_static.rank_from_shape(self._override_batch_shape),<EOL>self._zero))<EOL>self._is_maybe_batch_override = bool(<EOL>tf.get_static_value(self._override_batch_shape) is None or<EOL>tf.get_static_value(self._override_batch_shape).size != <NUM_LIT:0>)<EOL>self._override_event_shape = self._maybe_validate_shape_override(<EOL>event_shape, distribution.is_scalar_event(), validate_args,<EOL>"<STR_LIT>")<EOL>self._is_event_override = prefer_static.logical_not(<EOL>prefer_static.equal(<EOL>prefer_static.rank_from_shape(self._override_event_shape),<EOL>self._zero))<EOL>self._is_maybe_event_override = bool(<EOL>tf.get_static_value(self._override_event_shape) is None or<EOL>tf.get_static_value(self._override_event_shape).size != <NUM_LIT:0>)<EOL>self._needs_rotation = prefer_static.reduce_all([<EOL>self._is_event_override,<EOL>prefer_static.logical_not(self._is_batch_override),<EOL>prefer_static.logical_not(distribution.is_scalar_batch())])<EOL>override_event_ndims = prefer_static.rank_from_shape(<EOL>self._override_event_shape)<EOL>self._rotate_ndims = _pick_scalar_condition(<EOL>self._needs_rotation, override_event_ndims, <NUM_LIT:0>)<EOL>self._reduce_event_indices = tf.range(<EOL>self._rotate_ndims - override_event_ndims, self._rotate_ndims)<EOL><DEDENT>self._distribution = distribution<EOL>self._bijector = bijector<EOL>super(TransformedDistribution, self).__init__(<EOL>dtype=self._distribution.dtype,<EOL>reparameterization_type=self._distribution.reparameterization_type,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=self._distribution.allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=(distribution._graph_parents +  <EOL>bijector.graph_parents),<EOL>name=name)<EOL>

Construct a Transformed Distribution.

        Args:
          distribution: The base distribution instance to transform. Typically an
            instance of `Distribution`.
          bijector: The object responsible for calculating the transformation.
            Typically an instance of `Bijector`.
          batch_shape: `integer` vector `Tensor` which overrides `distribution`
            `batch_shape`; valid only if `distribution.is_scalar_batch()`.
          event_shape: `integer` vector `Tensor` which overrides `distribution`
            `event_shape`; valid only if `distribution.is_scalar_event()`.
          kwargs_split_fn: Python `callable` which takes a kwargs `dict` and returns
            a tuple of kwargs `dict`s for each of the `distribution` and `bijector`
            parameters respectively.
            Default value: `_default_kwargs_split_fn` (i.e.,
                `lambda kwargs: (kwargs.get('distribution_kwargs', {}),
                                 kwargs.get('bijector_kwargs', {}))`)
          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.
          parameters: Locals dict captured by subclass constructor, to be used for
            copy/slice re-instantiation operations.
          name: Python `str` name prefixed to Ops created by this class. Default:
            `bijector.name + distribution.name`.


return self._distribution<EOL>


return self._bijector<EOL>


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

x = self._maybe_rotate_dims(x, rotate_right=True)<EOL>log_prob = self.distribution.log_prob(x, **distribution_kwargs)<EOL>if self._is_maybe_event_override:<EOL><INDENT>log_prob = tf.reduce_sum(<EOL>input_tensor=log_prob, axis=self._reduce_event_indices)<EOL><DEDENT>log_prob += tf.cast(ildj, log_prob.dtype)<EOL>if self._is_maybe_event_override and isinstance(event_ndims, int):<EOL><INDENT>tensorshape_util.set_shape(<EOL>log_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 log_prob<EOL>

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


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 _maybe_validate_shape_override(self, override_shape, base_is_scalar,<EOL>validate_args, name):

if override_shape is None:<EOL><INDENT>override_shape = []<EOL><DEDENT>override_shape = tf.convert_to_tensor(<EOL>value=override_shape, dtype=tf.int32, name=name)<EOL>if not dtype_util.is_integer(override_shape.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>override_is_scalar = _is_scalar_from_shape_tensor(override_shape)<EOL>if tf.get_static_value(override_is_scalar):<EOL><INDENT>return self._empty<EOL><DEDENT>dynamic_assertions = []<EOL>if tensorshape_util.rank(override_shape.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(override_shape.shape) != <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>dynamic_assertions += [<EOL>assert_util.assert_rank(<EOL>override_shape, <NUM_LIT:1>, message="<STR_LIT>")<EOL>]<EOL><DEDENT>if tf.get_static_value(override_shape) is not None:<EOL><INDENT>if any(s < <NUM_LIT:0> for s in tf.get_static_value(override_shape)):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>dynamic_assertions += [<EOL>assert_util.assert_non_negative(<EOL>override_shape,<EOL>message="<STR_LIT>")<EOL>]<EOL><DEDENT>is_both_nonscalar = prefer_static.logical_and(<EOL>prefer_static.logical_not(base_is_scalar),<EOL>prefer_static.logical_not(override_is_scalar))<EOL>if tf.get_static_value(is_both_nonscalar) is not None:<EOL><INDENT>if tf.get_static_value(is_both_nonscalar):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>dynamic_assertions += [<EOL>assert_util.assert_equal(<EOL>is_both_nonscalar, False, message="<STR_LIT>")<EOL>]<EOL><DEDENT>if not dynamic_assertions:<EOL><INDENT>return override_shape<EOL><DEDENT>return distribution_util.with_dependencies(<EOL>dynamic_assertions, override_shape)<EOL>

Helper to __init__ which ensures override batch/event_shape are valid.


def _maybe_rotate_dims(self, x, rotate_right=False):

needs_rotation_const = tf.get_static_value(self._needs_rotation)<EOL>if needs_rotation_const is not None and not needs_rotation_const:<EOL><INDENT>return x<EOL><DEDENT>ndims = prefer_static.rank(x)<EOL>n = (ndims - self._rotate_ndims) if rotate_right else self._rotate_ndims<EOL>perm = prefer_static.concat([<EOL>prefer_static.range(n, ndims), prefer_static.range(<NUM_LIT:0>, n)], axis=<NUM_LIT:0>)<EOL>return tf.transpose(a=x, perm=perm)<EOL>

Helper which rolls left event_dims left or right event_dims right.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([logits, probs, temperature], tf.float32)<EOL>self._temperature = tf.convert_to_tensor(<EOL>value=temperature, name="<STR_LIT>", dtype=dtype)<EOL>if validate_args:<EOL><INDENT>with tf.control_dependencies(<EOL>[assert_util.assert_positive(temperature)]):<EOL><INDENT>self._temperature = tf.identity(self._temperature)<EOL><DEDENT><DEDENT>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>logits=logits, probs=probs, validate_args=validate_args, dtype=dtype)<EOL>super(RelaxedBernoulli, self).__init__(<EOL>distribution=logistic.Logistic(<EOL>self._logits / self._temperature,<EOL><NUM_LIT:1.> / self._temperature,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name + "<STR_LIT>"),<EOL>bijector=sigmoid_bijector.Sigmoid(validate_args=validate_args),<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>self._parameters = parameters<EOL>

Construct RelaxedBernoulli distributions.

        Args:
          temperature: An 0-D `Tensor`, representing the temperature
            of a set of RelaxedBernoulli distributions. The temperature should be
            positive.
          logits: An N-D `Tensor` representing the log-odds
            of a positive event. Each entry in the `Tensor` parametrizes
            an independent RelaxedBernoulli distribution where the probability of an
            event is sigmoid(logits). Only one of `logits` or `probs` should be
            passed in.
          probs: An N-D `Tensor` representing the probability of a positive event.
            Each entry in the `Tensor` parameterizes an independent Bernoulli
            distribution. 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.

        Raises:
          ValueError: If both `probs` and `logits` are passed, or if neither.


return self._temperature<EOL>

Distribution parameter for the location.


return self._logits<EOL>


return self._probs<EOL>


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

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>final_batch_shape = distribution_util.get_broadcast_shape(<EOL>a.low, b.low, a.high, b.high)<EOL>dtype = dtype_util.common_dtype(<EOL>[a.low, a.high, b.low, b.high], tf.float32)<EOL>return tf.where((b.low <= a.low) & (a.high <= b.high),<EOL>tf.math.log(b.high - b.low) - tf.math.log(a.high - a.low),<EOL>tf.broadcast_to(<EOL>dtype_util.as_numpy_dtype(dtype)(np.inf),<EOL>final_batch_shape))<EOL><DEDENT>

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

    Note that the KL divergence is infinite if the support of `a` is not a subset
    of the support of `b`.

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

    Returns:
      Batchwise KL(a || b)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([low, high], tf.float32)<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>with tf.control_dependencies([<EOL>assert_util.assert_less(<EOL>low, high, message="<STR_LIT>")<EOL>] if validate_args else []):<EOL><INDENT>self._low = tf.identity(low)<EOL>self._high = tf.identity(high)<EOL>dtype_util.assert_same_float_dtype([self._low, self._high])<EOL><DEDENT><DEDENT>super(Uniform, self).__init__(<EOL>dtype=self._low.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._low,<EOL>self._high],<EOL>name=name)<EOL>

Initialize a batch of Uniform distributions.

        Args:
          low: Floating point tensor, lower boundary of the output interval. Must
            have `low < high`.
          high: Floating point tensor, upper boundary of the output interval. Must
            have `low < high`.
          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:
          InvalidArgumentError: if `low >= high` and `validate_args=False`.


return self._low<EOL>


return self._high<EOL>


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


def _broadcast_event_and_samples(event, samples, event_ndims):

<EOL>samples_shape = tf.concat(<EOL>[tf.shape(input=samples)[:-event_ndims - <NUM_LIT:1>],<EOL>tf.shape(input=samples)[tf.rank(samples) - event_ndims:]],<EOL>axis=<NUM_LIT:0>)<EOL>event *= tf.ones(samples_shape, dtype=event.dtype)<EOL>event = tf.expand_dims(event, axis=-event_ndims - <NUM_LIT:1>)<EOL>samples *= tf.ones_like(event, dtype=samples.dtype)<EOL>return event, samples<EOL>


parameters = dict(locals())<EOL>with tf.name_scope(name):<EOL><INDENT>self._samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>')<EOL>self._event_ndims = event_ndims<EOL>self._samples_axis = (<EOL>(tensorshape_util.rank(self.samples.shape) or tf.rank(self.samples)) -<EOL>self._event_ndims - <NUM_LIT:1>)<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_rank_at_least(self._samples, event_ndims + <NUM_LIT:1>)]):<EOL><INDENT>samples_shape = distribution_util.prefer_static_shape(self._samples)<EOL>self._num_samples = samples_shape[self._samples_axis]<EOL><DEDENT><DEDENT>super(Empirical, self).__init__(<EOL>dtype=self._samples.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._samples],<EOL>name=name)<EOL>

Initialize `Empirical` distributions.

        Args:
          samples: Numeric `Tensor` of shape [B1, ..., Bk, S, E1, ..., En]`,
            `k, n >= 0`. Samples or batches of samples on which the distribution
            is based. The first `k` dimensions index into a batch of independent
            distributions. Length of `S` dimension determines number of samples
            in each multiset. The last `n` dimension represents samples for each
            distribution. n is specified by argument event_ndims.
          event_ndims: Python `int32`, default `0`. number of dimensions for each
            event. When `0` this distribution has scalar samples. When `1` this
            distribution has vector-like samples.
          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 the rank of `samples` < event_ndims + 1.


return self._samples<EOL>


return self._num_samples<EOL>


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


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

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>final_batch_shape = distribution_util.get_broadcast_shape(<EOL>a.concentration, b.concentration, a.scale, b.scale)<EOL>common_type = dtype_util.common_dtype(<EOL>[a.concentration, b.concentration, a.scale, b.scale], tf.float32)<EOL>return tf.where(<EOL>a.scale >= b.scale,<EOL>b.concentration * (tf.math.log(a.scale) - tf.math.log(b.scale)) +<EOL>tf.math.log(a.concentration) - tf.math.log(b.concentration) +<EOL>b.concentration / a.concentration - <NUM_LIT:1.0>,<EOL>tf.broadcast_to(tf.cast(np.inf, common_type), final_batch_shape))<EOL><DEDENT>

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

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

    Returns:
      Batchwise KL(a || b)


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

Construct Pareto distribution with `concentration` and `scale`.

        Args:
          concentration: Floating point tensor. Must contain only positive values.
          scale: Floating point tensor, equivalent to `mode`. `scale` also
            restricts the domain of this distribution to be in `[scale, inf)`.
            Must contain only positive values. Default value: `1`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs. Default value: `False` (i.e. do not validate args).
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
            Default value: `True`.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: 'Pareto'.


return self._scale<EOL>

Scale parameter and also the lower bound of the support.


return self._concentration<EOL>

Concentration parameter for this distribution.


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


@kullback_leibler.RegisterKL(Normal, Normal)<EOL>def _kl_normal_normal(n_a, n_b, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>one = tf.constant(<NUM_LIT:1>, dtype=n_a.dtype)<EOL>two = tf.constant(<NUM_LIT:2>, dtype=n_a.dtype)<EOL>half = tf.constant(<NUM_LIT:0.5>, dtype=n_a.dtype)<EOL>s_a_squared = tf.square(n_a.scale)<EOL>s_b_squared = tf.square(n_b.scale)<EOL>ratio = s_a_squared / s_b_squared<EOL>return (tf.square(n_a.loc - n_b.loc) / (two * s_b_squared) + half *<EOL>(ratio - one - tf.math.log(ratio)))<EOL><DEDENT>

Calculate the batched KL divergence KL(n_a || n_b) with n_a and n_b Normal.

    Args:
      n_a: instance of a Normal distribution object.
      n_b: instance of a Normal distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_normal_normal".

    Returns:
      Batchwise KL(n_a || n_b)


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

Construct Normal distributions with mean and stddev `loc` and `scale`.

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

        Args:
          loc: Floating point tensor; the means of the distribution(s).
          scale: Floating point tensor; the stddevs of the distribution(s).
            Must contain only positive values.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`,
            statistics (e.g., mean, mode, variance) use the value "`NaN`" to
            indicate the result is undefined. When `False`, an exception is raised
            if one or more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

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


return self._loc<EOL>


return self._scale<EOL>

Distribution parameter for standard deviation.


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


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

Reconstruct input `x` from a its normalized version.


parameters = dict(locals())<EOL>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.


parameters = dict(locals())<EOL>with tf.name_scope(name):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[concentration, mixing_concentration, mixing_rate],<EOL>preferred_dtype=tf.float32)<EOL>concentration = tf.convert_to_tensor(<EOL>value=concentration, name="<STR_LIT>", dtype=dtype)<EOL>mixing_concentration = tf.convert_to_tensor(<EOL>value=mixing_concentration, name="<STR_LIT>", dtype=dtype)<EOL>mixing_rate = tf.convert_to_tensor(<EOL>value=mixing_rate, name="<STR_LIT>", dtype=dtype)<EOL>with tf.control_dependencies([<EOL>assert_util.assert_positive(concentration),<EOL>assert_util.assert_positive(mixing_concentration),<EOL>assert_util.assert_positive(mixing_rate),<EOL>] if validate_args else []):<EOL><INDENT>self._concentration = tf.identity(concentration, name="<STR_LIT>")<EOL>self._mixing_concentration = tf.identity(<EOL>mixing_concentration, name="<STR_LIT>")<EOL>self._mixing_rate = tf.identity(mixing_rate, name="<STR_LIT>")<EOL><DEDENT>dtype_util.assert_same_float_dtype(<EOL>[self._concentration, self._mixing_concentration, self._mixing_rate])<EOL><DEDENT>super(GammaGamma, self).__init__(<EOL>dtype=self._concentration.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._concentration, self._mixing_concentration, self._mixing_rate<EOL>],<EOL>name=name)<EOL>

Initializes a batch of Gamma-Gamma distributions.

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

        Args:
          concentration: Floating point tensor, the concentration params of the
            distribution(s). Must contain only positive values.
          mixing_concentration: Floating point tensor, the concentration params of
            the mixing Gamma distribution(s). Must contain only positive values.
          mixing_rate: Floating point tensor, the rate params of the mixing Gamma
            distribution(s). Must contain only positive values.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or more
            of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          TypeError: if `concentration` and `rate` are different dtypes.


return self._concentration<EOL>


return self._mixing_concentration<EOL>

Concentration parameter for the mixing Gamma distribution.


return self._mixing_rate<EOL>

Rate parameter for the mixing Gamma distribution.