def _compute_inverse_log_det_jacobian_with_caching(<EOL>self, x, y, prefer_inverse_ldj_fn, event_ndims, kwargs):

<EOL>if not (hasattr(self, "<STR_LIT>") or<EOL>hasattr(self, "<STR_LIT>")):<EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>use_inverse_ldj_fn = (<EOL>(prefer_inverse_ldj_fn and hasattr(self, "<STR_LIT>"))<EOL>or not hasattr(self, "<STR_LIT>"))<EOL>if use_inverse_ldj_fn:<EOL><INDENT>tensor_to_use = y if y is not None else self.forward(x, **kwargs)<EOL>min_event_ndims = self.inverse_min_event_ndims<EOL><DEDENT>else:<EOL><INDENT>tensor_to_use = x if x is not None else self.inverse(y, **kwargs)<EOL>min_event_ndims = self.forward_min_event_ndims<EOL><DEDENT>if self.is_constant_jacobian:<EOL><INDENT>unreduced_ildj = self._compute_unreduced_constant_ildj_with_caching(<EOL>tensor_to_use, use_inverse_ldj_fn, kwargs)<EOL><DEDENT>else:<EOL><INDENT>unreduced_ildj = self._compute_unreduced_nonconstant_ildj_with_caching(<EOL>x, y, tensor_to_use, use_inverse_ldj_fn, kwargs)<EOL><DEDENT>return self._reduce_jacobian_det_over_event(<EOL>tf.shape(input=tensor_to_use), unreduced_ildj, min_event_ndims,<EOL>event_ndims)<EOL>

Compute ILDJ by the best available means, and ensure it is cached.

        Sub-classes of Bijector may implement either `_forward_log_det_jacobian` or
        `_inverse_log_det_jacobian`, or both, and may prefer one or the other for
        reasons of computational efficiency and/or numerics. In general, to compute
        the [I]LDJ, we need one of `x` or `y` and one of `_forward_log_det_jacobian`
        or `_inverse_log_det_jacobian` (all bijectors implement `forward` and
        `inverse`, so either of `x` and `y` may always computed from the other).

        This method encapsulates the logic of selecting the best possible method of
        computing the inverse log det jacobian, and executing it. Possible avenues
        to obtaining this value are:
          - recovering it from the cache,
          - computing it as `-_forward_log_det_jacobian(x)`, where `x` may need to
            be computed as `inverse(y)`, or
          - computing it as `_inverse_log_det_jacobian(y)`, where `y` may need to
            be computed as `forward(x)`.

        To make things more interesting, we may be able to take advantage of the
        knowledge that the jacobian is constant, for given `event_ndims`, and may
        even cache that result.

        To make things even more interesting still, we may need to perform an
        additional `reduce_sum` over the event dims of an input after computing the
        ILDJ (see `_reduce_jacobian_det_over_event` for the reduction logic). In
        order to prevent spurious TF graph dependencies on past inputs in cached
        results, we need to take care to do this reduction after the cache lookups.

        This method takes care of all the above concerns.

        Args:
          x: a `Tensor`, the pre-Bijector transform value at whose post-Bijector
            transform value the ILDJ is to be computed. Can be `None` as long as
            `y` is not `None`.
          y: a `Tensor`, a point in the output space of the Bijector's `forward`
            transformation, at whose value the ILDJ is to be computed. Can be
            `None` as long as `x` is not `None`.
          prefer_inverse_ldj_fn: Python `bool`, if `True`, will strictly prefer to
            use the `_inverse_log_det_jacobian` to compute ILDJ; else, will
            strictly prefer to use `_forward_log_det_jacobian`.
          event_ndims: int-like `Tensor`, the number of dims of an event (in the
            pre- or post-transformed space, as appropriate). These need to be summed
            over to compute the total ildj.
          kwargs: dictionary of keyword args that will be passed to calls to
            `self.forward` or `self.inverse` (if either of those are required), and
            to `self._compute_unreduced_nonconstant_ildj_with_caching` or
            `self._compute_unreduced_constant_ildj_with_caching`, as appropriate (
            those functions use the conditioning kwargs for caching and for
            their underlying computations).

        Returns:
          ildj: a Tensor of ILDJ['s] at the given input (as specified by the args).
            Also updates the cache as needed.


def _compute_unreduced_nonconstant_ildj_with_caching(<EOL>self, x, y, tensor_to_use, use_inverse_ldj_fn, kwargs):

mapping = self._lookup(x=x, y=y, kwargs=kwargs)<EOL>if mapping.ildj is not None:<EOL><INDENT>return mapping.ildj<EOL><DEDENT>if use_inverse_ldj_fn:<EOL><INDENT>ildj = self._inverse_log_det_jacobian(tensor_to_use, **kwargs)<EOL><DEDENT>else:<EOL><INDENT>ildj = -self._forward_log_det_jacobian(tensor_to_use, **kwargs)<EOL><DEDENT>mapping = mapping.merge(x=x, y=y, ildj=ildj)<EOL>self._cache_update(mapping)<EOL>return ildj<EOL>

Helper for computing ILDJ, with caching, in the non-constant case.

        Does not do the "reduce" step which is necessary in some cases; this is left
        to the caller.

        Args:
          x: a `Tensor`, the pre-Bijector transform value at whose post-Bijector
            transform value the ILDJ is to be computed. Can be `None` as long as
            `y` is not `None`. This method only uses the value for cache
            lookup/updating.
          y: a `Tensor`, a point in the output space of the Bijector's `forward`
            transformation, at whose value the ILDJ is to be computed. Can be
            `None` as long as `x` is not `None`. This method only uses the value
            for cache lookup/updating.
          tensor_to_use: a `Tensor`, the one to actually pass to the chosen compute
            function (`_inverse_log_det_jacobian` or `_forward_log_det_jacobian`).
            It is presumed that the caller has already figured out what input to use
            (it will either be the x or y value corresponding to the location where
            we are computing the ILDJ).
          use_inverse_ldj_fn: Python `bool`, if `True`, will use the
            `_inverse_log_det_jacobian` to compute ILDJ; else, will use
            `_forward_log_det_jacobian`.
          kwargs: dictionary of keyword args that will be passed to calls to to
            `_inverse_log_det_jacobian` or `_forward_log_det_jacobian`, as well as
            for lookup/updating of the result in the cache.

        Returns:
          ildj: the (un-reduce_sum'ed) value of the ILDJ at the specified input
          location. Also updates the cache as needed.


def _compute_unreduced_constant_ildj_with_caching(<EOL>self, tensor_to_use, use_inverse_ldj_fn, kwargs):

current_graph = _get_current_graph()<EOL>if current_graph in self._constant_ildj:<EOL><INDENT>return self._constant_ildj[current_graph]<EOL><DEDENT>if use_inverse_ldj_fn:<EOL><INDENT>ildj = self._inverse_log_det_jacobian(tensor_to_use, **kwargs)<EOL><DEDENT>else:<EOL><INDENT>ildj = -self._forward_log_det_jacobian(tensor_to_use, **kwargs)<EOL><DEDENT>self._constant_ildj[current_graph] = ildj<EOL>return ildj<EOL>

Helper for computing ILDJ, with caching, in the constant-ILDJ case.

        Does not do the "reduce" step which is necessary in some cases; this is left
        to the caller.

        Args:
          tensor_to_use: a `Tensor`, to pass to the chosen compute
            function (`_inverse_log_det_jacobian` or `_forward_log_det_jacobian`).
            Since we know the ILDJ is input-independent, the actual value is not
            meaningful, although the shape and dtype may be.
          use_inverse_ldj_fn: Python `bool`, if `True`, will use the
            `_inverse_log_det_jacobian` to compute ILDJ; else, will use
            `_forward_log_det_jacobian`.
          kwargs: dictionary of keyword args that will be passed to calls to
            to `_inverse_log_det_jacobian` or `_forward_log_det_jacobian`, as well
            as for lookup/updating of the result in the cache.

        Returns:
          ildj: the (un-reduce_sum'ed) value of the ILDJ for the specified input.
            Also updates the cache as needed.


def _call_inverse_log_det_jacobian(self, y, event_ndims, name, **kwargs):

with self._name_scope(name), tf.control_dependencies(<EOL>self._check_valid_event_ndims(<EOL>min_event_ndims=self.inverse_min_event_ndims,<EOL>event_ndims=event_ndims)):<EOL><INDENT>y = tf.convert_to_tensor(value=y, name="<STR_LIT:y>")<EOL>self._maybe_assert_dtype(y)<EOL>if not self._is_injective:<EOL><INDENT>ildjs = self._inverse_log_det_jacobian(y, **kwargs)<EOL>return tuple(<EOL>self._reduce_jacobian_det_over_event(  <EOL>tf.shape(input=y),<EOL>ildj,<EOL>self.inverse_min_event_ndims, event_ndims)<EOL>for ildj in ildjs)<EOL><DEDENT>return self._compute_inverse_log_det_jacobian_with_caching(<EOL>x=None,<EOL>y=y,<EOL>prefer_inverse_ldj_fn=True,<EOL>event_ndims=event_ndims,<EOL>kwargs=kwargs)<EOL><DEDENT>

Wraps call to _inverse_log_det_jacobian, allowing extra shared logic.

        Specifically, this method
          - adds a name scope,
          - performs validations,
          - handles the special case of non-injective Bijector (skip caching and
            reduce_sum over the values at the multiple points in the preimage of `y`
            under the non-injective transformation)

        so that sub-classes don't have to worry about this stuff.

        Args:
          y: same as in `inverse_log_det_jacobian`
          event_ndims: same as in `inverse_log_det_jacobian`
          name: same as in `inverse_log_det_jacobian`
          **kwargs: same as in `inverse_log_det_jacobian`

        Returns:
          ildj: the inverse log det jacobian at `y`. Also updates the cache as
            needed.


return self._call_inverse_log_det_jacobian(y, event_ndims, name, **kwargs)<EOL>

Returns the (log o det o Jacobian o inverse)(y).

        Mathematically, returns: `log(det(dX/dY))(Y)`. (Recall that: `X=g^{-1}(Y)`.)

        Note that `forward_log_det_jacobian` is the negative of this function,
        evaluated at `g^{-1}(y)`.

        Args:
          y: `Tensor`. The input to the "inverse" Jacobian determinant evaluation.
          event_ndims: Number of dimensions in the probabilistic events being
            transformed. Must be greater than or equal to
            `self.inverse_min_event_ndims`. The result is summed over the final
            dimensions to produce a scalar Jacobian determinant for each event, i.e.
            it has shape `rank(y) - event_ndims` dimensions.
          name: The name to give this op.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          ildj: `Tensor`, if this bijector is injective.
            If not injective, returns the tuple of local log det
            Jacobians, `log(det(Dg_i^{-1}(y)))`, where `g_i` is the restriction
            of `g` to the `ith` partition `Di`.

        Raises:
          TypeError: if `self.dtype` is specified and `y.dtype` is not
            `self.dtype`.
          NotImplementedError: if `_inverse_log_det_jacobian` is not implemented.


def _call_forward_log_det_jacobian(self, x, event_ndims, name, **kwargs):

Wraps call to _forward_log_det_jacobian, allowing extra shared logic.

        Specifically, this method
          - adds a name scope,
          - performs validations,
          - handles the special case of non-injective Bijector (forward jacobian is
            ill-defined in this case and we raise an exception)

        so that sub-classes don't have to worry about this stuff.

        Args:
          x: same as in `forward_log_det_jacobian`
          event_ndims: same as in `forward_log_det_jacobian`
          name: same as in `forward_log_det_jacobian`
          **kwargs: same as in `forward_log_det_jacobian`

        Returns:
          fldj: the forward log det jacobian at `x`. Also updates the cache as
          needed.


return self._call_forward_log_det_jacobian(x, event_ndims, name, **kwargs)<EOL>

Returns both the forward_log_det_jacobian.

        Args:
          x: `Tensor`. The input to the "forward" Jacobian determinant evaluation.
          event_ndims: Number of dimensions in the probabilistic events being
            transformed. Must be greater than or equal to
            `self.forward_min_event_ndims`. The result is summed over the final
            dimensions to produce a scalar Jacobian determinant for each event, i.e.
            it has shape `rank(x) - event_ndims` dimensions.
          name: The name to give this op.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          `Tensor`, if this bijector is injective.
            If not injective this is not implemented.

        Raises:
          TypeError: if `self.dtype` is specified and `y.dtype` is not
            `self.dtype`.
          NotImplementedError: if neither `_forward_log_det_jacobian`
            nor {`_inverse`, `_inverse_log_det_jacobian`} are implemented, or
            this is a non-injective bijector.


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

Helper function to standardize op scope.


if (self.dtype is not None and<EOL>not dtype_util.base_equal(self.dtype, x.dtype)):<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>" % (x.dtype, self.dtype))<EOL><DEDENT>

Helper to check dtype when self.dtype is known.


<EOL>mapping = mapping.merge(<EOL>mapping=self._lookup(mapping.x, mapping.y, mapping.kwargs))<EOL>if mapping.x is None:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>self._from_x[mapping.x][mapping.subkey] = mapping.remove("<STR_LIT:x>")<EOL>

Helper which stores new mapping info in the forward dict.


<EOL>mapping = mapping.merge(<EOL>mapping=self._lookup(mapping.x, mapping.y, mapping.kwargs))<EOL>if mapping.y is None:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>self._from_y[mapping.y][mapping.subkey] = mapping.remove("<STR_LIT:y>")<EOL>

Helper which stores new mapping info in the inverse dict.


if mapping.x is not None and mapping.subkey in self._from_x[mapping.x]:<EOL><INDENT>self._cache_by_x(mapping)<EOL><DEDENT>if mapping.y is not None and mapping.subkey in self._from_y[mapping.y]:<EOL><INDENT>self._cache_by_y(mapping)<EOL><DEDENT>

Helper which updates only those cached entries that already exist.


def _lookup(self, x=None, y=None, kwargs=None):

mapping = _Mapping(x=x, y=y, kwargs=kwargs)<EOL>subkey = mapping.subkey<EOL>if x is not None:<EOL><INDENT>mapping = self._from_x[x].get(subkey, mapping).merge(x=x)<EOL><DEDENT>if y is not None:<EOL><INDENT>mapping = self._from_y[y].get(subkey, mapping).merge(y=y)<EOL><DEDENT>return mapping<EOL>

Helper which retrieves mapping info from forward/inverse dicts.


def _reduce_jacobian_det_over_event(<EOL>self, shape_tensor, ildj, min_event_ndims, event_ndims):

<EOL>rank = tf.size(input=shape_tensor)<EOL>shape_tensor = shape_tensor[rank - event_ndims:rank - min_event_ndims]<EOL>ones = tf.ones(shape_tensor, ildj.dtype)<EOL>reduced_ildj = tf.reduce_sum(<EOL>input_tensor=ones * ildj,<EOL>axis=self._get_event_reduce_dims(min_event_ndims, event_ndims))<EOL>return reduced_ildj<EOL>

Reduce jacobian over event_ndims - min_event_ndims.


def _get_event_reduce_dims(self, min_event_ndims, event_ndims):

event_ndims_ = self._maybe_get_static_event_ndims(event_ndims)<EOL>if event_ndims_ is not None:<EOL><INDENT>return [-index for index in range(<NUM_LIT:1>, event_ndims_ - min_event_ndims + <NUM_LIT:1>)]<EOL><DEDENT>else:<EOL><INDENT>reduce_ndims = event_ndims - min_event_ndims<EOL>return tf.range(-reduce_ndims, <NUM_LIT:0>)<EOL><DEDENT>

Compute the reduction dimensions given event_ndims.


def _check_valid_event_ndims(self, min_event_ndims, event_ndims):

event_ndims = tf.convert_to_tensor(value=event_ndims, name="<STR_LIT>")<EOL>event_ndims_ = tf.get_static_value(event_ndims)<EOL>assertions = []<EOL>if not dtype_util.is_integer(event_ndims.dtype):<EOL><INDENT>raise ValueError("<STR_LIT>".format(<EOL>event_ndims.dtype))<EOL><DEDENT>if event_ndims_ is not None:<EOL><INDENT>if tensorshape_util.rank(event_ndims.shape) != <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>".format(<EOL>event_ndims.shape))<EOL><DEDENT>if min_event_ndims > event_ndims_:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(event_ndims_,<EOL>min_event_ndims))<EOL><DEDENT><DEDENT>elif self.validate_args:<EOL><INDENT>assertions += [<EOL>assert_util.assert_greater_equal(event_ndims, min_event_ndims)<EOL>]<EOL><DEDENT>if tensorshape_util.is_fully_defined(event_ndims.shape):<EOL><INDENT>if tensorshape_util.rank(event_ndims.shape) != <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>".format(<EOL>tensorshape_util.rank(event_ndims.shape)))<EOL><DEDENT><DEDENT>elif self.validate_args:<EOL><INDENT>assertions += [<EOL>assert_util.assert_rank(event_ndims, <NUM_LIT:0>, message="<STR_LIT>")<EOL>]<EOL><DEDENT>return assertions<EOL>

Check whether event_ndims is atleast min_event_ndims.


event_ndims_ = distribution_util.maybe_get_static_value(event_ndims)<EOL>if isinstance(event_ndims_, (np.generic, np.ndarray)):<EOL><INDENT>if event_ndims_.dtype not in (np.int32, np.int64):<EOL><INDENT>raise ValueError("<STR_LIT>".format(<EOL>event_ndims_.dtype))<EOL><DEDENT>if isinstance(event_ndims_, np.ndarray) and len(event_ndims_.shape):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>".format(event_ndims_))<EOL><DEDENT>event_ndims_ = int(event_ndims_)<EOL><DEDENT>return event_ndims_<EOL>

Helper which returns tries to return an integer static value.


return None if x is None else tf.convert_to_tensor(<EOL>value=x, name=name, dtype=dtype)<EOL>

Convenience to convert to `Tensor` or leave as `None`.


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>if scale_perturb_diag is not None and scale_perturb_factor is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>self._is_only_identity_multiplier = (scale_tril is None and<EOL>scale_diag is None and<EOL>scale_perturb_factor is None)<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>if dtype is None:<EOL><INDENT>dtype = dtype_util.common_dtype([<EOL>shift, scale_identity_multiplier, scale_diag, scale_tril,<EOL>scale_perturb_diag, scale_perturb_factor<EOL>], tf.float32)<EOL><DEDENT>if shift is not None:<EOL><INDENT>shift = tf.convert_to_tensor(value=shift, name="<STR_LIT>", dtype=dtype)<EOL><DEDENT>self._shift = shift<EOL>if (self._is_only_identity_multiplier and<EOL>scale_identity_multiplier is None):<EOL><INDENT>scale_identity_multiplier = tf.convert_to_tensor(value=<NUM_LIT:1.>, dtype=dtype)<EOL><DEDENT>scale = self._create_scale_operator(<EOL>identity_multiplier=scale_identity_multiplier,<EOL>diag=scale_diag,<EOL>tril=scale_tril,<EOL>perturb_diag=scale_perturb_diag,<EOL>perturb_factor=scale_perturb_factor,<EOL>shift=shift,<EOL>validate_args=validate_args,<EOL>dtype=dtype)<EOL>if scale is not None and not self._is_only_identity_multiplier:<EOL><INDENT>if (shift is not None and<EOL>not dtype_util.base_equal(shift.dtype, scale.dtype)):<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>".format(<EOL>shift.dtype, scale.dtype))<EOL><DEDENT><DEDENT>self._scale = scale<EOL>self._adjoint = adjoint<EOL>super(Affine, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:1>,<EOL>graph_parents=([self._scale] if tf.is_tensor(<EOL>self._scale) else self._scale.graph_parents +<EOL>[self._shift] if self._shift is not None else []),<EOL>is_constant_jacobian=True,<EOL>dtype=dtype,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Instantiates the `Affine` bijector.

        This `Bijector` is initialized with `shift` `Tensor` and `scale` arguments,
        giving the forward operation:

        ```none
        Y = g(X) = scale @ X + shift
        ```

        where the `scale` term is logically equivalent to:

        ```python
        scale = (
          scale_identity_multiplier * tf.diag(tf.ones(d)) +
          tf.diag(scale_diag) +
          scale_tril +
          scale_perturb_factor @ diag(scale_perturb_diag) @
            tf.transpose([scale_perturb_factor])
        )
        ```

        If none of `scale_identity_multiplier`, `scale_diag`, or `scale_tril` are
        specified then `scale += IdentityMatrix`. Otherwise specifying a
        `scale` argument has the semantics of `scale += Expand(arg)`, i.e.,
        `scale_diag != None` means `scale += tf.diag(scale_diag)`.

        Args:
          shift: Floating-point `Tensor`. If this is set to `None`, no shift is
            applied.
          scale_identity_multiplier: floating point rank 0 `Tensor` representing a
            scaling done to the identity matrix.
            When `scale_identity_multiplier = scale_diag = scale_tril = None` then
            `scale += IdentityMatrix`. Otherwise no scaled-identity-matrix is added
            to `scale`.
          scale_diag: Floating-point `Tensor` representing the diagonal matrix.
            `scale_diag` has shape `[N1, N2, ...  k]`, which represents a k x k
            diagonal matrix.
            When `None` no diagonal term is added to `scale`.
          scale_tril: Floating-point `Tensor` representing the lower triangular
            matrix. `scale_tril` has shape `[N1, N2, ...  k, k]`, which represents a
            k x k lower triangular matrix.
            When `None` no `scale_tril` term is added to `scale`.
            The upper triangular elements above the diagonal are ignored.
          scale_perturb_factor: Floating-point `Tensor` representing factor matrix
            with last two dimensions of shape `(k, r)`. When `None`, no rank-r
            update is added to `scale`.
          scale_perturb_diag: Floating-point `Tensor` representing the diagonal
            matrix. `scale_perturb_diag` has shape `[N1, N2, ...  r]`, which
            represents an `r x r` diagonal matrix. When `None` low rank updates will
            take the form `scale_perturb_factor * scale_perturb_factor.T`.
          adjoint: Python `bool` indicating whether to use the `scale` matrix as
            specified or its adjoint.
            Default value: `False`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.
          dtype: `tf.DType` to prefer when converting args to `Tensor`s. Else, we
            fall back to a common dtype inferred from the args, finally falling back
            to float32.

        Raises:
          ValueError: if `perturb_diag` is specified but not `perturb_factor`.
          TypeError: if `shift` has different `dtype` from `scale` arguments.


def _create_scale_operator(self, identity_multiplier, diag, tril,<EOL>perturb_diag, perturb_factor, shift, validate_args,<EOL>dtype):

identity_multiplier = _as_tensor(identity_multiplier, "<STR_LIT>",<EOL>dtype)<EOL>diag = _as_tensor(diag, "<STR_LIT>", dtype)<EOL>tril = _as_tensor(tril, "<STR_LIT>", dtype)<EOL>perturb_diag = _as_tensor(perturb_diag, "<STR_LIT>", dtype)<EOL>perturb_factor = _as_tensor(perturb_factor, "<STR_LIT>", dtype)<EOL>shape_hint = None<EOL>if perturb_factor is not None:<EOL><INDENT>shape_hint = distribution_util.dimension_size(perturb_factor, axis=-<NUM_LIT:2>)<EOL><DEDENT>if self._is_only_identity_multiplier:<EOL><INDENT>if validate_args:<EOL><INDENT>return distribution_util.with_dependencies([<EOL>assert_util.assert_none_equal(<EOL>identity_multiplier, tf.zeros([], identity_multiplier.dtype),<EOL>["<STR_LIT>"])<EOL>], identity_multiplier)<EOL><DEDENT>return identity_multiplier<EOL><DEDENT>scale = distribution_util.make_tril_scale(<EOL>loc=shift,<EOL>scale_tril=tril,<EOL>scale_diag=diag,<EOL>scale_identity_multiplier=identity_multiplier,<EOL>validate_args=validate_args,<EOL>assert_positive=False,<EOL>shape_hint=shape_hint)<EOL>if perturb_factor is not None:<EOL><INDENT>return tf.linalg.LinearOperatorLowRankUpdate(<EOL>scale,<EOL>u=perturb_factor,<EOL>diag_update=perturb_diag,<EOL>is_diag_update_positive=perturb_diag is None,<EOL>is_non_singular=True,  <EOL>is_self_adjoint=True,<EOL>is_positive_definite=True,<EOL>is_square=True)<EOL><DEDENT>return scale<EOL>

Construct `scale` from various components.

        Args:
          identity_multiplier: floating point rank 0 `Tensor` representing a scaling
            done to the identity matrix.
          diag: Floating-point `Tensor` representing the diagonal matrix.`diag` has
            shape `[N1, N2, ...  k]`, which represents a k x k diagonal matrix.
          tril: Floating-point `Tensor` representing the lower triangular matrix.
           `tril` has shape `[N1, N2, ...  k, k]`, which represents a k x k lower
           triangular matrix.
          perturb_diag: Floating-point `Tensor` representing the diagonal matrix of
            the low rank update.
          perturb_factor: Floating-point `Tensor` representing factor matrix.
          shift: Floating-point `Tensor` representing `shift in `scale @ X + shift`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          dtype: `DType` for arg `Tensor` conversions.

        Returns:
          scale. In the case of scaling by a constant, scale is a
          floating point `Tensor`. Otherwise, scale is a `LinearOperator`.

        Raises:
          ValueError: if all of `tril`, `diag` and `identity_multiplier` are `None`.


return self._shift<EOL>

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


return self._scale<EOL>

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


return self._adjoint<EOL>

`bool` indicating `scale` should be used as conjugate transpose.


def __init__(self, validate_args=False, name="<STR_LIT>"):

self._graph_parents = []<EOL>self._name = name<EOL>super(CholeskyOuterProduct, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:2>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `CholeskyOuterProduct` bijector.

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


if tensorshape_util.rank(x.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(x.shape) == <NUM_LIT:1>:<EOL><INDENT>x = x[tf.newaxis, :]<EOL><DEDENT>return x<EOL><DEDENT>shape = tf.shape(input=x)<EOL>maybe_expanded_shape = tf.concat([<EOL>shape[:-<NUM_LIT:1>],<EOL>distribution_util.pick_vector(<EOL>tf.equal(tf.rank(x), <NUM_LIT:1>), [<NUM_LIT:1>], np.array([], dtype=np.int32)),<EOL>shape[-<NUM_LIT:1>:],<EOL>], <NUM_LIT:0>)<EOL>return tf.reshape(x, maybe_expanded_shape)<EOL>

Ensures non-scalar input has at least one column.

        Example:
          If `x = [1, 2, 3]` then the output is `[[1], [2], [3]]`.

          If `x = [[1, 2, 3], [4, 5, 6]]` then the output is unchanged.

          If `x = 1` then the output is unchanged.

        Args:
          x: `Tensor`.

        Returns:
          columnar_x: `Tensor` with at least two dimensions.


with tf.compat.v2.name_scope(name or "<STR_LIT>"):<EOL><INDENT>rescale = tf.sqrt(variance + variance_epsilon)<EOL>if scale is not None:<EOL><INDENT>rescale /= scale<EOL><DEDENT>batch_unnormalized = x * rescale + (<EOL>mean - offset * rescale if offset is not None else mean)<EOL>return batch_unnormalized<EOL><DEDENT>

r"""Inverse of tf.nn.batch_normalization.

    Args:
      x: Input `Tensor` of arbitrary dimensionality.
      mean: A mean `Tensor`.
      variance: A variance `Tensor`.
      offset: An offset `Tensor`, often denoted `beta` in equations, or
        None. If present, will be added to the normalized tensor.
      scale: A scale `Tensor`, often denoted `gamma` in equations, or
        `None`. If present, the scale is applied to the normalized tensor.
      variance_epsilon: A small `float` added to the minibatch `variance` to
        prevent dividing by zero.
      name: A name for this operation (optional).

    Returns:
      batch_unnormalized: The de-normalized, de-scaled, de-offset `Tensor`.


<EOL>g_constraint = lambda x: tf.nn.relu(x) + <NUM_LIT><EOL>self.batchnorm = batchnorm_layer or tf.keras.layers.BatchNormalization(<EOL>gamma_constraint=g_constraint)<EOL>self._validate_bn_layer(self.batchnorm)<EOL>self._training = training<EOL>if isinstance(self.batchnorm.axis, int):<EOL><INDENT>forward_min_event_ndims = <NUM_LIT:1><EOL><DEDENT>else:<EOL><INDENT>forward_min_event_ndims = len(self.batchnorm.axis)<EOL><DEDENT>super(BatchNormalization, self).__init__(<EOL>forward_min_event_ndims=forward_min_event_ndims,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `BatchNormalization` bijector.

        Args:
          batchnorm_layer: `tf.layers.BatchNormalization` layer object. If `None`,
            defaults to
            `tf.layers.BatchNormalization(gamma_constraint=tf.nn.relu(x) + 1e-6)`.
            This ensures positivity of the scale variable.

          training: If True, updates running-average statistics during call to
            `inverse()`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.
        Raises:
          ValueError: If bn_layer is not an instance of
            `tf.layers.BatchNormalization`, or if it is specified with `renorm=True`
            or a virtual batch size.


if (not isinstance(layer, tf.keras.layers.BatchNormalization) and<EOL>not isinstance(layer, tf.compat.v1.layers.BatchNormalization)):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>if layer.renorm:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if layer.virtual_batch_size:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>

Check for valid BatchNormalization layer.

        Args:
          layer: Instance of `tf.layers.BatchNormalization`.
        Raises:
          ValueError: If batchnorm_layer argument is not an instance of
          `tf.layers.BatchNormalization`, or if `batchnorm_layer.renorm=True` or
          if `batchnorm_layer.virtual_batch_size` is specified.


assertions = []<EOL>message = '<STR_LIT>'<EOL>if tensorshape_util.rank(map_values.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(map_values.shape) != <NUM_LIT:1>:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(assert_util.assert_rank(map_values, <NUM_LIT:1>, message=message))<EOL><DEDENT>message = '<STR_LIT>'<EOL>if tensorshape_util.num_elements(map_values.shape) is not None:<EOL><INDENT>if tensorshape_util.num_elements(map_values.shape) == <NUM_LIT:0>:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>assert_util.assert_greater(<EOL>tf.size(input=map_values), <NUM_LIT:0>, message=message))<EOL><DEDENT>if validate_args:<EOL><INDENT>assertions.append(<EOL>assert_util.assert_equal(<EOL>tf.math.is_strictly_increasing(map_values),<EOL>True,<EOL>message='<STR_LIT>'))<EOL><DEDENT>return assertions<EOL>

Validate `map_values` if `validate_args`==True.


with tf.name_scope(name):<EOL><INDENT>map_values = tf.convert_to_tensor(value=map_values, name='<STR_LIT>')<EOL>assertions = _maybe_check_valid_map_values(map_values, validate_args)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>map_values = tf.identity(map_values)<EOL><DEDENT><DEDENT>self._map_values = map_values<EOL>super(CategoricalToDiscrete, self).__init__(<EOL>graph_parents=[map_values],<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>is_constant_jacobian=True,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Instantiates `CategoricalToDiscrete` bijector.

        Args:
          map_values: 1D numerical tensor of discrete values to map to, sorted in
            strictly increasing order.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


Instantiates the `TransformDiagonal` bijector.

        Args:
          diag_bijector: `Bijector` instance used to transform the diagonal.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


def _get_fldj_numerical(self, bijector, x, event_ndims,<EOL>eps=<NUM_LIT>,<EOL>input_to_vector=tfb.Identity,<EOL>output_to_vector=tfb.Identity):

x_vector = input_to_vector.forward(x)                         <EOL>n = tf.shape(input=x_vector)[-<NUM_LIT:1>]<EOL>x_plus_eps_vector = (<EOL>x_vector[..., tf.newaxis, :] +<EOL>eps * tf.eye(n, dtype=x_vector.dtype))                    <EOL>x_plus_eps = input_to_vector.inverse(x_plus_eps_vector)       <EOL>f_x_plus_eps = bijector.forward(x_plus_eps)                   <EOL>f_x_plus_eps_vector = output_to_vector.forward(f_x_plus_eps)  <EOL>f_x = bijector.forward(x)                                     <EOL>f_x_vector = output_to_vector.forward(f_x)                    <EOL>jacobian_numerical = (f_x_plus_eps_vector -<EOL>f_x_vector[..., tf.newaxis, :]) / eps<EOL>return (<EOL>tf.math.log(tf.abs(tf.linalg.det(jacobian_numerical))) +<EOL>input_to_vector.forward_log_det_jacobian(x, event_ndims=event_ndims) -<EOL>output_to_vector.forward_log_det_jacobian(f_x, event_ndims=event_ndims)<EOL>)<EOL>

Numerically approximate the forward log det Jacobian of a bijector.

        Args:
          bijector: the bijector whose Jacobian we wish to approximate
          x: the value for which we want to approximate the Jacobian
          event_ndims: number of dimensions in an event
          eps: epsilon to add when forming (f(x+eps)-f(x)) / eps
          input_to_vector: a bijector that maps the input value to a vector
          output_to_vector: a bijector that maps the output value to a vector

        Returns:
          A numerical approximation to the log det Jacobian of bijector.forward
          evaluated at x.


return -<NUM_LIT> * np.log1p(-np.abs(y))<EOL>

Inverse log det jacobian, before being reduced.


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>self._scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>")<EOL>self._concentration = tf.convert_to_tensor(<EOL>value=concentration, name="<STR_LIT>")<EOL>dtype_util.assert_same_float_dtype([self._scale, self._concentration])<EOL>if validate_args:<EOL><INDENT>self._scale = distribution_util.with_dependencies([<EOL>assert_util.assert_positive(<EOL>self._scale, message="<STR_LIT>")<EOL>], self._scale)<EOL>self._concentration = distribution_util.with_dependencies([<EOL>assert_util.assert_positive(<EOL>self._concentration,<EOL>message="<STR_LIT>")<EOL>], self._concentration)<EOL><DEDENT><DEDENT>super(Weibull, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `Weibull` bijector.

        Args:
          scale: Positive Float-type `Tensor` that is the same dtype and is
            broadcastable with `concentration`.
            This is `l` in `Y = g(X) = 1 - exp((-x / l) ** k)`.
          concentration: Positive Float-type `Tensor` that is the same dtype and is
            broadcastable with `scale`.
            This is `k` in `Y = g(X) = 1 - exp((-x / l) ** k)`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


return self._scale<EOL>


return self._concentration<EOL>


<EOL>super(Exp, self).__init__(<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `Exp` bijector.

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


c = scale_identity_multiplier<EOL>d1 = scale_diag<EOL>tril = scale_tril<EOL>v = scale_perturb_factor<EOL>d2 = scale_perturb_diag<EOL>if v is None and d2 is not None:<EOL><INDENT>return None<EOL><DEDENT>if c is None and d1 is None and tril is None:<EOL><INDENT>c = <NUM_LIT:1.><EOL><DEDENT>matrix = np.float32(<NUM_LIT:0.>)<EOL>if c is not None:<EOL><INDENT>matrix += c * self._matrix_diag(np.ones_like(x))<EOL><DEDENT>if d1 is not None:<EOL><INDENT>matrix += self._matrix_diag(np.array(d1, dtype=np.float32))<EOL><DEDENT>if tril is not None:<EOL><INDENT>matrix += np.array(tril, dtype=np.float32)<EOL><DEDENT>if v is not None:<EOL><INDENT>v = np.array(v, dtype=np.float32)<EOL>if v.ndim < <NUM_LIT:2>:<EOL><INDENT>vt = v.T<EOL><DEDENT>else:<EOL><INDENT>vt = np.swapaxes(v, axis1=v.ndim - <NUM_LIT:2>, axis2=v.ndim - <NUM_LIT:1>)<EOL><DEDENT>if d2 is not None:<EOL><INDENT>d2 = self._matrix_diag(np.array(d2, dtype=np.float32))<EOL>right = np.matmul(d2, vt)<EOL><DEDENT>else:<EOL><INDENT>right = vt<EOL><DEDENT>matrix += np.matmul(v, right)<EOL><DEDENT>return matrix<EOL>

Create a scale matrix. Return None if it can not be created.


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>


sqrt_eps = np.sqrt(np.finfo(dtype_util.as_numpy_dtype(x.dtype)).eps)<EOL>return tf.where(<EOL>tf.abs(x) * sqrt_eps <= <NUM_LIT:1.>,<EOL>tf.sqrt(x**<NUM_LIT> + <NUM_LIT:1.>),<EOL>tf.abs(x))<EOL>

Implementation of `sqrt(1 + x**2)` which is stable despite large `x`.


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>tailweight = <NUM_LIT:1.> if tailweight is None else tailweight<EOL>skewness = <NUM_LIT:0.> if skewness is None else skewness<EOL>self._skewness = tf.convert_to_tensor(value=skewness, name="<STR_LIT>")<EOL>self._tailweight = tf.convert_to_tensor(<EOL>value=tailweight, name="<STR_LIT>", dtype=self._skewness.dtype)<EOL>dtype_util.assert_same_float_dtype([self._skewness, self._tailweight])<EOL>if validate_args:<EOL><INDENT>self._tailweight = distribution_util.with_dependencies([<EOL>assert_util.assert_positive(<EOL>self._tailweight,<EOL>message="<STR_LIT>")<EOL>], self._tailweight)<EOL><DEDENT><DEDENT>super(SinhArcsinh, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `SinhArcsinh` bijector.

        Args:
          skewness:  Skewness parameter.  Float-type `Tensor`.  Default is `0`
            of type `float32`.
          tailweight:  Tailweight parameter.  Positive `Tensor` of same `dtype` as
            `skewness` and broadcastable `shape`.  Default is `1` of type `float32`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


return self._skewness<EOL>

The `skewness` in: `Y  = Sinh((Arcsinh(X) + skewness) * tailweight)`.


return self._tailweight<EOL>

The `tailweight` in: `Y = Sinh((Arcsinh(X) + skewness) * tailweight)`.


self._name = name<EOL>super(Reciprocal, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `Reciprocal`.

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


with tf.compat.v2.name_scope(name or "<STR_LIT>"):<EOL><INDENT>def _fn(x, output_units, **condition_kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>if condition_kwargs:<EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>")<EOL><DEDENT>if tensorshape_util.rank(x.shape) == <NUM_LIT:1>:<EOL><INDENT>x = x[tf.newaxis, ...]<EOL>reshape_output = lambda x: x[<NUM_LIT:0>]<EOL><DEDENT>else:<EOL><INDENT>reshape_output = lambda x: x<EOL><DEDENT>for units in hidden_layers:<EOL><INDENT>x = tf.compat.v1.layers.dense(<EOL>inputs=x,<EOL>units=units,<EOL>activation=activation,<EOL>*args,  <EOL>**kwargs)<EOL><DEDENT>x = tf.compat.v1.layers.dense(<EOL>inputs=x,<EOL>units=(<NUM_LIT:1> if shift_only else <NUM_LIT:2>) * output_units,<EOL>activation=None,<EOL>*args,  <EOL>**kwargs)<EOL>if shift_only:<EOL><INDENT>return reshape_output(x), None<EOL><DEDENT>shift, log_scale = tf.split(x, <NUM_LIT:2>, axis=-<NUM_LIT:1>)<EOL>return reshape_output(shift), reshape_output(log_scale)<EOL><DEDENT>return tf.compat.v1.make_template("<STR_LIT>", _fn)<EOL><DEDENT>

Build a scale-and-shift function using a multi-layer neural network.

    This will be wrapped in a make_template to ensure the variables are only
    created once. It takes the `d`-dimensional input x[0:d] and returns the `D-d`
    dimensional outputs `loc` ("mu") and `log_scale` ("alpha").

    The default template does not support conditioning and will raise an
    exception if `condition_kwargs` are passed to it. To use conditioning in
    real nvp bijector, implement a conditioned shift/scale template that
    handles the `condition_kwargs`.

    Arguments:
      hidden_layers: Python `list`-like of non-negative integer, scalars
        indicating the number of units in each hidden layer. Default: `[512, 512].
      shift_only: Python `bool` indicating if only the `shift` term shall be
        computed (i.e. NICE bijector). Default: `False`.
      activation: Activation function (callable). Explicitly setting to `None`
        implies a linear activation.
      name: A name for ops managed by this function. Default:
        "real_nvp_default_template".
      *args: `tf.layers.dense` arguments.
      **kwargs: `tf.layers.dense` keyword arguments.

    Returns:
      shift: `Float`-like `Tensor` of shift terms ("mu" in
        [Papamakarios et al.  (2016)][1]).
      log_scale: `Float`-like `Tensor` of log(scale) terms ("alpha" in
        [Papamakarios et al. (2016)][1]).

    Raises:
      NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
        graph execution, or if `condition_kwargs` is not empty.

    #### References

    [1]: George Papamakarios, Theo Pavlakou, and Iain Murray. Masked
         Autoregressive Flow for Density Estimation. In _Neural Information
         Processing Systems_, 2017. https://arxiv.org/abs/1705.07057


name = name or "<STR_LIT>"<EOL>if num_masked <= <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>self._num_masked = num_masked<EOL>self._input_depth = None<EOL>self._shift_and_log_scale_fn = shift_and_log_scale_fn<EOL>super(RealNVP, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:1>,<EOL>is_constant_jacobian=is_constant_jacobian,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Creates the Real NVP or NICE bijector.

        Args:
          num_masked: Python `int` indicating that the first `d` units of the event
            should be masked. Must be in the closed interval `[1, D-1]`, where `D`
            is the event size of the base distribution.
          shift_and_log_scale_fn: Python `callable` which computes `shift` and
            `log_scale` from both the forward domain (`x`) and the inverse domain
            (`y`). Calculation must respect the "autoregressive property" (see class
            docstring). Suggested default
            `masked_autoregressive_default_template(hidden_layers=...)`.
            Typically the function contains `tf.Variables` and is wrapped using
            `tf.make_template`. Returning `None` for either (both) `shift`,
            `log_scale` is equivalent to (but more efficient than) returning zero.
          is_constant_jacobian: Python `bool`. Default: `False`. When `True` the
            implementation assumes `log_scale` does not depend on the forward domain
            (`x`) or inverse domain (`y`) values. (No validation is made;
            `is_constant_jacobian=False` is always safe but possibly computationally
            inefficient.)
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str`, name given to ops managed by this object.

        Raises:
          ValueError: If num_masked < 1.


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>super(NormalCDF, self).__init__(<EOL>validate_args=validate_args,<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>name=name)<EOL>

Instantiates the `NormalCDF` bijector.

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>self._loc = tf.convert_to_tensor(value=loc, name="<STR_LIT>")<EOL>self._scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>")<EOL>dtype_util.assert_same_float_dtype([self._loc, self._scale])<EOL>if validate_args:<EOL><INDENT>self._scale = distribution_util.with_dependencies([<EOL>assert_util.assert_positive(<EOL>self._scale, message="<STR_LIT>")<EOL>], self._scale)<EOL><DEDENT><DEDENT>super(Gumbel, self).__init__(<EOL>validate_args=validate_args,<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>name=name)<EOL>

Instantiates the `Gumbel` bijector.

        Args:
          loc: Float-like `Tensor` that is the same dtype and is
            broadcastable with `scale`.
            This is `loc` in `Y = g(X) = exp(-exp(-(X - loc) / scale))`.
          scale: Positive Float-like `Tensor` that is the same dtype and is
            broadcastable with `loc`.
            This is `scale` in `Y = g(X) = exp(-exp(-(X - loc) / scale))`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


return self._loc<EOL>

The `loc` in `Y = g(X) = exp(-exp(-(X - loc) / scale))`.


return self._scale<EOL>

This is `scale` in `Y = g(X) = exp(-exp(-(X - loc) / scale))`.


with tf.name_scope(name or '<STR_LIT>') as name:<EOL><INDENT>self._lower_upper = tf.convert_to_tensor(<EOL>value=lower_upper, dtype_hint=tf.float32, name='<STR_LIT>')<EOL>self._permutation = tf.convert_to_tensor(<EOL>value=permutation, dtype_hint=tf.int32, name='<STR_LIT>')<EOL><DEDENT>super(MatvecLU, self).__init__(<EOL>is_constant_jacobian=True,<EOL>forward_min_event_ndims=<NUM_LIT:1>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Creates the MatvecLU bijector.

        Args:
          lower_upper: The LU factorization as returned by `tf.linalg.lu`.
          permutation: The LU factorization permutation as returned by
            `tf.linalg.lu`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
            Default value: `False`.
          name: Python `str` name given to ops managed by this object.
            Default value: `None` (i.e., "MatvecLU").

        Raises:
          ValueError: If both/neither `channels` and `lower_upper`/`permutation` are
            specified.


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>self._shift = shift<EOL>self._scale = scale<EOL>if self._shift is not None:<EOL><INDENT>self._shift = tf.convert_to_tensor(value=shift, name="<STR_LIT>")<EOL><DEDENT>if self._scale is not None:<EOL><INDENT>self._scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>")<EOL>if validate_args:<EOL><INDENT>self._scale = distribution_util.with_dependencies([<EOL>assert_util.assert_none_equal(<EOL>self._scale, tf.zeros([], dtype=self._scale.dtype))<EOL>], self._scale)<EOL><DEDENT><DEDENT>dtype = dtype_util.common_dtype([self._shift, self._scale])<EOL>super(AffineScalar, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>is_constant_jacobian=True,<EOL>validate_args=validate_args,<EOL>dtype=dtype,<EOL>name=name)<EOL><DEDENT>

Instantiates the `AffineScalar` bijector.

        This `Bijector` is initialized with `shift` `Tensor` and `scale` arguments,
        giving the forward operation:

        ```none
        Y = g(X) = scale * X + shift
        ```

        if `scale` is not specified, then the bijector has the semantics of
        `scale = 1.`. Similarly, if `shift` is not specified, then the bijector
        has the semantics of `shift = 0.`.

        Args:
          shift: Floating-point `Tensor`. If this is set to `None`, no shift is
            applied.
          scale: Floating-point `Tensor`. If this is set to `None`, no scale is
            applied.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


return self._shift<EOL>

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


return self._scale<EOL>

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


def _gen_slices(num_blocks, n_in, n_out, mask_type=MASK_EXCLUSIVE):

<EOL>slices = []<EOL>col = <NUM_LIT:0><EOL>d_in = n_in // num_blocks<EOL>d_out = n_out // num_blocks<EOL>row = d_out if mask_type == MASK_EXCLUSIVE else <NUM_LIT:0><EOL>for _ in range(num_blocks):<EOL><INDENT>row_slice = slice(row, None)<EOL>col_slice = slice(col, col + d_in)<EOL>slices.append([row_slice, col_slice])<EOL>col += d_in<EOL>row += d_out<EOL><DEDENT>return slices<EOL>

Generate the slices for building an autoregressive mask.


<EOL>mask = np.zeros([n_out, n_in], dtype=dtype.as_numpy_dtype())<EOL>slices = _gen_slices(num_blocks, n_in, n_out, mask_type=mask_type)<EOL>for [row_slice, col_slice] in slices:<EOL><INDENT>mask[row_slice, col_slice] = <NUM_LIT:1><EOL><DEDENT>return mask<EOL>

Generate the mask for building an autoregressive dense layer.


<EOL>input_depth = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(inputs.shape, <NUM_LIT:1>)[-<NUM_LIT:1>])<EOL>if input_depth is None:<EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>")<EOL><DEDENT>mask = _gen_mask(num_blocks, input_depth, units,<EOL>MASK_EXCLUSIVE if exclusive else MASK_INCLUSIVE).T<EOL>if kernel_initializer is None:<EOL><INDENT>kernel_initializer = tf.compat.v1.glorot_normal_initializer()<EOL><DEDENT>def masked_initializer(shape, dtype=None, partition_info=None):<EOL><INDENT>return mask * kernel_initializer(shape, dtype, partition_info)<EOL><DEDENT>with tf.compat.v2.name_scope(name or "<STR_LIT>"):<EOL><INDENT>layer = tf.compat.v1.layers.Dense(<EOL>units,<EOL>kernel_initializer=masked_initializer,<EOL>kernel_constraint=lambda x: mask * x,<EOL>name=name,<EOL>dtype=dtype_util.base_dtype(inputs.dtype),<EOL>_scope=name,<EOL>_reuse=reuse,<EOL>*args,  <EOL>**kwargs)<EOL>return layer.apply(inputs)<EOL><DEDENT>

A autoregressively masked dense layer. Analogous to `tf.layers.dense`.

    See [Germain et al. (2015)][1] for detailed explanation.

    Arguments:
      inputs: Tensor input.
      units: Python `int` scalar representing the dimensionality of the output
        space.
      num_blocks: Python `int` scalar representing the number of blocks for the
        MADE masks.
      exclusive: Python `bool` scalar representing whether to zero the diagonal of
        the mask, used for the first layer of a MADE.
      kernel_initializer: Initializer function for the weight matrix.
        If `None` (default), weights are initialized using the
        `tf.glorot_random_initializer`.
      reuse: Python `bool` scalar representing whether to reuse the weights of a
        previous layer by the same name.
      name: Python `str` used to describe ops managed by this function.
      *args: `tf.layers.dense` arguments.
      **kwargs: `tf.layers.dense` keyword arguments.

    Returns:
      Output tensor.

    Raises:
      NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
        graph execution.

    #### References

    [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
         Masked Autoencoder for Distribution Estimation. In _International
         Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509


name = name or "<STR_LIT>"<EOL>with tf.compat.v2.name_scope(name):<EOL><INDENT>def _fn(x):<EOL><INDENT>"""<STR_LIT>"""<EOL>input_depth = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(x.shape, <NUM_LIT:1>)[-<NUM_LIT:1>])<EOL>if input_depth is None:<EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>")<EOL><DEDENT>input_shape = (<EOL>np.int32(tensorshape_util.as_list(x.shape))<EOL>if tensorshape_util.is_fully_defined(x.shape) else tf.shape(input=x))<EOL>if tensorshape_util.rank(x.shape) == <NUM_LIT:1>:<EOL><INDENT>x = x[tf.newaxis, ...]<EOL><DEDENT>for i, units in enumerate(hidden_layers):<EOL><INDENT>x = masked_dense(<EOL>inputs=x,<EOL>units=units,<EOL>num_blocks=input_depth,<EOL>exclusive=True if i == <NUM_LIT:0> else False,<EOL>activation=activation,<EOL>*args,  <EOL>**kwargs)<EOL><DEDENT>x = masked_dense(<EOL>inputs=x,<EOL>units=(<NUM_LIT:1> if shift_only else <NUM_LIT:2>) * input_depth,<EOL>num_blocks=input_depth,<EOL>activation=None,<EOL>*args,  <EOL>**kwargs)<EOL>if shift_only:<EOL><INDENT>x = tf.reshape(x, shape=input_shape)<EOL>return x, None<EOL><DEDENT>x = tf.reshape(x, shape=tf.concat([input_shape, [<NUM_LIT:2>]], axis=<NUM_LIT:0>))<EOL>shift, log_scale = tf.unstack(x, num=<NUM_LIT:2>, axis=-<NUM_LIT:1>)<EOL>which_clip = (<EOL>tf.clip_by_value<EOL>if log_scale_clip_gradient else clip_by_value_preserve_gradient)<EOL>log_scale = which_clip(log_scale, log_scale_min_clip, log_scale_max_clip)<EOL>return shift, log_scale<EOL><DEDENT>return tf.compat.v1.make_template(name, _fn)<EOL><DEDENT>

Build the Masked Autoregressive Density Estimator (Germain et al., 2015).

    This will be wrapped in a make_template to ensure the variables are only
    created once. It takes the input and returns the `loc` ("mu" in [Germain et
    al. (2015)][1]) and `log_scale` ("alpha" in [Germain et al. (2015)][1]) from
    the MADE network.

    Warning: This function uses `masked_dense` to create randomly initialized
    `tf.Variables`. It is presumed that these will be fit, just as you would any
    other neural architecture which uses `tf.layers.dense`.

    #### About Hidden Layers

    Each element of `hidden_layers` should be greater than the `input_depth`
    (i.e., `input_depth = tf.shape(input)[-1]` where `input` is the input to the
    neural network). This is necessary to ensure the autoregressivity property.

    #### About Clipping

    This function also optionally clips the `log_scale` (but possibly not its
    gradient). This is useful because if `log_scale` is too small/large it might
    underflow/overflow making it impossible for the `MaskedAutoregressiveFlow`
    bijector to implement a bijection. Additionally, the `log_scale_clip_gradient`
    `bool` indicates whether the gradient should also be clipped. The default does
    not clip the gradient; this is useful because it still provides gradient
    information (for fitting) yet solves the numerical stability problem. I.e.,
    `log_scale_clip_gradient = False` means
    `grad[exp(clip(x))] = grad[x] exp(clip(x))` rather than the usual
    `grad[clip(x)] exp(clip(x))`.

    Args:
      hidden_layers: Python `list`-like of non-negative integer, scalars
        indicating the number of units in each hidden layer. Default: `[512, 512].
      shift_only: Python `bool` indicating if only the `shift` term shall be
        computed. Default: `False`.
      activation: Activation function (callable). Explicitly setting to `None`
        implies a linear activation.
      log_scale_min_clip: `float`-like scalar `Tensor`, or a `Tensor` with the
        same shape as `log_scale`. The minimum value to clip by. Default: -5.
      log_scale_max_clip: `float`-like scalar `Tensor`, or a `Tensor` with the
        same shape as `log_scale`. The maximum value to clip by. Default: 3.
      log_scale_clip_gradient: Python `bool` indicating that the gradient of
        `tf.clip_by_value` should be preserved. Default: `False`.
      name: A name for ops managed by this function. Default:
        "masked_autoregressive_default_template".
      *args: `tf.layers.dense` arguments.
      **kwargs: `tf.layers.dense` keyword arguments.

    Returns:
      shift: `Float`-like `Tensor` of shift terms (the "mu" in
        [Germain et al.  (2015)][1]).
      log_scale: `Float`-like `Tensor` of log(scale) terms (the "alpha" in
        [Germain et al. (2015)][1]).

    Raises:
      NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
        graph execution.

    #### References

    [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
         Masked Autoencoder for Distribution Estimation. In _International
         Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509


try:<EOL><INDENT>return list(xs)<EOL><DEDENT>except TypeError:<EOL><INDENT>return [xs]<EOL><DEDENT>


if isinstance(input_order, six.string_types):<EOL><INDENT>if input_order == "<STR_LIT>":<EOL><INDENT>return np.arange(start=<NUM_LIT:1>, stop=input_size + <NUM_LIT:1>)<EOL><DEDENT>elif input_order == "<STR_LIT>":<EOL><INDENT>return np.arange(start=input_size, stop=<NUM_LIT:0>, step=-<NUM_LIT:1>)<EOL><DEDENT>elif input_order == "<STR_LIT>":<EOL><INDENT>ret = np.arange(start=<NUM_LIT:1>, stop=input_size + <NUM_LIT:1>)<EOL>np.random.shuffle(ret)<EOL>return ret<EOL><DEDENT><DEDENT>elif np.all(np.sort(input_order) == np.arange(<NUM_LIT:1>, input_size + <NUM_LIT:1>)):<EOL><INDENT>return np.array(input_order)<EOL><DEDENT>raise ValueError("<STR_LIT>".format(input_order))<EOL>


input_order = _create_input_order(input_size, input_order)<EOL>degrees = [input_order]<EOL>if hidden_units is None:<EOL><INDENT>hidden_units = []<EOL><DEDENT>for units in hidden_units:<EOL><INDENT>if isinstance(hidden_degrees, six.string_types):<EOL><INDENT>if hidden_degrees == "<STR_LIT>":<EOL><INDENT>degrees.append(<EOL>np.random.randint(low=min(np.min(degrees[-<NUM_LIT:1>]), input_size - <NUM_LIT:1>),<EOL>high=input_size,<EOL>size=units))<EOL><DEDENT>elif hidden_degrees == "<STR_LIT>":<EOL><INDENT>min_degree = min(np.min(degrees[-<NUM_LIT:1>]), input_size - <NUM_LIT:1>)<EOL>degrees.append(np.maximum(<EOL>min_degree,<EOL>np.ceil(np.arange(<NUM_LIT:1>, units + <NUM_LIT:1>)<EOL>* (input_size - <NUM_LIT:1>) / float(units + <NUM_LIT:1>)).astype(np.int32)))<EOL><DEDENT><DEDENT>else:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(hidden_degrees))<EOL><DEDENT><DEDENT>return degrees<EOL>

Returns a list of degree vectors, one for each input and hidden layer.

    A unit with degree d can only receive input from units with degree < d. Output
    units always have the same degree as their associated input unit.

    Args:
      input_size: Number of inputs.
      hidden_units: list with the number of hidden units per layer. It does not
        include the output layer. Each hidden unit size must be at least the size
        of length (otherwise autoregressivity is not possible).
      input_order: Order of degrees to the input units: 'random', 'left-to-right',
        'right-to-left', or an array of an explicit order. For example,
        'left-to-right' builds an autoregressive model
        p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
      hidden_degrees: Method for assigning degrees to the hidden units:
        'equal', 'random'.  If 'equal', hidden units in each layer are allocated
        equally (up to a remainder term) to each degree.  Default: 'equal'.

    Raises:
      ValueError: invalid input order.
      ValueError: invalid hidden degrees.


return [<EOL>inp[:, np.newaxis] <= out<EOL>for inp, out in zip(degrees[:-<NUM_LIT:1>], degrees[<NUM_LIT:1>:])<EOL>] + [<EOL>degrees[-<NUM_LIT:1>][:, np.newaxis] < degrees[<NUM_LIT:0>]<EOL>]<EOL>

Returns a list of binary mask matrices enforcing autoregressivity.


initializer = tf.keras.initializers.get(initializer)<EOL>def masked_initializer(shape, dtype=None, partition_info=None):<EOL><INDENT>if partition_info is None:<EOL><INDENT>x = initializer(shape, dtype)<EOL><DEDENT>else:<EOL><INDENT>x = initializer(shape, dtype, partition_info)<EOL><DEDENT>return tf.cast(mask, x.dtype) * x<EOL><DEDENT>return masked_initializer<EOL>

Returns a masked version of the given initializer.


name = name or "<STR_LIT>"<EOL>self._shift_and_log_scale_fn = shift_and_log_scale_fn<EOL>self._unroll_loop = unroll_loop<EOL>self._event_ndims = event_ndims<EOL>super(MaskedAutoregressiveFlow, self).__init__(<EOL>forward_min_event_ndims=self._event_ndims,<EOL>is_constant_jacobian=is_constant_jacobian,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Creates the MaskedAutoregressiveFlow bijector.

        Args:
          shift_and_log_scale_fn: Python `callable` which computes `shift` and
            `log_scale` from both the forward domain (`x`) and the inverse domain
            (`y`). Calculation must respect the "autoregressive property" (see class
            docstring). Suggested default
            `masked_autoregressive_default_template(hidden_layers=...)`.
            Typically the function contains `tf.Variables` and is wrapped using
            `tf.make_template`. Returning `None` for either (both) `shift`,
            `log_scale` is equivalent to (but more efficient than) returning zero.
          is_constant_jacobian: Python `bool`. Default: `False`. When `True` the
            implementation assumes `log_scale` does not depend on the forward domain
            (`x`) or inverse domain (`y`) values. (No validation is made;
            `is_constant_jacobian=False` is always safe but possibly computationally
            inefficient.)
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          unroll_loop: Python `bool` indicating whether the `tf.while_loop` in
            `_forward` should be replaced with a static for loop. Requires that
            the final dimension of `x` be known at graph construction time. Defaults
            to `False`.
          event_ndims: Python `integer`, the intrinsic dimensionality of this
            bijector. 1 corresponds to a simple vector autoregressive bijector as
            implemented by the `masked_autoregressive_default_template`, 2 might be
            useful for a 2D convolutional `shift_and_log_scale_fn` and so on.
          name: Python `str`, name given to ops managed by this object.


super(AutoregressiveLayer, self).__init__(**kwargs)<EOL>self._params = params<EOL>self._event_shape = _list(event_shape) if event_shape is not None else None<EOL>self._hidden_units = hidden_units if hidden_units is not None else []<EOL>self._input_order_param = input_order<EOL>self._hidden_degrees = hidden_degrees<EOL>self._activation = activation<EOL>self._use_bias = use_bias<EOL>self._kernel_initializer = kernel_initializer<EOL>self._validate_args = validate_args<EOL>self._kwargs = kwargs<EOL>if self._event_shape is not None:<EOL><INDENT>self._event_size = self._event_shape[-<NUM_LIT:1>]<EOL>self._event_ndims = len(self._event_shape)<EOL>if self._event_ndims != <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(event_shape))<EOL><DEDENT><DEDENT>self._input_order = None<EOL>self._masks = None<EOL>self._network = None<EOL>

Constructs the MADE layer.

        Arguments:
          params: Python integer specifying the number of parameters to output
            per input.
          event_shape: Python `list`-like of positive integers (or a single int),
            specifying the shape of the input to this layer, which is also the
            event_shape of the distribution parameterized by this layer.  Currently
            only rank-1 shapes are supported.  That is, event_shape must be a single
            integer.  If not specified, the event shape is inferred when this layer
            is first called or built.
          hidden_units: Python `list`-like of non-negative integers, specifying
            the number of units in each hidden layer.
          input_order: Order of degrees to the input units: 'random',
            'left-to-right', 'right-to-left', or an array of an explicit order. For
            example, 'left-to-right' builds an autoregressive model:
            `p(x) = p(x1) p(x2 | x1) ... p(xD | x<D)`.  Default: 'left-to-right'.
          hidden_degrees: Method for assigning degrees to the hidden units:
            'equal', 'random'.  If 'equal', hidden units in each layer are allocated
            equally (up to a remainder term) to each degree.  Default: 'equal'.
          activation: An activation function.  See `tf.keras.layers.Dense`. Default:
            `None`.
          use_bias: Whether or not the dense layers constructed in this layer
            should have a bias term.  See `tf.keras.layers.Dense`.  Default: `True`.
          kernel_initializer: Initializer for the kernel weights matrix.  Default:
            'glorot_uniform'.
          validate_args: Python `bool`, default `False`. When `True`, layer
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          **kwargs: Addition keyword arguments passed to the `tf.keras.layers.Dense`
            constructed by this layer.


if self._event_shape is None:<EOL><INDENT>self._event_shape = [tf.compat.dimension_value(input_shape[-<NUM_LIT:1>])]<EOL>self._event_size = self._event_shape[-<NUM_LIT:1>]<EOL>self._event_ndims = len(self._event_shape)<EOL><DEDENT>if input_shape[-<NUM_LIT:1>] != self._event_shape[-<NUM_LIT:1>]:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>self._event_shape[-<NUM_LIT:1>], input_shape[-<NUM_LIT:1>]))<EOL><DEDENT>self._input_order = _create_input_order(<EOL>self._event_size, self._input_order_param)<EOL>self._masks = _create_masks(_create_degrees(<EOL>input_size=self._event_size,<EOL>hidden_units=self._hidden_units,<EOL>input_order=self._input_order,<EOL>hidden_degrees=self._hidden_degrees))<EOL>self._masks[-<NUM_LIT:1>] = np.reshape(<EOL>np.tile(self._masks[-<NUM_LIT:1>][..., tf.newaxis], [<NUM_LIT:1>, <NUM_LIT:1>, self._params]),<EOL>[self._masks[-<NUM_LIT:1>].shape[<NUM_LIT:0>], self._event_size * self._params])<EOL>self._network = tf.keras.Sequential([<EOL>tf.keras.layers.InputLayer((self._event_size,), dtype=self.dtype)<EOL>])<EOL>layer_output_sizes = self._hidden_units + [self._event_size * self._params]<EOL>for k in range(len(self._masks)):<EOL><INDENT>self._network.add(tf.keras.layers.Dense(<EOL>layer_output_sizes[k],<EOL>kernel_initializer=_make_masked_initializer(<EOL>self._masks[k], self._kernel_initializer),<EOL>kernel_constraint=_make_masked_constraint(self._masks[k]),<EOL>activation=self._activation if k + <NUM_LIT:1> < len(self._masks) else None,<EOL>use_bias=self._use_bias,<EOL>**self._kwargs))<EOL><DEDENT>super(AutoregressiveLayer, self).build(input_shape)<EOL>


with tf.compat.v2.name_scope(self.name or "<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, dtype=self.dtype, name="<STR_LIT:x>")<EOL>input_shape = tf.shape(input=x)<EOL>if tensorshape_util.rank(x.shape) == <NUM_LIT:1>:<EOL><INDENT>x = x[tf.newaxis, ...]<EOL><DEDENT>return tf.reshape(self._network(x),<EOL>tf.concat([input_shape, [self._params]], axis=<NUM_LIT:0>))<EOL><DEDENT>


return input_shape + (self._params,)<EOL>


return -np.log(<NUM_LIT:1> - np.exp(-y))<EOL>

Inverse log det jacobian, before being reduced.


<EOL>ten_x_pts = np.linspace(lower_x, upper_x, num=<NUM_LIT:10>).astype(np.float32)<EOL>if bijector.dtype is not None:<EOL><INDENT>ten_x_pts = ten_x_pts.astype(dtype_util.as_numpy_dtype(bijector.dtype))<EOL><DEDENT>forward_on_10_pts = bijector.forward(ten_x_pts)<EOL>lower_y, upper_y = eval_func(<EOL>[bijector.forward(lower_x),<EOL>bijector.forward(upper_x)])<EOL>if upper_y < lower_y:  <EOL><INDENT>lower_y, upper_y = upper_y, lower_y<EOL><DEDENT>uniform_x_samps = uniform_distribution.Uniform(<EOL>low=lower_x, high=upper_x).sample(<EOL>n, seed=<NUM_LIT:0>)<EOL>uniform_y_samps = uniform_distribution.Uniform(<EOL>low=lower_y, high=upper_y).sample(<EOL>n, seed=<NUM_LIT:1>)<EOL>inverse_forward_x = bijector.inverse(bijector.forward(uniform_x_samps))<EOL>forward_inverse_y = bijector.forward(bijector.inverse(uniform_y_samps))<EOL>dy_dx = tf.exp(<EOL>bijector.inverse_log_det_jacobian(uniform_y_samps, event_ndims=<NUM_LIT:0>))<EOL>expectation_of_dy_dx_under_uniform = tf.reduce_mean(input_tensor=dy_dx)<EOL>change_measure_dy_dx = (<EOL>(upper_y - lower_y) * expectation_of_dy_dx_under_uniform)<EOL>dx_dy = tf.exp(<EOL>bijector.forward_log_det_jacobian(<EOL>bijector.inverse(uniform_y_samps), event_ndims=<NUM_LIT:0>))<EOL>[<EOL>forward_on_10_pts_v,<EOL>dy_dx_v,<EOL>dx_dy_v,<EOL>change_measure_dy_dx_v,<EOL>uniform_x_samps_v,<EOL>uniform_y_samps_v,<EOL>inverse_forward_x_v,<EOL>forward_inverse_y_v,<EOL>] = eval_func([<EOL>forward_on_10_pts,<EOL>dy_dx,<EOL>dx_dy,<EOL>change_measure_dy_dx,<EOL>uniform_x_samps,<EOL>uniform_y_samps,<EOL>inverse_forward_x,<EOL>forward_inverse_y,<EOL>])<EOL>assert_strictly_monotonic(forward_on_10_pts_v)<EOL>np.testing.assert_allclose(<EOL>inverse_forward_x_v, uniform_x_samps_v, atol=<NUM_LIT>, rtol=<NUM_LIT>)<EOL>np.testing.assert_allclose(<EOL>forward_inverse_y_v, uniform_y_samps_v, atol=<NUM_LIT>, rtol=<NUM_LIT>)<EOL>np.testing.assert_allclose(<EOL>upper_x - lower_x, change_measure_dy_dx_v, atol=<NUM_LIT:0>, rtol=rtol)<EOL>np.testing.assert_allclose(<EOL>dy_dx_v, np.divide(<NUM_LIT:1.>, dx_dy_v), atol=<NUM_LIT>, rtol=<NUM_LIT>)<EOL>

Assert `bijector`'s forward/inverse/inverse_log_det_jacobian are congruent.

    We draw samples `X ~ U(lower_x, upper_x)`, then feed these through the
    `bijector` in order to check that:

    1. the forward is strictly monotonic.
    2. the forward/inverse methods are inverses of each other.
    3. the jacobian is the correct change of measure.

    This can only be used for a Bijector mapping open subsets of the real line
    to themselves.  This is due to the fact that this test compares the `prob`
    before/after transformation with the Lebesgue measure on the line.

    Args:
      bijector:  Instance of Bijector
      lower_x:  Python scalar.
      upper_x:  Python scalar.  Must have `lower_x < upper_x`, and both must be in
        the domain of the `bijector`.  The `bijector` should probably not produce
        huge variation in values in the interval `(lower_x, upper_x)`, or else the
        variance based check of the Jacobian will require small `rtol` or huge
        `n`.
      eval_func: Function to evaluate any intermediate results.
      n:  Number of samples to draw for the checks.
      rtol:  Positive number.  Used for the Jacobian check.

    Raises:
      AssertionError:  If tests fail.


if inverse_event_ndims is None:<EOL><INDENT>inverse_event_ndims = event_ndims<EOL><DEDENT>assert_finite(x)<EOL>assert_finite(y)<EOL>f_x = bijector.forward(x)<EOL>g_y = bijector.inverse(y)<EOL>[<EOL>x_from_x,<EOL>y_from_y,<EOL>ildj_f_x,<EOL>fldj_x,<EOL>ildj_y,<EOL>fldj_g_y,<EOL>f_x_v,<EOL>g_y_v,<EOL>] = eval_func([<EOL>bijector.inverse(f_x),<EOL>bijector.forward(g_y),<EOL>bijector.inverse_log_det_jacobian(f_x, event_ndims=inverse_event_ndims),<EOL>bijector.forward_log_det_jacobian(x, event_ndims=event_ndims),<EOL>bijector.inverse_log_det_jacobian(y, event_ndims=inverse_event_ndims),<EOL>bijector.forward_log_det_jacobian(g_y, event_ndims=event_ndims),<EOL>f_x,<EOL>g_y,<EOL>])<EOL>assert_finite(x_from_x)<EOL>assert_finite(y_from_y)<EOL>assert_finite(ildj_f_x)<EOL>assert_finite(fldj_x)<EOL>assert_finite(ildj_y)<EOL>assert_finite(fldj_g_y)<EOL>assert_finite(f_x_v)<EOL>assert_finite(g_y_v)<EOL>np.testing.assert_allclose(x_from_x, x, atol=atol, rtol=rtol)<EOL>np.testing.assert_allclose(y_from_y, y, atol=atol, rtol=rtol)<EOL>np.testing.assert_allclose(-ildj_f_x, fldj_x, atol=atol, rtol=rtol)<EOL>np.testing.assert_allclose(-ildj_y, fldj_g_y, atol=atol, rtol=rtol)<EOL>

Assert that forward/inverse (along with jacobians) are inverses and finite.

    It is recommended to use x and y values that are very very close to the edge
    of the Bijector's domain.

    Args:
      bijector:  A Bijector instance.
      x:  np.array of values in the domain of bijector.forward.
      y:  np.array of values in the domain of bijector.inverse.
      event_ndims: Integer describing the number of event dimensions this bijector
        operates on.
      eval_func: Function to evaluate any intermediate results.
      inverse_event_ndims: Integer describing the number of event dimensions for
        the bijector codomain. If None, then the value of `event_ndims` is used.
      atol:  Absolute tolerance.
      rtol:  Relative tolerance.

    Raises:
      AssertionError:  If tests fail.


if input_to_unconstrained is None:<EOL><INDENT>input_to_unconstrained = reshape_bijector.Reshape(<EOL>event_shape_in=x.shape[tensorshape_util.rank(x.shape) - event_ndims:],<EOL>event_shape_out=[-<NUM_LIT:1>])<EOL><DEDENT>if output_to_unconstrained is None:<EOL><INDENT>output_to_unconstrained = reshape_bijector.Reshape(<EOL>event_shape_in=x.shape[tensorshape_util.rank(x.shape) - event_ndims:],<EOL>event_shape_out=[-<NUM_LIT:1>])<EOL><DEDENT>x = tf.convert_to_tensor(value=x)<EOL>x_unconstrained = <NUM_LIT:1> * input_to_unconstrained.forward(x)<EOL>with tf.GradientTape(persistent=True) as tape:<EOL><INDENT>tape.watch(x_unconstrained)<EOL>f_x = bijector.forward(input_to_unconstrained.inverse(x_unconstrained))<EOL>f_x_unconstrained = output_to_unconstrained.forward(f_x)<EOL><DEDENT>jacobian = tape.batch_jacobian(<EOL>f_x_unconstrained, x_unconstrained, experimental_use_pfor=False)<EOL>return (<EOL>tf.linalg.slogdet(jacobian).log_abs_determinant +<EOL>input_to_unconstrained.forward_log_det_jacobian(<EOL>x, event_ndims=event_ndims) -<EOL>output_to_unconstrained.forward_log_det_jacobian(<EOL>f_x, event_ndims=event_ndims))<EOL>

Numerically approximate the forward log det Jacobian of a bijector.

    We compute the Jacobian of the chain
    output_to_unconst_vec(bijector(inverse(input_to_unconst_vec))) so that
    we're working with a full rank matrix.  We then adjust the resulting Jacobian
    for the unconstraining bijectors.

    Bijectors that constrain / unconstrain their inputs/outputs may not be
    testable with this method, since the composition above may reduce the test
    to something trivial.  However, bijectors that map within constrained spaces
    should be fine.

    Args:
      bijector: the bijector whose Jacobian we wish to approximate
      x: the value for which we want to approximate the Jacobian.  x must have
        a a single batch dimension for compatibility with tape.batch_jacobian.
      event_ndims: number of dimensions in an event
      input_to_unconstrained: bijector that maps the input to the above bijector
        to an unconstrained 1-D vector.  If unspecified, flatten the input into
        a 1-D vector according to its event_ndims.
      output_to_unconstrained: bijector that maps the output of the above bijector
        to an unconstrained 1-D vector.  If unspecified, flatten the input into
        a 1-D vector according to its event_ndims.

    Returns:
      A numerical approximation to the log det Jacobian of bijector.forward
      evaluated at x.


min_event_ndims = <NUM_LIT:0><EOL>rank_changed_adjusted_max_min_event_ndims = <NUM_LIT:0><EOL>if compute_forward:<EOL><INDENT>bijector_list = reversed(bijector_list)<EOL><DEDENT>for b in bijector_list:<EOL><INDENT>if compute_forward:<EOL><INDENT>current_min_event_ndims = b.forward_min_event_ndims<EOL>current_inverse_min_event_ndims = b.inverse_min_event_ndims<EOL><DEDENT>else:<EOL><INDENT>current_min_event_ndims = b.inverse_min_event_ndims<EOL>current_inverse_min_event_ndims = b.forward_min_event_ndims<EOL><DEDENT>if rank_changed_adjusted_max_min_event_ndims < current_min_event_ndims:<EOL><INDENT>min_event_ndims += (<EOL>current_min_event_ndims - rank_changed_adjusted_max_min_event_ndims)<EOL><DEDENT>rank_changed_adjusted_max_min_event_ndims = max(<EOL>current_min_event_ndims, rank_changed_adjusted_max_min_event_ndims)<EOL>number_of_changed_dimensions = (<EOL>current_min_event_ndims - current_inverse_min_event_ndims)<EOL>rank_changed_adjusted_max_min_event_ndims -= number_of_changed_dimensions<EOL><DEDENT>return min_event_ndims<EOL>

Computes the min_event_ndims associated with the give list of bijectors.

    Given a list `bijector_list` of bijectors, compute the min_event_ndims that is
    associated with the composition of bijectors in that list.

    min_event_ndims is the # of right most dimensions for which the bijector has
    done necessary computation on (i.e. the non-broadcastable part of the
    computation).

    We can derive the min_event_ndims for a chain of bijectors as follows:

    In the case where there are no rank changing bijectors, this will simply be
    `max(b.forward_min_event_ndims for b in bijector_list)`. This is because the
    bijector with the most forward_min_event_ndims requires the most dimensions,
    and hence the chain also requires operating on those dimensions.

    However in the case of rank changing, more care is needed in determining the
    exact amount of dimensions. Padding dimensions causes subsequent bijectors to
    operate on the padded dimensions, and Removing dimensions causes bijectors to
    operate more left.

    Args:
      bijector_list: List of bijectors to be composed by chain.
      compute_forward: Boolean. If True, computes the min_event_ndims associated
        with a forward call to Chain, and otherwise computes the min_event_ndims
        associated with an inverse call to Chain. The latter is the same as the
        min_event_ndims associated with a forward call to Invert(Chain(....)).

    Returns:
      min_event_ndims


def __init__(self, bijectors=None, validate_args=False, name=None):

if bijectors is None:<EOL><INDENT>bijectors = ()<EOL><DEDENT>self._bijectors = bijectors<EOL>for a_bijector in bijectors:<EOL><INDENT>if not a_bijector._is_injective:  <EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>".format(<EOL>a_bijector.name))<EOL><DEDENT><DEDENT>dtype = list(set([b.dtype for b in bijectors if b.dtype is not None]))<EOL>if len(dtype) > <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>" % dtype)<EOL><DEDENT>elif len(dtype) == <NUM_LIT:1>:<EOL><INDENT>dtype = dtype[<NUM_LIT:0>]<EOL><DEDENT>else:<EOL><INDENT>dtype = None<EOL><DEDENT>inverse_min_event_ndims = _compute_min_event_ndims(<EOL>bijectors, compute_forward=False)<EOL>forward_min_event_ndims = _compute_min_event_ndims(<EOL>bijectors, compute_forward=True)<EOL>super(Chain, self).__init__(<EOL>graph_parents=list(itertools.chain.from_iterable(<EOL>b.graph_parents for b in bijectors)),<EOL>forward_min_event_ndims=forward_min_event_ndims,<EOL>inverse_min_event_ndims=inverse_min_event_ndims,<EOL>is_constant_jacobian=all(b.is_constant_jacobian for b in bijectors),<EOL>validate_args=validate_args,<EOL>dtype=dtype,<EOL>name=name or ("<STR_LIT>" if not bijectors else<EOL>"<STR_LIT>".join(["<STR_LIT>"] + [b.name for b in bijectors])))<EOL>

Instantiates `Chain` bijector.

        Args:
          bijectors: Python `list` of bijector instances. An empty list makes this
            bijector equivalent to the `Identity` bijector.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str`, name given to ops managed by this object. Default:
            E.g., `Chain([Exp(), Softplus()]).name == "chain_of_exp_of_softplus"`.

        Raises:
          ValueError: if bijectors have different dtypes.


ndarray = np.asarray(ndarray).astype(self.dtype)<EOL>return tf.compat.v1.placeholder_with_default(<EOL>input=ndarray, shape=ndarray.shape if self.use_static_shape else None)<EOL>

Convert a numpy array to a TF placeholder.

        Args:
          ndarray: any object convertible to a numpy array via `np.asarray()`.

        Returns:
          placeholder: a TensorFlow `placeholder` with default value given by the
          provided `ndarray`, dtype given by `self.dtype`, and shape specified
          statically only if `self.use_static_shape` is `True`.


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[level_scale, slope_scale]) as name:<EOL><INDENT>dtype = initial_state_prior.dtype<EOL>level_scale = tf.convert_to_tensor(<EOL>value=level_scale, name='<STR_LIT>', dtype=dtype)<EOL>slope_scale = tf.convert_to_tensor(<EOL>value=slope_scale, name='<STR_LIT>', dtype=dtype)<EOL>observation_noise_scale = tf.convert_to_tensor(<EOL>value=observation_noise_scale,<EOL>name='<STR_LIT>',<EOL>dtype=dtype)<EOL>broadcast_batch_shape = dist_util.get_broadcast_shape(<EOL>level_scale, slope_scale)<EOL>broadcast_ones = tf.ones(broadcast_batch_shape, dtype=dtype)<EOL>self._level_scale = level_scale<EOL>self._slope_scale = slope_scale<EOL>self._observation_noise_scale = observation_noise_scale<EOL>super(LocalLinearTrendStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=tf.constant(<EOL>[[<NUM_LIT:1.>, <NUM_LIT:1.>], [<NUM_LIT:0.>, <NUM_LIT:1.>]], dtype=dtype, name='<STR_LIT>'),<EOL>transition_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.stack(<EOL>[level_scale * broadcast_ones, slope_scale * broadcast_ones],<EOL>axis=-<NUM_LIT:1>),<EOL>name='<STR_LIT>'),<EOL>observation_matrix=tf.constant(<EOL>[[<NUM_LIT:1.>, <NUM_LIT:0.>]], dtype=dtype, name='<STR_LIT>'),<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=observation_noise_scale[..., tf.newaxis],<EOL>name='<STR_LIT>'),<EOL>initial_state_prior=initial_state_prior,<EOL>initial_step=initial_step,<EOL>allow_nan_stats=allow_nan_stats,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Build a state space model implementing a local linear trend.

        Args:
          num_timesteps: Scalar `int` `Tensor` number of timesteps to model
            with this distribution.
          level_scale: Scalar (any additional dimensions are treated as batch
            dimensions) `float` `Tensor` indicating the standard deviation of the
            level transitions.
          slope_scale: Scalar (any additional dimensions are treated as batch
            dimensions) `float` `Tensor` indicating the standard deviation of the
            slope transitions.
          initial_state_prior: instance of `tfd.MultivariateNormal`
            representing the prior distribution on latent states; must
            have event shape `[2]`.
          observation_noise_scale: Scalar (any additional dimensions are
            treated as batch dimensions) `float` `Tensor` indicating the standard
            deviation of the observation noise.
          initial_step: Optional scalar `int` `Tensor` specifying the starting
            timestep.
            Default value: 0.
          validate_args: Python `bool`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are
            invalid, correct behavior is not guaranteed.
            Default value: `False`.
          allow_nan_stats: Python `bool`. If `False`, raise an
            exception if a statistic (e.g. mean/mode/etc...) is undefined for any
            batch member. If `True`, batch members with valid parameters leading to
            undefined statistics will return NaN for this statistic.
            Default value: `True`.
          name: Python `str` name prefixed to ops created by this class.
            Default value: "LocalLinearTrendStateSpaceModel".


return self._level_scale<EOL>

Standard deviation of the level transitions.


return self._slope_scale<EOL>

Standard deviation of the slope transitions.


return self._observation_noise_scale<EOL>

Standard deviation of the observation noise.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>_, observed_stddev, observed_initial = (<EOL>sts_util.empirical_statistics(observed_time_series)<EOL>if observed_time_series is not None else (<NUM_LIT:0.>, <NUM_LIT:1.>, <NUM_LIT:0.>))<EOL>if level_scale_prior is None:<EOL><INDENT>level_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev),<EOL>scale=<NUM_LIT>,<EOL>name='<STR_LIT>')<EOL><DEDENT>if slope_scale_prior is None:<EOL><INDENT>slope_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev),<EOL>scale=<NUM_LIT>,<EOL>name='<STR_LIT>')<EOL><DEDENT>if initial_level_prior is None:<EOL><INDENT>initial_level_prior = tfd.Normal(<EOL>loc=observed_initial,<EOL>scale=tf.abs(observed_initial) + observed_stddev,<EOL>name='<STR_LIT>')<EOL><DEDENT>if initial_slope_prior is None:<EOL><INDENT>initial_slope_prior = tfd.Normal(<EOL>loc=<NUM_LIT:0.>, scale=observed_stddev, name='<STR_LIT>')<EOL><DEDENT>tf.debugging.assert_same_float_dtype([<EOL>level_scale_prior, slope_scale_prior, initial_level_prior,<EOL>initial_slope_prior<EOL>])<EOL>self._initial_state_prior = tfd.MultivariateNormalDiag(<EOL>loc=tf.stack(<EOL>[initial_level_prior.mean(),<EOL>initial_slope_prior.mean()<EOL>], axis=-<NUM_LIT:1>),<EOL>scale_diag=tf.stack([<EOL>initial_level_prior.stddev(),<EOL>initial_slope_prior.stddev()<EOL>], axis=-<NUM_LIT:1>))<EOL>super(LocalLinearTrend, self).__init__(<EOL>parameters=[<EOL>Parameter('<STR_LIT>', level_scale_prior, tfb.Softplus()),<EOL>Parameter('<STR_LIT>', slope_scale_prior, tfb.Softplus())<EOL>],<EOL>latent_size=<NUM_LIT:2>,<EOL>name=name)<EOL><DEDENT>

Specify a local linear trend model.

        Args:
          level_scale_prior: optional `tfd.Distribution` instance specifying a prior
            on the `level_scale` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          slope_scale_prior: optional `tfd.Distribution` instance specifying a prior
            on the `slope_scale` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          initial_level_prior: optional `tfd.Distribution` instance specifying a
            prior on the initial level. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          initial_slope_prior: optional `tfd.Distribution` instance specifying a
            prior on the initial slope. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          observed_time_series: optional `float` `Tensor` of shape
            `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
            supported when `T > 1`), specifying an observed time series.
            Any priors not explicitly set will be given default values according to
            the scale of the observed time series (or batch of time series). May
            optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
            a mask `Tensor` to specify timesteps with missing observations.
            Default value: `None`.
          name: the name of this model component.
            Default value: 'LocalLinearTrend'.


return self._initial_state_prior<EOL>

Prior distribution on the initial latent state (level and scale).


def _build_tensor(self, ndarray, dtype=None):

ndarray = np.asarray(ndarray).astype(self.dtype if dtype is None else dtype)<EOL>return tf.compat.v1.placeholder_with_default(<EOL>input=ndarray, shape=ndarray.shape if self.use_static_shape else None)<EOL>

Convert a numpy array to a TF placeholder.

        Args:
          ndarray: any object convertible to a numpy array via `np.asarray()`.
          dtype: optional `dtype`.

        Returns:
          placeholder: a TensorFlow `placeholder` with default value given by the
          provided `ndarray`, dtype given by `self.dtype` (if not specified), and
          shape specified statically only if `self.use_static_shape` is `True`.


def one_step_predictive(model, observed_time_series, parameter_samples):

with tf.compat.v1.name_scope(<EOL>'<STR_LIT>', values=[observed_time_series, parameter_samples]):<EOL><INDENT>[<EOL>observed_time_series,<EOL>is_missing<EOL>] = sts_util.canonicalize_observed_time_series_with_mask(<EOL>observed_time_series)<EOL>num_timesteps = dist_util.prefer_static_value(<EOL>tf.shape(input=observed_time_series))[-<NUM_LIT:2>]<EOL>lgssm = model.make_state_space_model(<EOL>num_timesteps=num_timesteps, param_vals=parameter_samples)<EOL>(_, _, _, _, _, observation_means, observation_covs<EOL>) = lgssm.forward_filter(observed_time_series, mask=is_missing)<EOL>return sts_util.mix_over_posterior_draws(<EOL>means=observation_means[..., <NUM_LIT:0>],<EOL>variances=observation_covs[..., <NUM_LIT:0>, <NUM_LIT:0>])<EOL><DEDENT>

Compute one-step-ahead predictive distributions for all timesteps.

    Given samples from the posterior over parameters, return the predictive
    distribution over observations at each time `T`, given observations up
    through time `T-1`.

    Args:
      model: An instance of `StructuralTimeSeries` representing a
        time-series model. This represents a joint distribution over
        time-series and their parameters with batch shape `[b1, ..., bN]`.
      observed_time_series: `float` `Tensor` of shape
        `concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
        `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
        dimension may (optionally) be omitted if `num_timesteps > 1`. May
        optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
        mask `Tensor` to encode the locations of missing observations.
      parameter_samples: Python `list` of `Tensors` representing posterior samples
        of model parameters, with shapes `[concat([[num_posterior_draws],
        param.prior.batch_shape, param.prior.event_shape]) for param in
        model.parameters]`. This may optionally also be a map (Python `dict`) of
        parameter names to `Tensor` values.

    Returns:
      forecast_dist: a `tfd.MixtureSameFamily` instance with event shape
        [num_timesteps] and
        batch shape `concat([sample_shape, model.batch_shape])`, with
        `num_posterior_draws` mixture components. The `t`th step represents the
        forecast distribution `p(observed_time_series[t] |
        observed_time_series[0:t-1], parameter_samples)`.

    #### Examples

    Suppose we've built a model and fit it to data using HMC:

    ```python
      day_of_week = tfp.sts.Seasonal(
          num_seasons=7,
          observed_time_series=observed_time_series,
          name='day_of_week')
      local_linear_trend = tfp.sts.LocalLinearTrend(
          observed_time_series=observed_time_series,
          name='local_linear_trend')
      model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                          observed_time_series=observed_time_series)

      samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
    ```

    Passing the posterior samples into `one_step_predictive`, we construct a
    one-step-ahead predictive distribution:

    ```python
      one_step_predictive_dist = tfp.sts.one_step_predictive(
        model, observed_time_series, parameter_samples=samples)

      predictive_means = one_step_predictive_dist.mean()
      predictive_scales = one_step_predictive_dist.stddev()
    ```

    If using variational inference instead of HMC, we'd construct a forecast using
    samples from the variational posterior:

    ```python
      (variational_loss,
       variational_distributions) = tfp.sts.build_factored_variational_loss(
         model=model, observed_time_series=observed_time_series)

      # OMITTED: take steps to optimize variational loss

      samples = {k: q.sample(30) for (k, q) in variational_distributions.items()}
      one_step_predictive_dist = tfp.sts.one_step_predictive(
        model, observed_time_series, parameter_samples=samples)
    ```

    We can visualize the forecast by plotting:

    ```python
      from matplotlib import pylab as plt
      def plot_one_step_predictive(observed_time_series,
                                   forecast_mean,
                                   forecast_scale):
        plt.figure(figsize=(12, 6))
        num_timesteps = forecast_mean.shape[-1]
        c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
        plt.plot(observed_time_series, label="observed time series", color=c1)
        plt.plot(forecast_mean, label="one-step prediction", color=c2)
        plt.fill_between(np.arange(num_timesteps),
                         forecast_mean - 2 * forecast_scale,
                         forecast_mean + 2 * forecast_scale,
                         alpha=0.1, color=c2)
        plt.legend()

      plot_one_step_predictive(observed_time_series,
                               forecast_mean=predictive_means,
                               forecast_scale=predictive_scales)
    ```

    To detect anomalous timesteps, we check whether the observed value at each
    step is within a 95% predictive interval, i.e., two standard deviations from
    the mean:

    ```python
      z_scores = ((observed_time_series[..., 1:] - predictive_means[..., :-1])
                   / predictive_scales[..., :-1])
      anomalous_timesteps = tf.boolean_mask(
          tf.range(1, num_timesteps),
          tf.abs(z_scores) > 2.0)
    ```