def _apply(self, x1, x2, param_expansion_ndims=<NUM_LIT:0>):

raise NotImplementedError(<EOL>'<STR_LIT>')<EOL>

Apply the kernel function to a pair of (batches of) inputs.

        Args:
          x1: `Tensor` input to the first positional parameter of the kernel, of
            shape `[b1, ..., bB, f1, ..., fF]`, where `B` may be zero (no batching)
            and `F` (number of feature dimensions) must equal the kernel's
            `feature_ndims` property. Batch shape must broadcast with the batch
            shape of `x2` and with the kernel's parameters *after* parameter
            expansion (see `param_expansion_ndims` argument).
          x2: `Tensor` input to the second positional parameter of the kernel,
            shape `[c1, ..., cC, f1, ..., fF]`, where `C` may be zero (no batching)
            and `F` (number of feature dimensions) must equal the kernel's
            `feature_ndims` property. Batch shape must broadcast with the batch
            shape of `x1` and with the kernel's parameters *after* parameter
            expansion (see `param_expansion_ndims` argument).
          param_expansion_ndims: Python `integer` enabling reshaping of kernel
            parameters by concatenating a list of 1's to the param shapes. This
            allows the caller to control how the parameters broadcast across the
            inputs.

        Returns:
          `Tensor` containing the (batch of) results of applying the kernel function
          to inputs `x1` and `x2`. If the kernel parameters' batch shape *after*
          parameter expansion (ie, concatenating `param_expansion_ndims` 1's onto
          the parameters' shapes) is `[k1, ..., kK, 1, ..., 1]` then the shape of
          the `Tensor` resulting from this method call is
          `broadcast([b1, ..., bB], [c1, ..., cC], [k1, ..., kK, 1, ..., 1])`.


return self._kernels<EOL>

The list of kernels this _SumKernel sums over.


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

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

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

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


return self._kernels<EOL>

The list of kernels this _ProductKernel multiplies over.


return self._name<EOL>

Name prepended to all ops created by this class.


with self._name_scope(self._name, values=[x1, x2]):<EOL><INDENT>x1 = tf.convert_to_tensor(value=x1, name='<STR_LIT>')<EOL>x2 = tf.convert_to_tensor(value=x2, name='<STR_LIT>')<EOL>x1 = tf.expand_dims(x1, -(self.feature_ndims + <NUM_LIT:1>))<EOL>x2 = tf.expand_dims(x2, -(self.feature_ndims + <NUM_LIT:2>))<EOL>return self._apply(x1, x2, param_expansion_ndims=<NUM_LIT:2>)<EOL><DEDENT>

Construct (batched) matrices from (batches of) collections of inputs.

        Args:
          x1: `Tensor` input to the first positional parameter of the kernel, of
            shape `[b1, ..., bB, e1, f1, ..., fF]`, where `B` may be zero (ie, no
            batching), e1 is an integer greater than zero, and `F` (number of
            feature dimensions) must equal the kernel's `feature_ndims` property.
            Batch shape must broadcast with the batch shape of `x2` and with the
            kernel's parameters *after* parameter expansion (see
            `param_expansion_ndims` argument).
          x2: `Tensor` input to the second positional parameter of the kernel,
            shape `[c1, ..., cC, e2, f1, ..., fF]`, where `C` may be zero (ie, no
            batching), e2 is an integer greater than zero,  and `F` (number of
            feature dimensions) must equal the kernel's `feature_ndims` property.
            Batch shape must broadcast with the batch shape of `x1` and with the
            kernel's parameters *after* parameter expansion (see
            `param_expansion_ndims` argument).

        Returns:
          `Tensor containing (batch of) matrices of kernel applications to pairs
          from inputs `x1` and `x2`. If the kernel parameters' batch shape is
          `[k1, ..., kK]`, then the shape of the resulting `Tensor` is
          `broadcast([b1, ..., bB], [c1, ..., cC], [k1, ..., kK]) + [e1, e2]`.

        Given inputs `x1` and `x2` of shapes

        ```none
        [b1, ..., bB, e1, f1, ..., fF]
        ```

        and

        ```none
        [c1, ..., cC, e2, f1, ..., fF]
        ```

        This method computes the batch of `e1 x e2` matrices resulting from applying
        the kernel function to all pairs of inputs from `x1` and `x2`. The shape
        of the batch of matrices is the result of broadcasting the batch shapes of
        `x1`, `x2`, and the kernel parameters (see examples below). As such, it's
        required that these shapes all be broadcast compatible. However, the kernel
        parameter batch shapes need not broadcast against the 'example shapes' (`e1`
        and `e2` above).

        When the two inputs are the (batches of) identical collections, the
        resulting matrix is the so-called Gram (or Gramian) matrix
        (https://en.wikipedia.org/wiki/Gramian_matrix).

        N.B., this method can only be used to compute the pairwise application of
        the kernel function on rank-1 collections. E.g., it *does* support inputs of
        shape `[e1, f]` and `[e2, f]`, yielding a matrix of shape `[e1, e2]`. It
        *does not* support inputs of shape `[e1, e2, f]` and `[e3, e4, f]`, yielding
        a `Tensor` of shape `[e1, e2, e3, e4]`. To do this, one should instead
        reshape the inputs and pass them to `apply`, e.g.:

        ```python
        k = tfpk.SomeKernel()
        t1 = tf.placeholder([4, 4, 3], tf.float32)
        t2 = tf.placeholder([5, 5, 3], tf.float32)
        k.apply(
            tf.reshape(t1, [4, 4, 1, 1, 3]),
            tf.reshape(t2, [1, 1, 5, 5, 3])).shape
        # ==> [4, 4, 5, 5, 3]
        ```

        `matrix` is a special case of the above, where there is only one example
        dimension; indeed, its implementation looks almost exactly like the above
        (reshaped inputs passed to the private version of `_apply`).

        #### Examples

        First, consider a kernel with a single scalar parameter.

        ```python
        import tensorflow_probability as tfp

        scalar_kernel = tfp.positive_semidefinite_kernels.SomeKernel(param=.5)
        scalar_kernel.batch_shape
        # ==> []

        # Our inputs are two lists of 3-D vectors
        x = np.ones([5, 3], np.float32)
        y = np.ones([4, 3], np.float32)
        scalar_kernel.matrix(x, y).shape
        # ==> [5, 4]
        ```

        The result comes from applying the kernel to the entries in `x` and `y`
        pairwise, across all pairs:

          ```none
          | k(x[0], y[0])    k(x[0], y[1])  ...  k(x[0], y[3]) |
          | k(x[1], y[0])    k(x[1], y[1])  ...  k(x[1], y[3]) |
          |      ...              ...                 ...      |
          | k(x[4], y[0])    k(x[4], y[1])  ...  k(x[4], y[3]) |
          ```

        Now consider a kernel with batched parameters with the same inputs

        ```python
        batch_kernel = tfp.positive_semidefinite_kernels.SomeKernel(param=[1., .5])
        batch_kernel.batch_shape
        # ==> [2]

        batch_kernel.matrix(x, y).shape
        # ==> [2, 5, 4]
        ```

        This results in a batch of 2 matrices, one computed from the kernel with
        `param = 1.` and the other with `param = .5`.

        We also support batching of the inputs. First, let's look at that with
        the scalar kernel again.

        ```python
        # Batch of 10 lists of 5 vectors of dimension 3
        x = np.ones([10, 5, 3], np.float32)

        # Batch of 10 lists of 4 vectors of dimension 3
        y = np.ones([10, 4, 3], np.float32)

        scalar_kernel.matrix(x, y).shape
        # ==> [10, 5, 4]
        ```

        The result is a batch of 10 matrices built from the batch of 10 lists of
        input vectors. These batch shapes have to be broadcastable. The following
        will *not* work:

        ```python
        x = np.ones([10, 5, 3], np.float32)
        y = np.ones([20, 4, 3], np.float32)
        scalar_kernel.matrix(x, y).shape
        # ==> Error! [10] and [20] can't broadcast.
        ```

        Now let's consider batches of inputs in conjunction with batches of kernel
        parameters. We require that the input batch shapes be broadcastable with
        the kernel parameter batch shapes, otherwise we get an error:

        ```python
        x = np.ones([10, 5, 3], np.float32)
        y = np.ones([10, 4, 3], np.float32)

        batch_kernel = tfp.positive_semidefinite_kernels.SomeKernel(params=[1., .5])
        batch_kernel.batch_shape
        # ==> [2]
        batch_kernel.matrix(x, y).shape
        # ==> Error! [2] and [10] can't broadcast.
        ```

        The fix is to make the kernel parameter shape broadcastable with `[10]` (or
        reshape the inputs to be broadcastable!):

        ```python
        x = np.ones([10, 5, 3], np.float32)
        y = np.ones([10, 4, 3], np.float32)

        batch_kernel = tfp.positive_semidefinite_kernels.SomeKernel(
            params=[[1.], [.5]])
        batch_kernel.batch_shape
        # ==> [2, 1]
        batch_kernel.matrix(x, y).shape
        # ==> [2, 10, 5, 4]

        # Or, make the inputs broadcastable:
        x = np.ones([10, 1, 5, 3], np.float32)
        y = np.ones([10, 1, 4, 3], np.float32)

        batch_kernel = tfp.positive_semidefinite_kernels.SomeKernel(
            params=[1., .5])
        batch_kernel.batch_shape
        # ==> [2]
        batch_kernel.matrix(x, y).shape
        # ==> [10, 2, 5, 4]

        ```

        Here, we have the result of applying the kernel, with 2 different
        parameters, to each of a batch of 10 pairs of input lists.


@contextlib.contextmanager<EOL><INDENT>def _name_scope(self, name=None, values=None):<DEDENT>

with tf.compat.v1.name_scope(self.name):<EOL><INDENT>values = [] if values is None else values<EOL>with tf.compat.v1.name_scope(name, values=values) as scope:<EOL><INDENT>yield scope<EOL><DEDENT><DEDENT>

Helper function to standardize op scope.


with tf.compat.v1.name_scope(<EOL>name, values=[amplitude, length_scale]) as name:<EOL><INDENT>dtype = super(MaternFiveHalves, self)._init_params(<EOL>amplitude, length_scale, validate_args)<EOL><DEDENT>super(MaternFiveHalves, self).__init__(<EOL>feature_ndims, dtype=dtype, name=name)<EOL>

Construct a MaternFiveHalves kernel instance.

        Args:
          amplitude: Positive floating point `Tensor` that controls the maximum
            value of the kernel. Must be broadcastable with `length_scale` and
            inputs to `apply` and `matrix` methods. A value of `None` is treated
            like 1.
          length_scale: Positive floating point `Tensor` that controls how sharp or
            wide the kernel shape is. This provides a characteristic "unit" of
            length against which `||x - y||` can be compared for scale. Must be
            broadcastable with `amplitude`, and inputs to `apply` and `matrix`
            methods. A value of `None` is treated like 1.
          feature_ndims: Python `int` number of rightmost dims to include in the
            squared difference norm in the exponential.
          validate_args: If `True`, parameters are checked for validity despite
            possibly degrading runtime performance
          name: Python `str` name prefixed to Ops created by this class.


return self._amplitude<EOL>


return self._shift<EOL>

Shift of linear function that is exponentiated.


return self._slope_variance<EOL>


@interceptable<EOL>@functools.wraps(distribution_cls, assigned=('<STR_LIT>', '<STR_LIT>'))<EOL>@docstring_util.expand_docstring(<EOL>cls=distribution_cls.__name__,<EOL>doc=inspect.cleandoc(distribution_cls.__init__.__doc__ or '<STR_LIT>'))<EOL>def func(*args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>sample_shape = kwargs.pop('<STR_LIT>', ())<EOL>value = kwargs.pop('<STR_LIT:value>', None)<EOL>return RandomVariable(distribution=distribution_cls(*args, **kwargs),<EOL>sample_shape=sample_shape,<EOL>value=value)<EOL><DEDENT>return func<EOL>

Factory function to make random variable given distribution class.


@interceptable<EOL>def _build_custom_rv(distribution, sample_shape, value, name):

<EOL>del name  <EOL>return RandomVariable(distribution=distribution,<EOL>sample_shape=sample_shape,<EOL>value=value)<EOL>

RandomVariable constructor with a dummy name argument.


simple_name = distribution.name<EOL>if simple_name.endswith('<STR_LIT:/>'):<EOL><INDENT>simple_name = simple_name.split('<STR_LIT:/>')[-<NUM_LIT:2>]<EOL><DEDENT>parts = simple_name.split('<STR_LIT:_>')<EOL>if parts[-<NUM_LIT:1>].isdigit():<EOL><INDENT>simple_name = '<STR_LIT:_>'.join(parts[:-<NUM_LIT:1>])<EOL><DEDENT>return simple_name<EOL>

Infer the original name passed into a distribution constructor.

    Distributions typically follow the pattern of
    with.name_scope(name) as name:
      super(name=name)
    so we attempt to reverse the name-scope transformation to allow
    addressing of RVs by the distribution's original, user-visible
    name kwarg.

    Args:
      distribution: a tfd.Distribution instance.
    Returns:
      simple_name: the original name passed into the Distribution.

    #### Example

    ```
    d1 = tfd.Normal(0., 1., name='x') # d1.name = 'x/'
    d2 = tfd.Normal(0., 1., name='x') # d2.name = 'x_2/'
    _simple_name(d2) # returns 'x'

    ```


return self.value.shape<EOL>


if self._value is None:<EOL><INDENT>try:<EOL><INDENT>self._value = self.distribution.sample(self.sample_shape_tensor())<EOL><DEDENT>except NotImplementedError:<EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>"<EOL>.format(self.distribution.__class__.__name__))<EOL><DEDENT><DEDENT>return self._value<EOL>

Get tensor that the random variable corresponds to.


return self._distribution<EOL>


def eval(self, session=None, feed_dict=None):

return self.value.eval(session=session, feed_dict=feed_dict)<EOL>

In a session, computes and returns the value of this random variable.

        This is not a graph construction method, it does not add ops to the graph.

        This convenience method requires a session where the graph
        containing this variable has been launched. If no session is
        passed, the default session is used.

        Args:
          session: tf.BaseSession.
            The `tf.Session` to use to evaluate this random variable. If
            none, the default session is used.
          feed_dict: dict.
            A dictionary that maps `tf.Tensor` objects to feed values. See
            `tf.Session.run()` for a description of the valid feed values.

        Returns:
          Value of the random variable.

        #### Examples

        ```python
        x = Normal(0.0, 1.0)
        with tf.Session() as sess:
          # Usage passing the session explicitly.
          print(x.eval(sess))
          # Usage with the default session.  The 'with' block
          # above makes 'sess' the default session.
          print(x.eval())
        ```


def set_values(f, *args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>name = kwargs.get("<STR_LIT:name>")<EOL>if name in model_kwargs:<EOL><INDENT>kwargs["<STR_LIT:value>"] = model_kwargs[name]<EOL><DEDENT>return interceptable(f)(*args, **kwargs)<EOL><DEDENT>return set_values<EOL>

Creates a value-setting interceptor.

    This function creates an interceptor that sets values of Edward2 random
    variable objects. This is useful for a range of tasks, including conditioning
    on observed data, sampling from posterior predictive distributions, and as a
    building block of inference primitives such as computing log joint
    probabilities (see examples below).

    Args:
      **model_kwargs: dict of str to Tensor. Keys are the names of random
        variables in the model to which this interceptor is being applied. Values
        are Tensors to set their value to. Variables not included in this dict
        will not be set and will maintain their existing value semantics (by
        default, a sample from the parent-conditional distribution).

    Returns:
      set_values: function that sets the value of intercepted ops.

    #### Examples

    Consider for illustration a model with latent `z` and
    observed `x`, and a corresponding trainable posterior model:

    ```python
    num_observations = 10
    def model():
      z = ed.Normal(loc=0, scale=1., name='z')  # log rate
      x = ed.Poisson(rate=tf.exp(z) * tf.ones(num_observations), name='x')
      return x

    def variational_model():
      return ed.Normal(loc=tf.Variable(0.),
                       scale=tf.nn.softplus(tf.Variable(-4.)),
                       name='z')  # for simplicity, match name of the model RV.
    ```

    We can use a value-setting interceptor to condition the model on observed
    data. This approach is slightly more cumbersome than that of partially
    evaluating the complete log-joint function, but has the potential advantage
    that it returns a new model callable, which may be used to sample downstream
    variables, passed into additional transformations, etc.

    ```python
    x_observed = np.array([6, 3, 1, 8, 7, 0, 6, 4, 7, 5])
    def observed_model():
      with ed.interception(make_value_setter(x=x_observed)):
        model()
    observed_log_joint_fn = ed.make_log_joint_fn(observed_model)

    # After fixing 'x', the observed log joint is now only a function of 'z'.
    # This enables us to define a variational lower bound,
    # `E_q[ log p(x, z) - log q(z)]`, simply by evaluating the observed and
    # variational log joints at variational samples.
    variational_log_joint_fn = ed.make_log_joint_fn(variational_model)
    with ed.tape() as variational_sample:  # Sample trace from variational model.
      variational_model()
    elbo_loss = -(observed_log_joint_fn(**variational_sample) -
                  variational_log_joint_fn(**variational_sample))
    ```

    After performing inference by minimizing the variational loss, a value-setting
    interceptor enables simulation from the posterior predictive distribution:

    ```python
    with ed.tape() as posterior_samples:  # tape is a map {rv.name : rv}
      variational_model()
    with ed.interception(ed.make_value_setter(**posterior_samples)):
      x = model()
    # x is a sample from p(X | Z = z') where z' ~ q(z) (the variational model)
    ```

    As another example, using a value setter inside of `ed.tape` enables
    computing the log joint probability, by setting all variables to
    posterior values and then accumulating the log probs of those values under
    the induced parent-conditional distributions. This is one way that we could
    have implemented `ed.make_log_joint_fn`:

    ```python
    def make_log_joint_fn_demo(model):
      def log_joint_fn(**model_kwargs):
        with ed.tape() as model_tape:
          with ed.make_value_setter(**model_kwargs):
            model()

        # accumulate sum_i log p(X_i = x_i | X_{:i-1} = x_{:i-1})
        log_prob = 0.
        for rv in model_tape.values():
          log_prob += tf.reduce_sum(rv.log_prob(rv.value))

        return log_prob
      return log_joint_fn
    ```


fn = getattr(cls, attr)<EOL>is_property = isinstance(fn, property)<EOL>if is_property:<EOL><INDENT>fn = fn.fget<EOL><DEDENT>@functools.wraps(fn)<EOL>def wrapped(self, *args, **kwargs):<EOL><INDENT>return fn(self._value(), *args, **kwargs)  <EOL><DEDENT>return property(wrapped) if is_property else wrapped<EOL>

Replaces member function's first arg, `self`, to `self._value()`.

    This function is used by `_get_tensor_like_attributes` to take existing
    `Tensor` member functions and make them operate on `self._value()`, i.e., the
    concretization of a `Distribution`.

    Args:
      cls: The `class` from which we will look up the `attr`.
      attr: Python `str` representing the `attr` to inject a new notion of `self`.

    Returns:
      dependency_injected_function: Python `callable` (or `property`)
        corresponding to `cls.attr` with `self` replaced as `self._value()`.


<EOL>attrs = dict()<EOL>attrs.update((attr, _wrap_method(tf.Tensor, attr))<EOL>for attr in tf.Tensor.OVERLOADABLE_OPERATORS.union({'<STR_LIT>'}))<EOL>attrs.update((attr, getattr(tf.Tensor, attr))<EOL>for attr in {'<STR_LIT>', '<STR_LIT>', '<STR_LIT>'})<EOL>return attrs<EOL>

Returns `Tensor` attributes related to shape and Python builtins.


def __init__(self, initializers, sizes, validate_args=False):

self._initializers = initializers<EOL>self._sizes = sizes<EOL>self._validate_args = validate_args<EOL>

Creates the `BlockwiseInitializer`.

        Arguments:
          initializers: `list` of Keras initializers, e.g., `"glorot_uniform"` or
            `tf.keras.initializers.Constant(0.5413)`.
          sizes: `list` of `int` scalars representing the number of elements
            associated with each initializer in `initializers`.
          validate_args: Python `bool` indicating we should do (possibly expensive)
            graph-time assertions, if necessary.


return {<EOL>'<STR_LIT>': [<EOL>tf.compat.v2.initializers.serialize(<EOL>tf.keras.initializers.get(init))<EOL>for init in self.initializers<EOL>],<EOL>'<STR_LIT>': self.sizes,<EOL>'<STR_LIT>': self.validate_args,<EOL>}<EOL>

Returns initializer configuration as a JSON-serializable dict.


super(CategoricalMixtureOfOneHotCategorical, self).__init__(<EOL>lambda t: CategoricalMixtureOfOneHotCategorical.new(<EOL>t, event_size, num_components, sample_dtype, validate_args),<EOL>convert_to_tensor_fn,<EOL>**kwargs)<EOL>

Initialize the `CategoricalMixtureOfOneHotCategorical` layer.

        Args:
          event_size: Scalar `int` representing the size of single draw from this
            distribution.
          num_components: Scalar `int` representing the number of mixture
            components. Must be at least 1. (If `num_components=1`, it's more
            efficient to use the `OneHotCategorical` layer.)
          convert_to_tensor_fn: Python `callable` that takes a `tfd.Distribution`
            instance and returns a `tf.Tensor`-like object. For examples, see
            `class` docstring.
            Default value: `tfd.Distribution.sample`.
          sample_dtype: `dtype` of samples produced by this distribution.
            Default value: `None` (i.e., previous layer's `dtype`).
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          **kwargs: Additional keyword arguments passed to `tf.keras.Layer`.


@staticmethod<EOL><INDENT>def new(params, event_size, validate_args=False, name=None):<DEDENT>

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[params, event_size]):<EOL><INDENT>params = tf.convert_to_tensor(value=params, name='<STR_LIT>')<EOL>scale_tril = tfb.ScaleTriL(<EOL>diag_shift=np.array(<NUM_LIT>, params.dtype.as_numpy_dtype()),<EOL>validate_args=validate_args)<EOL>return tfd.MultivariateNormalTriL(<EOL>loc=params[..., :event_size],<EOL>scale_tril=scale_tril(params[..., event_size:]),<EOL>validate_args=validate_args)<EOL><DEDENT>

Create the distribution instance from a `params` vector.


@staticmethod<EOL><INDENT>def new(params, event_shape=(), validate_args=False, name=None):<DEDENT>

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[params, event_shape]):<EOL><INDENT>params = tf.convert_to_tensor(value=params, name='<STR_LIT>')<EOL>event_shape = dist_util.expand_to_vector(<EOL>tf.convert_to_tensor(<EOL>value=event_shape, name='<STR_LIT>', dtype_hint=tf.int32),<EOL>tensor_name='<STR_LIT>')<EOL>output_shape = tf.concat([<EOL>tf.shape(input=params)[:-<NUM_LIT:1>],<EOL>event_shape,<EOL>],<EOL>axis=<NUM_LIT:0>)<EOL>loc_params, scale_params = tf.split(params, <NUM_LIT:2>, axis=-<NUM_LIT:1>)<EOL>return tfd.Independent(<EOL>tfd.Logistic(<EOL>loc=tf.reshape(loc_params, output_shape),<EOL>scale=tf.math.softplus(tf.reshape(scale_params, output_shape)),<EOL>validate_args=validate_args),<EOL>reinterpreted_batch_ndims=tf.size(input=event_shape),<EOL>validate_args=validate_args)<EOL><DEDENT>


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[event_shape]):<EOL><INDENT>event_shape = tf.convert_to_tensor(<EOL>value=event_shape, name='<STR_LIT>', dtype_hint=tf.int32)<EOL>return <NUM_LIT:2> * _event_size(<EOL>event_shape, name=name or '<STR_LIT>')<EOL><DEDENT>

The number of `params` needed to create a single distribution.


config = {<EOL>'<STR_LIT>': self._event_shape,<EOL>'<STR_LIT>': _serialize(self._convert_to_tensor_fn),<EOL>'<STR_LIT>': self._sample_dtype,<EOL>'<STR_LIT>': self._validate_args<EOL>}<EOL>base_config = super(IndependentBernoulli, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>

Returns the config of this layer.

        NOTE: At the moment, this configuration can only be serialized if the
        Layer's `convert_to_tensor_fn` is a serializable Keras object (i.e.,
        implements `get_config`) or one of the standard values:
         - `Distribution.sample` (or `"sample"`)
         - `Distribution.mean` (or `"mean"`)
         - `Distribution.mode` (or `"mode"`)
         - `Distribution.stddev` (or `"stddev"`)
         - `Distribution.variance` (or `"variance"`)


convert_to_tensor_fn = _get_convert_to_tensor_fn(convert_to_tensor_fn)<EOL>kwargs.pop('<STR_LIT>', None)<EOL>super(IndependentLogistic, self).__init__(<EOL>lambda t: IndependentLogistic.new(t, event_shape, validate_args),<EOL>convert_to_tensor_fn,<EOL>**kwargs)<EOL>self._event_shape = event_shape<EOL>self._convert_to_tensor_fn = convert_to_tensor_fn<EOL>self._validate_args = validate_args<EOL>

Initialize the `IndependentLogistic` layer.

        Args:
          event_shape: integer vector `Tensor` representing the shape of single
            draw from this distribution.
          convert_to_tensor_fn: Python `callable` that takes a `tfd.Distribution`
            instance and returns a `tf.Tensor`-like object.
            Default value: `tfd.Distribution.sample`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          **kwargs: Additional keyword arguments passed to `tf.keras.Layer`.


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[params, event_shape]):<EOL><INDENT>params = tf.convert_to_tensor(value=params, name='<STR_LIT>')<EOL>event_shape = dist_util.expand_to_vector(<EOL>tf.convert_to_tensor(<EOL>value=event_shape, name='<STR_LIT>', dtype_hint=tf.int32),<EOL>tensor_name='<STR_LIT>')<EOL>output_shape = tf.concat([<EOL>tf.shape(input=params)[:-<NUM_LIT:1>],<EOL>event_shape,<EOL>],<EOL>axis=<NUM_LIT:0>)<EOL>return tfd.Independent(<EOL>tfd.Poisson(<EOL>log_rate=tf.reshape(params, output_shape),<EOL>validate_args=validate_args),<EOL>reinterpreted_batch_ndims=tf.size(input=event_shape),<EOL>validate_args=validate_args)<EOL><DEDENT>


@staticmethod<EOL><INDENT>def params_size(num_components, event_shape=(), name=None):<DEDENT>

return MixtureSameFamily.params_size(<EOL>num_components,<EOL>IndependentNormal.params_size(event_shape, name=name),<EOL>name=name)<EOL>


convert_to_tensor_fn = _get_convert_to_tensor_fn(convert_to_tensor_fn)<EOL>kwargs.pop('<STR_LIT>', None)<EOL>super(IndependentPoisson, self).__init__(<EOL>lambda t: IndependentPoisson.new(t, event_shape, validate_args),<EOL>convert_to_tensor_fn,<EOL>**kwargs)<EOL>self._event_shape = event_shape<EOL>self._convert_to_tensor_fn = convert_to_tensor_fn<EOL>self._validate_args = validate_args<EOL>

Initialize the `IndependentPoisson` layer.

        Args:
          event_shape: integer vector `Tensor` representing the shape of single
            draw from this distribution.
          convert_to_tensor_fn: Python `callable` that takes a `tfd.Distribution`
            instance and returns a `tf.Tensor`-like object.
            Default value: `tfd.Distribution.sample`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          **kwargs: Additional keyword arguments passed to `tf.keras.Layer`.


convert_to_tensor_fn = _get_convert_to_tensor_fn(convert_to_tensor_fn)<EOL>kwargs.pop('<STR_LIT>', None)<EOL>super(IndependentBernoulli, self).__init__(<EOL>lambda t: IndependentBernoulli.new(  <EOL>t, event_shape, sample_dtype, validate_args),<EOL>convert_to_tensor_fn,<EOL>**kwargs)<EOL>self._event_shape = event_shape<EOL>self._convert_to_tensor_fn = convert_to_tensor_fn<EOL>self._sample_dtype = sample_dtype<EOL>self._validate_args = validate_args<EOL>

Initialize the `IndependentBernoulli` layer.

        Args:
          event_shape: integer vector `Tensor` representing the shape of single
            draw from this distribution.
          convert_to_tensor_fn: Python `callable` that takes a `tfd.Distribution`
            instance and returns a `tf.Tensor`-like object. For examples, see
            `class` docstring.
            Default value: `tfd.Distribution.sample`.
          sample_dtype: `dtype` of samples produced by this distribution.
            Default value: `None` (i.e., previous layer's `dtype`).
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          **kwargs: Additional keyword arguments passed to `tf.keras.Layer`.


@staticmethod<EOL><INDENT>def new(params, event_size, dtype=None, validate_args=False, name=None):<DEDENT>

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[params, event_size]):<EOL><INDENT>return tfd.OneHotCategorical(<EOL>logits=params,<EOL>dtype=dtype or params.dtype.base_dtype,<EOL>validate_args=validate_args)<EOL><DEDENT>


config = {<EOL>'<STR_LIT>': self._event_shape,<EOL>'<STR_LIT>': _serialize(self._convert_to_tensor_fn),<EOL>'<STR_LIT>': self._validate_args<EOL>}<EOL>base_config = super(IndependentPoisson, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[event_shape]):<EOL><INDENT>event_shape = tf.convert_to_tensor(<EOL>value=event_shape, name='<STR_LIT>', dtype_hint=tf.int32)<EOL>return _event_size(<EOL>event_shape, name=name or '<STR_LIT>')<EOL><DEDENT>


@staticmethod<EOL><INDENT>def new(params, num_components, component_layer,<EOL>validate_args=False, name=None):<DEDENT>

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[params, num_components, component_layer]):<EOL><INDENT>params = tf.convert_to_tensor(value=params, name='<STR_LIT>')<EOL>num_components = tf.convert_to_tensor(<EOL>value=num_components, name='<STR_LIT>', dtype_hint=tf.int32)<EOL>components_dist = component_layer(<EOL>tf.reshape(<EOL>params[..., num_components:],<EOL>tf.concat([tf.shape(input=params)[:-<NUM_LIT:1>], [num_components, -<NUM_LIT:1>]],<EOL>axis=<NUM_LIT:0>)))<EOL>mixture_dist = tfd.Categorical(logits=params[..., :num_components])<EOL>return tfd.MixtureSameFamily(<EOL>mixture_dist,<EOL>components_dist,<EOL>validate_args=False)<EOL><DEDENT>


Construct layer.

        Args:
          filters: Integer, the dimensionality of the output space (i.e. the number
            of filters in the convolution).
          kernel_size: An integer or tuple/list of a single integer, specifying the
            length of the 1D convolution window.
          strides: An integer or tuple/list of a single integer,
            specifying the stride length of the convolution.
            Specifying any stride value != 1 is incompatible with specifying
            any `dilation_rate` value != 1.
          padding: One of `"valid"` or `"same"` (case-insensitive).
          data_format: A string, one of `channels_last` (default) or
            `channels_first`. The ordering of the dimensions in the inputs.
            `channels_last` corresponds to inputs with shape `(batch, length,
            channels)` while `channels_first` corresponds to inputs with shape
            `(batch, channels, length)`.
          dilation_rate: An integer or tuple/list of a single integer, specifying
            the dilation rate to use for dilated convolution.
            Currently, specifying any `dilation_rate` value != 1 is
            incompatible with specifying any `strides` value != 1.
          ${args}
          seed: Python scalar `int` which initializes the random number
            generator. Default value: `None` (i.e., use global seed).


Construct layer.

        Args:
          filters: Integer, the dimensionality of the output space (i.e. the number
            of filters in the convolution).
          kernel_size: An integer or tuple/list of 3 integers, specifying the
            depth, height and width of the 3D convolution window.
            Can be a single integer to specify the same value for
            all spatial dimensions.
          strides: An integer or tuple/list of 3 integers,
            specifying the strides of the convolution along the depth,
            height and width.
            Can be a single integer to specify the same value for
            all spatial dimensions.
            Specifying any stride value != 1 is incompatible with specifying
            any `dilation_rate` value != 1.
          padding: One of `"valid"` or `"same"` (case-insensitive).
          data_format: A string, one of `channels_last` (default) or
            `channels_first`. The ordering of the dimensions in the inputs.
            `channels_last` corresponds to inputs with shape `(batch, depth,
            height, width, channels)` while `channels_first` corresponds to inputs
            with shape `(batch, channels, depth, height, width)`.
          dilation_rate: An integer or tuple/list of 3 integers, specifying
            the dilation rate to use for dilated convolution.
            Can be a single integer to specify the same value for
            all spatial dimensions.
            Currently, specifying any `dilation_rate` value != 1 is
            incompatible with specifying any stride value != 1.
          ${args}
          seed: Python scalar `int` which initializes the random number
            generator. Default value: `None` (i.e., use global seed).


Construct layer.

        Args:
          filters: Integer, the dimensionality of the output space (i.e. the number
            of filters in the convolution).
          kernel_size: An integer or tuple/list of 3 integers, specifying the
            depth, height and width of the 3D convolution window.
            Can be a single integer to specify the same value for
            all spatial dimensions.
          strides: An integer or tuple/list of 3 integers,
            specifying the strides of the convolution along the depth,
            height and width.
            Can be a single integer to specify the same value for
            all spatial dimensions.
            Specifying any stride value != 1 is incompatible with specifying
            any `dilation_rate` value != 1.
          padding: One of `"valid"` or `"same"` (case-insensitive).
          data_format: A string, one of `channels_last` (default) or
            `channels_first`. The ordering of the dimensions in the inputs.
            `channels_last` corresponds to inputs with shape `(batch, depth,
            height, width, channels)` while `channels_first` corresponds to inputs
            with shape `(batch, channels, depth, height, width)`.
          dilation_rate: An integer or tuple/list of 3 integers, specifying
            the dilation rate to use for dilated convolution.
            Can be a single integer to specify the same value for
            all spatial dimensions.
            Currently, specifying any `dilation_rate` value != 1 is
            incompatible with specifying any stride value != 1.
          ${args}


input_shape = tf.TensorShape(input_shape).as_list()<EOL>if self.data_format == '<STR_LIT>':<EOL><INDENT>space = input_shape[<NUM_LIT:1>:-<NUM_LIT:1>]<EOL>new_space = []<EOL>for i in range(len(space)):<EOL><INDENT>new_dim = tf_layers_util.conv_output_length(<EOL>space[i],<EOL>self.kernel_size[i],<EOL>padding=self.padding,<EOL>stride=self.strides[i],<EOL>dilation=self.dilation_rate[i])<EOL>new_space.append(new_dim)<EOL><DEDENT>return tf.TensorShape([input_shape[<NUM_LIT:0>]] + new_space + [self.filters])<EOL><DEDENT>else:<EOL><INDENT>space = input_shape[<NUM_LIT:2>:]<EOL>new_space = []<EOL>for i in range(len(space)):<EOL><INDENT>new_dim = tf_layers_util.conv_output_length(<EOL>space[i],<EOL>self.kernel_size[i],<EOL>padding=self.padding,<EOL>stride=self.strides[i],<EOL>dilation=self.dilation_rate[i])<EOL>new_space.append(new_dim)<EOL><DEDENT>return tf.TensorShape([input_shape[<NUM_LIT:0>], self.filters] + new_space)<EOL><DEDENT>

Computes the output shape of the layer.

        Args:
          input_shape: Shape tuple (tuple of integers) or list of shape tuples
            (one per output tensor of the layer). Shape tuples can include None for
            free dimensions, instead of an integer.

        Returns:
          output_shape: A tuple representing the output shape.


def _distribution_to_params(self, distribution, batch_shape):

params_shape = tf.concat([batch_shape, [-<NUM_LIT:1>]], axis=<NUM_LIT:0>)<EOL>batch_and_n_shape = tf.concat(<EOL>[tf.shape(input=distribution.mixture_distribution.logits), [-<NUM_LIT:1>]],<EOL>axis=<NUM_LIT:0>)<EOL>cd = distribution.components_distribution.distribution<EOL>return tf.concat([<EOL>distribution.mixture_distribution.logits,<EOL>tf.reshape(tf.concat([<EOL>tf.reshape(cd.loc, batch_and_n_shape),<EOL>tf.reshape(tfd.softplus_inverse(cd.scale), batch_and_n_shape)<EOL>], axis=-<NUM_LIT:1>), params_shape),<EOL>], axis=-<NUM_LIT:1>)<EOL>

Given a self.layer_class instance, return a tensor of its parameters.


def assertExportable(self, model, batch_size=<NUM_LIT:1>):

batch_shape = [batch_size]<EOL>input_shape = batch_shape + model.input.shape[<NUM_LIT:1>:].as_list()<EOL>dtype = model.input.dtype.as_numpy_dtype()<EOL>model_dir = self.create_tempdir()<EOL>tfk.experimental.export_saved_model(model, model_dir.full_path)<EOL>model_copy = tfk.experimental.load_from_saved_model(model_dir.full_path)<EOL>x = np.random.uniform(-<NUM_LIT>, <NUM_LIT>, input_shape).astype(dtype)<EOL>self.assertAllEqual(self.evaluate(model(x)), self.evaluate(model_copy(x)))<EOL>self.assertAllEqual(model.predict(x), model_copy.predict(x))<EOL>

Assert a Keras model supports export_saved_model/load_from_saved_model.

        Arguments:
          model: A Keras model with Tensor output.
          batch_size: The batch size to use when checking that the model produces
            the same results as a serialized/deserialized copy.  Default value: 1.


def assertSerializable(self, model, batch_size=<NUM_LIT:1>):

batch_shape = [batch_size]<EOL>input_shape = batch_shape + model.input.shape[<NUM_LIT:1>:].as_list()<EOL>dtype = model.input.dtype.as_numpy_dtype()<EOL>model_file = self.create_tempfile()<EOL>model.save(model_file.full_path)<EOL>model_copy = tfk.models.load_model(model_file.full_path)<EOL>x = np.random.uniform(-<NUM_LIT>, <NUM_LIT>, input_shape).astype(dtype)<EOL>model_x_mean = self.evaluate(model(x).mean())<EOL>self.assertAllEqual(model_x_mean, self.evaluate(model_copy(x).mean()))<EOL>output_shape = model_x_mean.shape<EOL>y = np.random.uniform(<NUM_LIT:0.>, <NUM_LIT:1.>, output_shape).astype(dtype)<EOL>self.assertAllEqual(self.evaluate(model(x).log_prob(y)),<EOL>self.evaluate(model_copy(x).log_prob(y)))<EOL>

Assert that a model can be saved/loaded via Keras Model.save/load_model.

        Arguments:
          model: A Keras model that outputs a `tfd.Distribution`.
          batch_size: The batch size to use when checking that the model produces
            the same results as a serialized/deserialized copy.  Default value: 1.


config = {<EOL>'<STR_LIT>': self.seed,<EOL>}<EOL>base_config = super(DenseFlipout, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>

Returns the config of the layer.

        A layer config is a Python dictionary (serializable) containing the
        configuration of a layer. The same layer can be reinstantiated later
        (without its trained weights) from this configuration.

        Returns:
          config: A Python dictionary of class keyword arguments and their
            serialized values.


@docstring_util.expand_docstring(args=doc_args)<EOL><INDENT>def __init__(<EOL>self,<EOL>units,<EOL>activation=None,<EOL>activity_regularizer=None,<EOL>kernel_posterior_fn=tfp_layers_util.default_mean_field_normal_fn(),<EOL>kernel_posterior_tensor_fn=lambda d: d.sample(),<EOL>kernel_prior_fn=tfp_layers_util.default_multivariate_normal_fn,<EOL>kernel_divergence_fn=lambda q, p, ignore: tfd.kl_divergence(q, p),<EOL>bias_posterior_fn=tfp_layers_util.default_mean_field_normal_fn(is_singular=True),  <EOL>bias_posterior_tensor_fn=lambda d: d.sample(),<EOL>bias_prior_fn=None,<EOL>bias_divergence_fn=lambda q, p, ignore: tfd.kl_divergence(q, p),<EOL>**kwargs):<EOL><DEDENT>

<EOL>super(_DenseVariational, self).__init__(<EOL>activity_regularizer=activity_regularizer,<EOL>**kwargs)<EOL>self.units = units<EOL>self.activation = tf.keras.activations.get(activation)<EOL>self.input_spec = tf.keras.layers.InputSpec(min_ndim=<NUM_LIT:2>)<EOL>self.kernel_posterior_fn = kernel_posterior_fn<EOL>self.kernel_posterior_tensor_fn = kernel_posterior_tensor_fn<EOL>self.kernel_prior_fn = kernel_prior_fn<EOL>self.kernel_divergence_fn = kernel_divergence_fn<EOL>self.bias_posterior_fn = bias_posterior_fn<EOL>self.bias_posterior_tensor_fn = bias_posterior_tensor_fn<EOL>self.bias_prior_fn = bias_prior_fn<EOL>self.bias_divergence_fn = bias_divergence_fn<EOL>

Construct layer.

        Args:
          ${args}


config = {<EOL>'<STR_LIT>': self.units,<EOL>'<STR_LIT>': (tf.keras.activations.serialize(self.activation)<EOL>if self.activation else None),<EOL>'<STR_LIT>':<EOL>tf.keras.initializers.serialize(self.activity_regularizer),<EOL>}<EOL>function_keys = [<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>]<EOL>for function_key in function_keys:<EOL><INDENT>function = getattr(self, function_key)<EOL>if function is None:<EOL><INDENT>function_name = None<EOL>function_type = None<EOL><DEDENT>else:<EOL><INDENT>function_name, function_type = tfp_layers_util.serialize_function(<EOL>function)<EOL><DEDENT>config[function_key] = function_name<EOL>config[function_key + '<STR_LIT>'] = function_type<EOL><DEDENT>base_config = super(_DenseVariational, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>


graph_deps = [model_matrix, response, model_coefficients_start,<EOL>predicted_linear_response_start, dispersion, offset,<EOL>learning_rate, maximum_iterations]<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', graph_deps):<EOL><INDENT>[<EOL>model_matrix,<EOL>response,<EOL>model_coefficients_start,<EOL>predicted_linear_response_start,<EOL>offset,<EOL>] = prepare_args(<EOL>model_matrix,<EOL>response,<EOL>model_coefficients_start,<EOL>predicted_linear_response_start,<EOL>offset)<EOL>if convergence_criteria_fn is None:<EOL><INDENT>convergence_criteria_fn = (<EOL>convergence_criteria_small_relative_norm_weights_change())<EOL><DEDENT>def _body(<EOL>is_converged_previous,<EOL>iter_,<EOL>model_coefficients_previous,<EOL>predicted_linear_response_previous):<EOL><INDENT>"""<STR_LIT>"""<EOL>model_coefficients_next, predicted_linear_response_next = fit_one_step(<EOL>model_matrix,<EOL>response,<EOL>model,<EOL>model_coefficients_previous,<EOL>predicted_linear_response_previous,<EOL>l2_regularizer,<EOL>dispersion,<EOL>offset,<EOL>learning_rate,<EOL>fast_unsafe_numerics)<EOL>is_converged_next = convergence_criteria_fn(<EOL>is_converged_previous=is_converged_previous,<EOL>iter_=iter_,<EOL>model_coefficients_previous=model_coefficients_previous,<EOL>predicted_linear_response_previous=predicted_linear_response_previous,<EOL>model_coefficients_next=model_coefficients_next,<EOL>predicted_linear_response_next=predicted_linear_response_next,<EOL>response=response,<EOL>model=model,<EOL>dispersion=dispersion)<EOL>return [<EOL>is_converged_next,<EOL>iter_ + <NUM_LIT:1>,<EOL>model_coefficients_next,<EOL>predicted_linear_response_next,<EOL>]<EOL><DEDENT>[<EOL>is_converged,<EOL>iter_,<EOL>model_coefficients,<EOL>predicted_linear_response,<EOL>] = tf.while_loop(<EOL>cond=lambda is_converged, *args: tf.logical_not(is_converged),<EOL>body=_body,<EOL>loop_vars=[<EOL>tf.zeros([], np.bool),   <EOL>tf.zeros([], np.int32),  <EOL>model_coefficients_start,<EOL>predicted_linear_response_start,<EOL>],<EOL>maximum_iterations=maximum_iterations)<EOL>return [<EOL>model_coefficients,<EOL>predicted_linear_response,<EOL>is_converged,<EOL>iter_<EOL>]<EOL><DEDENT>

Runs multiple Fisher scoring steps.

    Args:
      model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
        represents a sample's features.
      response: (Batch of) vector-shaped `Tensor` where each element represents a
        sample's observed response (to the corresponding row of features). Must
        have same `dtype` as `model_matrix`.
      model: `tfp.glm.ExponentialFamily`-like instance which implicitly
        characterizes a negative log-likelihood loss by specifying the
        distribuion's `mean`, `gradient_mean`, and `variance`.
      model_coefficients_start: Optional (batch of) vector-shaped `Tensor`
        representing the initial model coefficients, one for each column in
        `model_matrix`. Must have same `dtype` as `model_matrix`.
        Default value: Zeros.
      predicted_linear_response_start: Optional `Tensor` with `shape`, `dtype`
        matching `response`; represents `offset` shifted initial linear
        predictions based on `model_coefficients_start`.
        Default value: `offset` if `model_coefficients is None`, and
        `tf.linalg.matvec(model_matrix, model_coefficients_start) + offset`
        otherwise.
      l2_regularizer: Optional scalar `Tensor` representing L2 regularization
        penalty, i.e.,
        `loss(w) = sum{-log p(y[i]|x[i],w) : i=1..n} + l2_regularizer ||w||_2^2`.
        Default value: `None` (i.e., no L2 regularization).
      dispersion: Optional (batch of) `Tensor` representing `response` dispersion,
        i.e., as in, `p(y|theta) := exp((y theta - A(theta)) / dispersion)`.
        Must broadcast with rows of `model_matrix`.
        Default value: `None` (i.e., "no dispersion").
      offset: Optional `Tensor` representing constant shift applied to
        `predicted_linear_response`.  Must broadcast to `response`.
        Default value: `None` (i.e., `tf.zeros_like(response)`).
      convergence_criteria_fn: Python `callable` taking:
        `is_converged_previous`, `iter_`, `model_coefficients_previous`,
        `predicted_linear_response_previous`, `model_coefficients_next`,
        `predicted_linear_response_next`, `response`, `model`, `dispersion` and
        returning a `bool` `Tensor` indicating that Fisher scoring has converged.
        See `convergence_criteria_small_relative_norm_weights_change` as an
        example function.
        Default value: `None` (i.e.,
        `convergence_criteria_small_relative_norm_weights_change`).
      learning_rate: Optional (batch of) scalar `Tensor` used to dampen iterative
        progress. Typically only needed if optimization diverges, should be no
        larger than `1` and typically very close to `1`.
        Default value: `None` (i.e., `1`).
      fast_unsafe_numerics: Optional Python `bool` indicating if faster, less
        numerically accurate methods can be employed for computing the weighted
        least-squares solution.
        Default value: `True` (i.e., "fast but possibly diminished accuracy").
      maximum_iterations: Optional maximum number of iterations of Fisher scoring
        to run; "and-ed" with result of `convergence_criteria_fn`.
        Default value: `None` (i.e., `infinity`).
      name: Python `str` used as name prefix to ops created by this function.
        Default value: `"fit"`.

    Returns:
      model_coefficients: (Batch of) vector-shaped `Tensor`; represents the
        fitted model coefficients, one for each column in `model_matrix`.
      predicted_linear_response: `response`-shaped `Tensor` representing linear
        predictions based on new `model_coefficients`, i.e.,
        `tf.linalg.matvec(model_matrix, model_coefficients) + offset`.
      is_converged: `bool` `Tensor` indicating that the returned
        `model_coefficients` met the `convergence_criteria_fn` criteria within the
        `maximum_iterations` limit.
      iter_: `int32` `Tensor` indicating the number of iterations taken.

    #### Example

    ```python
    from __future__ import print_function
    import numpy as np
    import tensorflow as tf
    import tensorflow_probability as tfp
    tfd = tfp.distributions

    def make_dataset(n, d, link, scale=1., dtype=np.float32):
      model_coefficients = tfd.Uniform(
          low=np.array(-1, dtype),
          high=np.array(1, dtype)).sample(d, seed=42)
      radius = np.sqrt(2.)
      model_coefficients *= radius / tf.linalg.norm(model_coefficients)
      model_matrix = tfd.Normal(
          loc=np.array(0, dtype),
          scale=np.array(1, dtype)).sample([n, d], seed=43)
      scale = tf.convert_to_tensor(scale, dtype)
      linear_response = tf.tensordot(
          model_matrix, model_coefficients, axes=[[1], [0]])
      if link == 'linear':
        response = tfd.Normal(loc=linear_response, scale=scale).sample(seed=44)
      elif link == 'probit':
        response = tf.cast(
            tfd.Normal(loc=linear_response, scale=scale).sample(seed=44) > 0,
            dtype)
      elif link == 'logit':
        response = tfd.Bernoulli(logits=linear_response).sample(seed=44)
      else:
        raise ValueError('unrecognized true link: {}'.format(link))
      return model_matrix, response, model_coefficients

    X, Y, w_true = make_dataset(n=int(1e6), d=100, link='probit')

    w, linear_response, is_converged, num_iter = tfp.glm.fit(
        model_matrix=X,
        response=Y,
        model=tfp.glm.BernoulliNormalCDF())
    log_likelihood = tfp.glm.BernoulliNormalCDF().log_prob(Y, linear_response)

    with tf.Session() as sess:
      [w_, linear_response_, is_converged_, num_iter_, Y_, w_true_,
       log_likelihood_] = sess.run([
          w, linear_response, is_converged, num_iter, Y, w_true,
          log_likelihood])

    print('is_converged: ', is_converged_)
    print('    num_iter: ', num_iter_)
    print('    accuracy: ', np.mean((linear_response_ > 0.) == Y_))
    print('    deviance: ', 2. * np.mean(log_likelihood_))
    print('||w0-w1||_2 / (1+||w0||_2): ', (np.linalg.norm(w_true_ - w_, ord=2) /
                                           (1. + np.linalg.norm(w_true_, ord=2))))

    # ==>
    # is_converged:  True
    #     num_iter:  6
    #     accuracy:  0.804382
    #     deviance:  -0.820746600628
    # ||w0-w1||_2 / (1+||w0||_2):  0.00619245105309
    ```


def _log_prob(self, response, predicted_linear_response):

raise NotImplementedError('<STR_LIT>')<EOL>

Implements the log-probability computation.


def log_prob(self, response, predicted_linear_response, name=None):

with self._name_scope(<EOL>name, '<STR_LIT>', [response, predicted_linear_response]):<EOL><INDENT>dtype = dtype_util.common_dtype([response, predicted_linear_response])<EOL>response = tf.convert_to_tensor(<EOL>value=response, dtype=dtype, name='<STR_LIT>')<EOL>predicted_linear_response = tf.convert_to_tensor(<EOL>value=predicted_linear_response, name='<STR_LIT>')<EOL>return self._log_prob(response, predicted_linear_response)<EOL><DEDENT>

Computes `D(param=mean(r)).log_prob(response)` for linear response, `r`.

        Args:
          response: `float`-like `Tensor` representing observed ("actual")
            responses.
          predicted_linear_response: `float`-like `Tensor` corresponding to
            `tf.matmul(model_matrix, weights)`.
          name: Python `str` used as TF namescope for ops created by member
            functions. Default value: `None` (i.e., 'log_prob').

        Returns:
          log_prob: `Tensor` with shape and dtype of `predicted_linear_response`
            representing the distribution prescribed log-probability of the observed
            `response`s.


def tril_with_diag_softplus_and_shift(x, diag_shift=<NUM_LIT>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[x, diag_shift]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>x = tfd.fill_triangular(x)<EOL>diag = softplus_and_shift(tf.linalg.diag_part(x), diag_shift)<EOL>x = tf.linalg.set_diag(x, diag)<EOL>return x<EOL><DEDENT>

Converts (batch of) vectors to (batch of) lower-triangular scale matrices.

    Args:
      x: (Batch of) `float`-like `Tensor` representing vectors which will be
        transformed into lower-triangular scale matrices with positive diagonal
        elements. Rightmost shape `n` must be such that
        `n = dims * (dims + 1) / 2` for some positive, integer `dims`.
      diag_shift: `Tensor` added to `softplus` transformation of diagonal
        elements.
        Default value: `1e-5`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "tril_with_diag_softplus_and_shift").

    Returns:
      scale_tril: (Batch of) lower-triangular `Tensor` with `x.dtype` and
        rightmost shape `[dims, dims]` where `n = dims * (dims + 1) / 2` where
        `n = x.shape[-1]`.


def normal(x,<EOL>layer_fn=tf.compat.v1.layers.dense,<EOL>loc_fn=lambda x: x,<EOL>scale_fn=<NUM_LIT:1.>,<EOL>name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if callable(scale_fn):<EOL><INDENT>y = layer_fn(x, <NUM_LIT:2>)<EOL>loc = loc_fn(y[..., <NUM_LIT:0>])<EOL>scale = scale_fn(y[..., <NUM_LIT:1>])<EOL><DEDENT>else:<EOL><INDENT>y = tf.squeeze(layer_fn(x, <NUM_LIT:1>), axis=-<NUM_LIT:1>)<EOL>loc = loc_fn(y)<EOL>scale = tf.cast(scale_fn, loc.dtype.base_dtype)<EOL><DEDENT>return tfd.Normal(loc=loc, scale=scale)<EOL><DEDENT>

Constructs a trainable `tfd.Normal` distribution.

    This function creates a Normal distribution parameterized by loc and scale.
    Using default args, this function is mathematically equivalent to:

    ```none
    Y = Normal(loc=matmul(W, x) + b, scale=1)

    where,
      W in R^[d, n]
      b in R^d
    ```

    #### Examples

    This function can be used as a [linear regression](
    https://en.wikipedia.org/wiki/Linear_regression) loss.

    ```python
    # This example fits a linear regression loss.
    import tensorflow as tf
    import tensorflow_probability as tfp

    # Create fictitious training data.
    dtype = np.float32
    n = 3000    # number of samples
    x_size = 4  # size of single x
    def make_training_data():
      np.random.seed(142)
      x = np.random.randn(n, x_size).astype(dtype)
      w = np.random.randn(x_size).astype(dtype)
      b = np.random.randn(1).astype(dtype)
      true_mean = np.tensordot(x, w, axes=[[-1], [-1]]) + b
      noise = np.random.randn(n).astype(dtype)
      y = true_mean + noise
      return y, x
    y, x = make_training_data()

    # Build TF graph for fitting Normal maximum likelihood estimator.
    normal = tfp.trainable_distributions.normal(x)
    loss = -tf.reduce_mean(normal.log_prob(y))
    train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
    mse = tf.reduce_mean(tf.squared_difference(y, normal.mean()))
    init_op = tf.global_variables_initializer()

    # Run graph 1000 times.
    num_steps = 1000
    loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
    mse_ = np.zeros(num_steps)
    with tf.Session() as sess:
      sess.run(init_op)
      for it in xrange(loss_.size):
        _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
        if it % 200 == 0 or it == loss_.size - 1:
          print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

    # ==> iteration:0    loss:6.34114170074  mse:10.8444051743
    #     iteration:200  loss:1.40146839619  mse:0.965059816837
    #     iteration:400  loss:1.40052902699  mse:0.963181257248
    #     iteration:600  loss:1.40052902699  mse:0.963181257248
    #     iteration:800  loss:1.40052902699  mse:0.963181257248
    #     iteration:999  loss:1.40052902699  mse:0.963181257248
    ```

    Args:
      x: `Tensor` with floating type. Must have statically defined rank and
        statically known right-most dimension.
      layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
        returns a transformation of `x` with shape
        `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
        Default value: `tf.layers.dense`.
      loc_fn: Python `callable` which transforms the `loc` parameter. Takes a
        (batch of) length-`dims` vectors and returns a `Tensor` of same shape and
        `dtype`.
        Default value: `lambda x: x`.
      scale_fn: Python `callable` or `Tensor`. If a `callable` transforms the
        `scale` parameters; if `Tensor` is the `tfd.Normal` `scale` argument.
        Takes a (batch of) length-`dims` vectors and returns a `Tensor` of same
        size. (Taking a `callable` or `Tensor` is how `tf.Variable` intializers
        behave.)
        Default value: `1`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "normal").

    Returns:
      normal: An instance of `tfd.Normal`.


def bernoulli(x, layer_fn=tf.compat.v1.layers.dense, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>logits = tf.squeeze(layer_fn(x, <NUM_LIT:1>), axis=-<NUM_LIT:1>)<EOL>return tfd.Bernoulli(logits=logits)<EOL><DEDENT>

Constructs a trainable `tfd.Bernoulli` distribution.

    This function creates a Bernoulli distribution parameterized by logits.
    Using default args, this function is mathematically equivalent to:

    ```none
    Y = Bernoulli(logits=matmul(W, x) + b)

    where,
      W in R^[d, n]
      b in R^d
    ```

    #### Examples

    This function can be used as a [logistic regression](
    https://en.wikipedia.org/wiki/Logistic_regression) loss.

    ```python
    # This example fits a logistic regression loss.
    import tensorflow as tf
    import tensorflow_probability as tfp

    # Create fictitious training data.
    dtype = np.float32
    n = 3000    # number of samples
    x_size = 4  # size of single x
    def make_training_data():
      np.random.seed(142)
      x = np.random.randn(n, x_size).astype(dtype)
      w = np.random.randn(x_size).astype(dtype)
      b = np.random.randn(1).astype(dtype)
      true_logits = np.tensordot(x, w, axes=[[-1], [-1]]) + b
      noise = np.random.logistic(size=n).astype(dtype)
      y = dtype(true_logits + noise > 0.)
      return y, x
    y, x = make_training_data()

    # Build TF graph for fitting Bernoulli maximum likelihood estimator.
    bernoulli = tfp.trainable_distributions.bernoulli(x)
    loss = -tf.reduce_mean(bernoulli.log_prob(y))
    train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
    mse = tf.reduce_mean(tf.squared_difference(y, bernoulli.mean()))
    init_op = tf.global_variables_initializer()

    # Run graph 1000 times.
    num_steps = 1000
    loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
    mse_ = np.zeros(num_steps)
    with tf.Session() as sess:
      sess.run(init_op)
      for it in xrange(loss_.size):
        _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
        if it % 200 == 0 or it == loss_.size - 1:
          print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

    # ==> iteration:0    loss:0.635675370693  mse:0.222526371479
    #     iteration:200  loss:0.440077394247  mse:0.143687799573
    #     iteration:400  loss:0.440077394247  mse:0.143687844276
    #     iteration:600  loss:0.440077394247  mse:0.143687844276
    #     iteration:800  loss:0.440077424049  mse:0.143687844276
    #     iteration:999  loss:0.440077424049  mse:0.143687844276
    ```

    Args:
      x: `Tensor` with floating type. Must have statically defined rank and
        statically known right-most dimension.
      layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
        returns a transformation of `x` with shape
        `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
        Default value: `tf.layers.dense`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "bernoulli").

    Returns:
      bernoulli: An instance of `tfd.Bernoulli`.


assertions = []<EOL>if not a.dtype.is_floating:<EOL><INDENT>raise TypeError('<STR_LIT>'<EOL>'<STR_LIT>'.format(a.dtype.name))<EOL><DEDENT>if a.shape.ndims is not None:<EOL><INDENT>if a.shape.ndims < <NUM_LIT:2>:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(a.shape.ndims))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(tf.compat.v1.assert_rank_at_least(<EOL>a, rank=<NUM_LIT:2>, message='<STR_LIT>'))<EOL><DEDENT>return assertions<EOL>


vecs = tf.convert_to_tensor(value=vecs, dtype=tf.int64, name='<STR_LIT>')<EOL>m = tf.convert_to_tensor(value=m, dtype=tf.int64, name='<STR_LIT:m>')<EOL>i = tf.convert_to_tensor(value=i, dtype=tf.int64, name='<STR_LIT:i>')<EOL>trailing_elts = tf.broadcast_to(<EOL>tf.range(m + <NUM_LIT:1>,<EOL>prefer_static.shape(vecs, out_type=tf.int64)[-<NUM_LIT:1>]),<EOL>prefer_static.shape(vecs[..., m + <NUM_LIT:1>:]))<EOL>shp = prefer_static.shape(trailing_elts)<EOL>trailing_elts = tf.where(<EOL>tf.equal(trailing_elts, tf.broadcast_to(i, shp)),<EOL>tf.broadcast_to(tf.gather(vecs, [m], axis=-<NUM_LIT:1>), shp),<EOL>tf.broadcast_to(vecs[..., m + <NUM_LIT:1>:], shp))<EOL>vecs_shape = vecs.shape<EOL>vecs = tf.concat([<EOL>vecs[..., :m],<EOL>tf.gather(vecs, i, batch_dims=prefer_static.rank(vecs) - <NUM_LIT:1>), trailing_elts<EOL>], axis=-<NUM_LIT:1>)<EOL>tensorshape_util.set_shape(vecs, vecs_shape)<EOL>return vecs<EOL>

Swaps `m` and `i` on axis -1. (Helper for pivoted_cholesky.)

    Given a batch of int64 vectors `vecs`, scalar index `m`, and compatibly shaped
    per-vector indices `i`, this function swaps elements `m` and `i` in each
    vector. For the use-case below, these are permutation vectors.

    Args:
      vecs: Vectors on which we perform the swap, int64 `Tensor`.
      m: Scalar int64 `Tensor`, the index into which the `i`th element is going.
      i: Batch int64 `Tensor`, shaped like vecs.shape[:-1] + [1]; the index into
        which the `m`th element is going.

    Returns:
      vecs: The updated vectors.


def pivoted_cholesky(matrix, max_rank, diag_rtol=<NUM_LIT>, name=None):

with tf.compat.v2.name_scope(name or '<STR_LIT>'):<EOL><INDENT>dtype = dtype_util.common_dtype([matrix, diag_rtol],<EOL>preferred_dtype=tf.float32)<EOL>matrix = tf.convert_to_tensor(value=matrix, name='<STR_LIT>', dtype=dtype)<EOL>if tensorshape_util.rank(matrix.shape) is None:<EOL><INDENT>raise NotImplementedError('<STR_LIT>')<EOL><DEDENT>max_rank = tf.convert_to_tensor(<EOL>value=max_rank, name='<STR_LIT>', dtype=tf.int64)<EOL>max_rank = tf.minimum(max_rank,<EOL>prefer_static.shape(matrix, out_type=tf.int64)[-<NUM_LIT:1>])<EOL>diag_rtol = tf.convert_to_tensor(<EOL>value=diag_rtol, dtype=dtype, name='<STR_LIT>')<EOL>matrix_diag = tf.linalg.diag_part(matrix)<EOL>orig_error = tf.reduce_max(input_tensor=matrix_diag, axis=-<NUM_LIT:1>)<EOL>def cond(m, pchol, perm, matrix_diag):<EOL><INDENT>"""<STR_LIT>"""<EOL>del pchol<EOL>del perm<EOL>error = tf.linalg.norm(tensor=matrix_diag, ord=<NUM_LIT:1>, axis=-<NUM_LIT:1>)<EOL>max_err = tf.reduce_max(input_tensor=error / orig_error)<EOL>return (m < max_rank) & (tf.equal(m, <NUM_LIT:0>) | (max_err > diag_rtol))<EOL><DEDENT>batch_dims = tensorshape_util.rank(matrix.shape) - <NUM_LIT:2><EOL>def batch_gather(params, indices, axis=-<NUM_LIT:1>):<EOL><INDENT>return tf.gather(params, indices, axis=axis, batch_dims=batch_dims)<EOL><DEDENT>def body(m, pchol, perm, matrix_diag):<EOL><INDENT>"""<STR_LIT>"""<EOL>permuted_diag = batch_gather(matrix_diag, perm[..., m:])<EOL>maxi = tf.argmax(<EOL>input=permuted_diag, axis=-<NUM_LIT:1>, output_type=tf.int64)[..., tf.newaxis]<EOL>maxval = batch_gather(permuted_diag, maxi)<EOL>maxi = maxi + m<EOL>maxval = maxval[..., <NUM_LIT:0>]<EOL>perm = _swap_m_with_i(perm, m, maxi)<EOL>row = batch_gather(matrix, perm[..., m:m + <NUM_LIT:1>], axis=-<NUM_LIT:2>)<EOL>row = batch_gather(row, perm[..., m + <NUM_LIT:1>:])<EOL>prev_rows = pchol[..., :m, :]<EOL>prev_rows_perm_m_onward = batch_gather(prev_rows, perm[..., m + <NUM_LIT:1>:])<EOL>prev_rows_pivot_col = batch_gather(prev_rows, perm[..., m:m + <NUM_LIT:1>])<EOL>row -= tf.reduce_sum(<EOL>input_tensor=prev_rows_perm_m_onward * prev_rows_pivot_col,<EOL>axis=-<NUM_LIT:2>)[..., tf.newaxis, :]<EOL>pivot = tf.sqrt(maxval)[..., tf.newaxis, tf.newaxis]<EOL>row = tf.concat([pivot, row / pivot], axis=-<NUM_LIT:1>)<EOL>paddings = tf.concat([<EOL>tf.zeros([prefer_static.rank(pchol) - <NUM_LIT:1>, <NUM_LIT:2>], dtype=tf.int32),<EOL>[[tf.cast(m, tf.int32), <NUM_LIT:0>]]], axis=<NUM_LIT:0>)<EOL>diag_update = tf.pad(tensor=row**<NUM_LIT:2>, paddings=paddings)[..., <NUM_LIT:0>, :]<EOL>reverse_perm = _invert_permutation(perm)<EOL>matrix_diag -= batch_gather(diag_update, reverse_perm)<EOL>row = tf.pad(tensor=row, paddings=paddings)<EOL>row = batch_gather(row, reverse_perm)<EOL>pchol_shape = pchol.shape<EOL>pchol = tf.concat([pchol[..., :m, :], row, pchol[..., m + <NUM_LIT:1>:, :]],<EOL>axis=-<NUM_LIT:2>)<EOL>tensorshape_util.set_shape(pchol, pchol_shape)<EOL>return m + <NUM_LIT:1>, pchol, perm, matrix_diag<EOL><DEDENT>m = np.int64(<NUM_LIT:0>)<EOL>pchol = tf.zeros_like(matrix[..., :max_rank, :])<EOL>matrix_shape = prefer_static.shape(matrix, out_type=tf.int64)<EOL>perm = tf.broadcast_to(<EOL>prefer_static.range(matrix_shape[-<NUM_LIT:1>]), matrix_shape[:-<NUM_LIT:1>])<EOL>_, pchol, _, _ = tf.while_loop(<EOL>cond=cond, body=body, loop_vars=(m, pchol, perm, matrix_diag))<EOL>pchol = tf.linalg.matrix_transpose(pchol)<EOL>tensorshape_util.set_shape(<EOL>pchol, tensorshape_util.concatenate(matrix_diag.shape, [None]))<EOL>return pchol<EOL><DEDENT>

Computes the (partial) pivoted cholesky decomposition of `matrix`.

    The pivoted Cholesky is a low rank approximation of the Cholesky decomposition
    of `matrix`, i.e. as described in [(Harbrecht et al., 2012)][1]. The
    currently-worst-approximated diagonal element is selected as the pivot at each
    iteration. This yields from a `[B1...Bn, N, N]` shaped `matrix` a `[B1...Bn,
    N, K]` shaped rank-`K` approximation `lr` such that `lr @ lr.T ~= matrix`.
    Note that, unlike the Cholesky decomposition, `lr` is not triangular even in
    a rectangular-matrix sense. However, under a permutation it could be made
    triangular (it has one more zero in each column as you move to the right).

    Such a matrix can be useful as a preconditioner for conjugate gradient
    optimization, i.e. as in [(Wang et al. 2019)][2], as matmuls and solves can be
    cheaply done via the Woodbury matrix identity, as implemented by
    `tf.linalg.LinearOperatorLowRankUpdate`.

    Args:
      matrix: Floating point `Tensor` batch of symmetric, positive definite
        matrices.
      max_rank: Scalar `int` `Tensor`, the rank at which to truncate the
        approximation.
      diag_rtol: Scalar floating point `Tensor` (same dtype as `matrix`). If the
        errors of all diagonal elements of `lr @ lr.T` are each lower than
        `element * diag_rtol`, iteration is permitted to terminate early.
      name: Optional name for the op.

    Returns:
      lr: Low rank pivoted Cholesky approximation of `matrix`.

    #### References

    [1]: H Harbrecht, M Peters, R Schneider. On the low-rank approximation by the
         pivoted Cholesky decomposition. _Applied numerical mathematics_,
         62(4):428-440, 2012.

    [2]: K. A. Wang et al. Exact Gaussian Processes on a Million Data Points.
         _arXiv preprint arXiv:1903.08114_, 2019. https://arxiv.org/abs/1903.08114


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[sparse_or_dense_matrix, dense_vector]):<EOL><INDENT>dense_vector = tf.convert_to_tensor(<EOL>value=dense_vector, dtype_hint=tf.float32, name='<STR_LIT>')<EOL>return tf.squeeze(<EOL>sparse_or_dense_matmul(<EOL>sparse_or_dense_matrix,<EOL>dense_vector[..., tf.newaxis],<EOL>validate_args=validate_args,<EOL>**kwargs),<EOL>axis=[-<NUM_LIT:1>])<EOL><DEDENT>

Returns (batched) matmul of a (sparse) matrix with a column vector.

    Args:
      sparse_or_dense_matrix: `SparseTensor` or `Tensor` representing a (batch of)
        matrices.
      dense_vector: `Tensor` representing a (batch of) vectors, with the same
        batch shape as `sparse_or_dense_matrix`. The shape must be compatible with
        the shape of `sparse_or_dense_matrix` and kwargs.
      validate_args: When `True`, additional assertions might be embedded in the
        graph.
        Default value: `False` (i.e., no graph assertions are added).
      name: Python `str` prefixed to ops created by this function.
        Default value: "sparse_or_dense_matvecmul".
      **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
        `tf.matmul`.

    Returns:
      product: A dense (batch of) vector-shaped Tensor of the same batch shape and
      dtype as `sparse_or_dense_matrix` and `dense_vector`.


def _lu_solve_assertions(lower_upper, perm, rhs, validate_args):

assertions = _lu_reconstruct_assertions(lower_upper, perm, validate_args)<EOL>message = '<STR_LIT>'<EOL>if rhs.shape.ndims is not None:<EOL><INDENT>if rhs.shape.ndims < <NUM_LIT:2>:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_rank_at_least(rhs, rank=<NUM_LIT:2>, message=message))<EOL><DEDENT>message = '<STR_LIT>'<EOL>if (tf.compat.dimension_value(lower_upper.shape[-<NUM_LIT:1>]) is not None and<EOL>tf.compat.dimension_value(rhs.shape[-<NUM_LIT:2>]) is not None):<EOL><INDENT>if lower_upper.shape[-<NUM_LIT:1>] != rhs.shape[-<NUM_LIT:2>]:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_equal(<EOL>tf.shape(input=lower_upper)[-<NUM_LIT:1>],<EOL>tf.shape(input=rhs)[-<NUM_LIT:2>],<EOL>message=message))<EOL><DEDENT>return assertions<EOL>

Returns list of assertions related to `lu_solve` assumptions.


<EOL>sp_a_shape = tf.convert_to_tensor(value=_get_shape(sp_a, tf.int64))<EOL>ind_mat = tf.concat([[sp_a_shape[-<NUM_LIT:2>:]], tf.eye(<NUM_LIT:2>, dtype=tf.int64)], axis=<NUM_LIT:0>)<EOL>indices = tf.matmul(sp_a.indices, ind_mat)<EOL>dense_shape = sp_a_shape[<NUM_LIT:0>] * sp_a_shape[<NUM_LIT:1>:]<EOL>return tf.SparseTensor(<EOL>indices=indices, values=sp_a.values, dense_shape=dense_shape)<EOL>

Returns a block diagonal rank 2 SparseTensor from a batch of SparseTensors.

    Args:
      sp_a: A rank 3 `SparseTensor` representing a batch of matrices.

    Returns:
      sp_block_diag_a: matrix-shaped, `float` `SparseTensor` with the same dtype
      as `sparse_or_matrix`, of shape [B * M, B * N] where `sp_a` has shape
      [B, M, N]. Each [M, N] batch of `sp_a` is lined up along the diagonal.


ndims = x.shape.ndims<EOL>if ndims is None:<EOL><INDENT>if expect_static:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(x))<EOL><DEDENT>return<EOL><DEDENT>if expect_ndims is not None and ndims != expect_ndims:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(expect_ndims, ndims))<EOL><DEDENT>if expect_ndims_at_least is not None and ndims < expect_ndims_at_least:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(<EOL>expect_ndims_at_least, ndims))<EOL><DEDENT>

Assert that Tensor x has expected number of dimensions.


Linear `1-D` interpolation on a regular (constant spacing) grid.

    Given [batch of] reference values, this function computes a piecewise linear
    interpolant and evaluates it on a [batch of] of new `x` values.

    The interpolant is built from `C` reference values indexed by one dimension
    of `y_ref` (specified by the `axis` kwarg).

    If `y_ref` is a vector, then each value `y_ref[i]` is considered to be equal
    to `f(x_ref[i])`, for `C` (implicitly defined) reference values between
    `x_ref_min` and `x_ref_max`:

    ```none
    x_ref[i] = x_ref_min + i * (x_ref_max - x_ref_min) / (C - 1),
    i = 0, ..., C - 1.
    ```

    In the general case, dimensions to the left of `axis` in `y_ref` are broadcast
    with leading dimensions in `x`, `x_ref_min`, `x_ref_max`.

    Args:
      x: Numeric `Tensor` The x-coordinates of the interpolated output values for
        each batch.  Shape broadcasts with `[A1, ..., AN, D]`, `N >= 0`.
      x_ref_min:  `Tensor` of same `dtype` as `x`.  The minimum value of the each
        batch of the (implicitly defined) reference `x_ref`. Shape broadcasts with
        `[A1, ..., AN]`, `N >= 0`.
      x_ref_max:  `Tensor` of same `dtype` as `x`.  The maximum value of the each
        batch of the (implicitly defined) reference `x_ref`. Shape broadcasts with
        `[A1, ..., AN]`, `N >= 0`.
      y_ref:  `Tensor` of same `dtype` as `x`.  The reference output values.
        `y_ref.shape[:axis]` broadcasts with the batch shape `[A1, ..., AN]`, and
        `y_ref.shape[axis:]` is `[C, B1, ..., BM]`, so the trailing dimensions
          index `C` reference values of a rank `M` `Tensor` (`M >= 0`).
      axis:  Scalar `Tensor` designating the dimension of `y_ref` that indexes
        values of the interpolation table.
        Default value: `-1`, the rightmost axis.
      fill_value:  Determines what values output should take for `x` values that
        are below `x_ref_min` or above `x_ref_max`. `Tensor` or one of the strings
        "constant_extension" ==> Extend as constant function. "extrapolate" ==>
        Extrapolate in a linear fashion.
        Default value: `"constant_extension"`
      fill_value_below:  Optional override of `fill_value` for `x < x_ref_min`.
      fill_value_above:  Optional override of `fill_value` for `x > x_ref_max`.
      grid_regularizing_transform:  Optional transformation `g` which regularizes
        the implied spacing of the x reference points.  In other words, if
        provided, we assume `g(x_ref_i)` is a regular grid between `g(x_ref_min)`
        and `g(x_ref_max)`.
      name:  A name to prepend to created ops.
        Default value: `"batch_interp_regular_1d_grid"`.

    Returns:
      y_interp:  Interpolation between members of `y_ref`, at points `x`.
        `Tensor` of same `dtype` as `x`, and shape `[A1, ..., AN, D, B1, ..., BM]`

    Raises:
      ValueError:  If `fill_value` is not an allowed string.
      ValueError:  If `axis` is not a scalar.

    #### Examples

    Interpolate a function of one variable:

    ```python
    y_ref = tf.exp(tf.linspace(start=0., stop=10., 20))

    batch_interp_regular_1d_grid(
        x=[6.0, 0.5, 3.3], x_ref_min=0., x_ref_max=1., y_ref=y_ref)
    ==> approx [exp(6.0), exp(0.5), exp(3.3)]
    ```

    Interpolate a batch of functions of one variable.

    ```python
    # First batch member is an exponential function, second is a log.
    implied_x_ref = [tf.linspace(-3., 3.2, 200), tf.linspace(0.5, 3., 200)]
    y_ref = tf.stack(  # Shape [2, 200], 2 batches, 200 reference values per batch
        [tf.exp(implied_x_ref[0]), tf.log(implied_x_ref[1])], axis=0)

    x = [[-1., 1., 0.],  # Shape [2, 3], 2 batches, 3 values per batch.
         [1., 2., 3.]]

    y = tfp.math.batch_interp_regular_1d_grid(  # Shape [2, 3]
        x,
        x_ref_min=[-3., 0.5],
        x_ref_max=[3.2, 3.],
        y_ref=y_ref,
        axis=-1)

    # y[0] approx tf.exp(x[0])
    # y[1] approx tf.log(x[1])
    ```

    Interpolate a function of one variable on a log-spaced grid:

    ```python
    x_ref = tf.exp(tf.linspace(tf.log(1.), tf.log(100000.), num_pts))
    y_ref = tf.log(x_ref + x_ref**2)

    batch_interp_regular_1d_grid(x=[1.1, 2.2], x_ref_min=1., x_ref_max=100000.,
        y_ref, grid_regularizing_transform=tf.log)
    ==> [tf.log(1.1 + 1.1**2), tf.log(2.2 + 2.2**2)]
    ```


def random_rademacher(shape, dtype=tf.float32, seed=None, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [shape, seed]):<EOL><INDENT>generation_dtype = tf.int64 if tf.as_dtype(dtype) != tf.int32 else tf.int32<EOL>random_bernoulli = tf.random.uniform(<EOL>shape, minval=<NUM_LIT:0>, maxval=<NUM_LIT:2>, dtype=generation_dtype, seed=seed)<EOL>return tf.cast(<NUM_LIT:2> * random_bernoulli - <NUM_LIT:1>, dtype)<EOL><DEDENT>

Generates `Tensor` consisting of `-1` or `+1`, chosen uniformly at random.

    For more details, see [Rademacher distribution](
    https://en.wikipedia.org/wiki/Rademacher_distribution).

    Args:
      shape: Vector-shaped, `int` `Tensor` representing shape of output.
      dtype: (Optional) TF `dtype` representing `dtype` of output.
      seed: (Optional) Python integer to seed the random number generator.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'random_rademacher').

    Returns:
      rademacher: `Tensor` with specified `shape` and `dtype` consisting of `-1`
        or `+1` chosen uniformly-at-random.


def value_and_gradient(f, xs, use_gradient_tape=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [xs]):<EOL><INDENT>is_xs_list_like = isinstance(xs, (tuple, list))<EOL>if not is_xs_list_like:<EOL><INDENT>xs = [xs]<EOL><DEDENT>xs = [tf.convert_to_tensor(value=x, name='<STR_LIT>'.format(i))<EOL>for i, x in enumerate(xs)]<EOL>if tf.executing_eagerly() or use_gradient_tape:<EOL><INDENT>with tf.GradientTape(watch_accessed_variables=False) as tape:<EOL><INDENT>for x in xs:<EOL><INDENT>tape.watch(x)<EOL><DEDENT>y = f(*xs)<EOL><DEDENT>dydx = tape.gradient(y, xs)<EOL><DEDENT>else:<EOL><INDENT>y = f(*xs)<EOL>dydx = tf.gradients(ys=y, xs=xs)<EOL><DEDENT>if not is_xs_list_like:<EOL><INDENT>dydx = dydx[<NUM_LIT:0>]<EOL><DEDENT>return y, dydx<EOL><DEDENT>

Computes `f(*xs)` and its gradients wrt to `*xs`.

    Args:
      f: Python `callable` to be differentiated. If `f` returns a scalar, this
        scalar will be differentiated. If `f` returns a tensor or list of tensors,
        by default a scalar will be computed by adding all their values to produce
        a single scalar. If desired, the tensors can be elementwise multiplied by
        the tensors passed as the `dy` keyword argument to the returned gradient
        function.
      xs: Python list of parameters of f for which to differentiate. (Can also
        be single `Tensor`.)
      use_gradient_tape: Python `bool` indicating that `tf.GradientTape`
        should be used regardless of `tf.executing_eagerly()` status.
        Default value: `False`.
      name: Python `str` name prefixed to ops created by this function.
        Default value: `None` (i.e., `'value_and_gradient'`).

    Returns:
      y: `y = f(*xs)`.
      dydx: Gradient of `y` wrt each of `xs`.


Solves an initial value problem.

        An initial value problem consists of a system of ODEs and an initial
        condition:

        ```none
        dy/dt(t) = ode_fn(t, y(t))
        y(initial_time) = initial_state
        ```

        Here, `t` (also called time) is a scalar `Tensor` of real `dtype` and `y(t)`
        (also called the state at time `t`) is an N-D float or complex `Tensor`.

        Args:
          ode_fn: Function of the form `ode_fn(t, y)`. The input `t` is a scalar
            `Tensor` of real `dtype`. The input `y` and output are both `Tensor`s
            with the same shape and `dtype` as `initial_state`.
          initial_state: N-D float or complex `Tensor` specifying the initial state.
            The `dtype` of `initial_state` must be complex for problems with
            complex-valued states (even if the initial state is real).
          initial_time: Scalar `Tensor` of real `dtype` specifying the initial time.
            Default value: `0.`.
          final_time: Optional scalar `Tensor` of real `dtype` specifying the final
            time to integrate up to. Must satisfy `initial_time < final_time`. If
            unspecified, the solver will use `solution_times[-1]` (see below) as the
            final time. Therefore, at least one of `final_time` or `solution_times`
            must be specified.
            Default value: `None` (i.e., `solution_times[-1]` or error).
          solution_times: Optional 1-D float `Tensor` specifying a list of times. If
            specified, the solver stores the computed state at each of these times
            in the returned `Solution` object. If unspecified, the solver will
            choose this list of times automatically. Must satisfy `initial_time <=
            solution_times[0]`, `solution_times[i] < solution_times[i+1]`, and
            `solution_times[-1] <= final_time` if `final_time` is specified.
            Default value: `None`.
          jacobian_fn: Optional function of the form `jacobian_fn(t, y)`. The input
            `t` is a scalar `Tensor` of real `dtype`. The input `y` has the same
            shape and `dtype` as `initial_state`. The output is a 2N-D `Tensor`
            whose shape is `initial_state.shape + initial_state.shape` and whose
            `dtype` is the same as `initial_state`. In particular, the `(i1, ...,
            iN, j1, ..., jN)`-th entry of `jacobian_fn(t, y)` is the derivative of
            the `(i1, ..., iN)`-th entry of `ode_fn(t, y)` with respect to the `(j1,
            ..., jN)`-th entry of `y`. If this argument is left unspecified, the
            solver will automatically compute the Jacobian if and when it is needed.
            Default value: `None`.
          jacobian_sparsity: Optional 2N-D boolean `Tensor` whose shape is
            `initial_state.shape + initial_state.shape` specifying the sparsity
            pattern of the Jacobian. This argument is ignored if `jacobian_fn` is
            specified.
            Default value: `None`.
          batch_ndims: Optional nonnegative integer. When specified, the first
            `batch_ndims` dimensions of `initial_state` are batch dimensions.
            Default value: `None`.
          previous_results: Optional solver-specific argument used to warm-start
            this invocation of `solve`.
            Default value: `None`.

        Returns:
          solution: Object of type `Solution` containing the solution.
          diagnostics: Object of type `Diagnostics` containing performance
            information.
          results: Solver-specific object which can be used to warm-start the solver
            on a future invocation of `solve`.


for i in sorted(axis):<EOL><INDENT>x = tf.expand_dims(x, axis=i)<EOL><DEDENT>return x<EOL>

Insert the dims in `axis` back as singletons after being removed.

    Args:
      x:  `Tensor`.
      axis:  Python list of integers.

    Returns:
      `Tensor` with same values as `x`, but additional singleton dimensions.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[x, num_quantiles, axis]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>return percentile(<EOL>x,<EOL>q=tf.linspace(<EOL>tf.convert_to_tensor(value=<NUM_LIT:0>, dtype=tf.float64),<EOL>tf.convert_to_tensor(value=<NUM_LIT:100>, dtype=tf.float64),<EOL>num=num_quantiles + <NUM_LIT:1>),<EOL>axis=axis,<EOL>interpolation=interpolation,<EOL>keep_dims=keep_dims,<EOL>validate_args=validate_args,<EOL>preserve_gradients=False)<EOL><DEDENT>

Compute quantiles of `x` along `axis`.

    The quantiles of a distribution are cut points dividing the range into
    intervals with equal probabilities.

    Given a vector `x` of samples, this function estimates the cut points by
    returning `num_quantiles + 1` cut points, `(c0, ..., cn)`, such that, roughly
    speaking, equal number of sample points lie in the `num_quantiles` intervals
    `[c0, c1), [c1, c2), ..., [c_{n-1}, cn]`.  That is,

    * About `1 / n` fraction of the data lies in `[c_{k-1}, c_k)`, `k = 1, ..., n`
    * About `k / n` fraction of the data lies below `c_k`.
    * `c0` is the sample minimum and `cn` is the maximum.

    The exact number of data points in each interval depends on the size of
    `x` (e.g. whether the size is divisible by `n`) and the `interpolation` kwarg.

    Args:
      x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not `None`,
        `x` must have statically known number of dimensions.
      num_quantiles:  Scalar `integer` `Tensor`.  The number of intervals the
        returned `num_quantiles + 1` cut points divide the range into.
      axis:  Optional `0-D` or `1-D` integer `Tensor` with constant values. The
        axis that index independent samples over which to return the desired
        percentile.  If `None` (the default), treat every dimension as a sample
        dimension, returning a scalar.
      interpolation : {'nearest', 'linear', 'lower', 'higher', 'midpoint'}.
        Default value: 'nearest'.  This specifies the interpolation method to
        use when the fractions `k / n` lie between two data points `i < j`:
          * linear: i + (j - i) * fraction, where fraction is the fractional part
            of the index surrounded by i and j.
          * lower: `i`.
          * higher: `j`.
          * nearest: `i` or `j`, whichever is nearest.
          * midpoint: (i + j) / 2. `linear` and `midpoint` interpolation do not
            work with integer dtypes.
      keep_dims:  Python `bool`. If `True`, the last dimension is kept with size 1
        If `False`, the last dimension is removed from the output shape.
      validate_args:  Whether to add runtime checks of argument validity. If
        False, and arguments are incorrect, correct behavior is not guaranteed.
      name:  A Python string name to give this `Op`.  Default is 'percentile'

    Returns:
      cut_points:  A `rank(x) + 1 - len(axis)` dimensional `Tensor` with same
      `dtype` as `x` and shape `[num_quantiles + 1, ...]` where the trailing shape
      is that of `x` without the dimensions in `axis` (unless `keep_dims is True`)

    Raises:
      ValueError:  If argument 'interpolation' is not an allowed type.
      ValueError:  If interpolation type not compatible with `dtype`.

    #### Examples

    ```python
    # Get quartiles of x with various interpolation choices.
    x = [0.,  1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.]

    tfp.stats.quantiles(x, num_quantiles=4, interpolation='nearest')
    ==> [  0.,   2.,   5.,   8.,  10.]

    tfp.stats.quantiles(x, num_quantiles=4, interpolation='linear')
    ==> [  0. ,   2.5,   5. ,   7.5,  10. ]

    tfp.stats.quantiles(x, num_quantiles=4, interpolation='lower')
    ==> [  0.,   2.,   5.,   7.,  10.]

    # Get deciles of columns of an R x C data set.
    data = load_my_columnar_data(...)
    tfp.stats.quantiles(data, num_quantiles=10)
    ==> Shape [11, C] Tensor
    ```


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[x, y, event_axis, sample_axis]):<EOL><INDENT>x /= stddev(x, sample_axis=sample_axis, keepdims=True)<EOL>if y is not None:<EOL><INDENT>y /= stddev(y, sample_axis=sample_axis, keepdims=True)<EOL><DEDENT>return covariance(<EOL>x=x,<EOL>y=y,<EOL>event_axis=event_axis,<EOL>sample_axis=sample_axis,<EOL>keepdims=keepdims)<EOL><DEDENT>

Sample correlation (Pearson) between observations indexed by `event_axis`.

    Given `N` samples of scalar random variables `X` and `Y`, correlation may be
    estimated as

    ```none
    Corr[X, Y] := Cov[X, Y] / Sqrt(Cov[X, X] * Cov[Y, Y]),
    where
    Cov[X, Y] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(Y_n - Ybar)}
    Xbar := N^{-1} sum_{n=1}^N X_n
    Ybar := N^{-1} sum_{n=1}^N Y_n
    ```

    Correlation is always in the interval `[-1, 1]`, and `Corr[X, X] == 1`.

    For vector-variate random variables `X = (X1, ..., Xd)`, `Y = (Y1, ..., Yd)`,
    one is often interested in the correlation matrix, `C_{ij} := Corr[Xi, Yj]`.

    ```python
    x = tf.random_normal(shape=(100, 2, 3))
    y = tf.random_normal(shape=(100, 2, 3))

    # corr[i, j] is the sample correlation between x[:, i, j] and y[:, i, j].
    corr = tfp.stats.correlation(x, y, sample_axis=0, event_axis=None)

    # corr_matrix[i, m, n] is the sample correlation of x[:, i, m] and y[:, i, n]
    corr_matrix = tfp.stats.correlation(x, y, sample_axis=0, event_axis=-1)
    ```

    Notice we divide by `N` (the numpy default), which does not create `NaN`
    when `N = 1`, but is slightly biased.

    Args:
      x:  A numeric `Tensor` holding samples.
      y:  Optional `Tensor` with same `dtype` and `shape` as `x`.
        Default value: `None` (`y` is effectively set to `x`).
      sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
        `None` (meaning all axis hold samples).
        Default value: `0` (leftmost dimension).
      event_axis:  Scalar or vector `Tensor`, or `None` (scalar events).
        Axis indexing random events, whose correlation we are interested in.
        If a vector, entries must form a contiguous block of dims. `sample_axis`
        and `event_axis` should not intersect.
        Default value: `-1` (rightmost axis holds events).
      keepdims:  Boolean.  Whether to keep the sample axis as singletons.
      name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., `'correlation'`).

    Returns:
      corr: A `Tensor` of same `dtype` as the `x`, and rank equal to
        `rank(x) - len(sample_axis) + 2 * len(event_axis)`.

    Raises:
      AssertionError:  If `x` and `y` are found to have different shape.
      ValueError:  If `sample_axis` and `event_axis` are found to overlap.
      ValueError:  If `event_axis` is found to not be contiguous.


class Dummy(object):<EOL><INDENT>pass<EOL><DEDENT>num_examples = <NUM_LIT:10><EOL>mnist_data = Dummy()<EOL>mnist_data.train = Dummy()<EOL>mnist_data.train.images = np.float32(np.random.randn(<EOL>num_examples, np.prod(IMAGE_SHAPE)))<EOL>mnist_data.train.labels = np.int32(np.random.permutation(<EOL>np.arange(num_examples)))<EOL>mnist_data.train.num_examples = num_examples<EOL>mnist_data.validation = Dummy()<EOL>mnist_data.validation.images = np.float32(np.random.randn(<EOL>num_examples, np.prod(IMAGE_SHAPE)))<EOL>mnist_data.validation.labels = np.int32(np.random.permutation(<EOL>np.arange(num_examples)))<EOL>mnist_data.validation.num_examples = num_examples<EOL>return mnist_data<EOL>

Builds fake MNIST-style data for unit testing.


Helper method to save images visualizing model reconstructions.

    Args:
      images_val: Numpy array containing a batch of input images.
      reconstructed_images_val: Numpy array giving the expected output
        (mean) of the decoder.
      random_images_val: Optionally, a Numpy array giving the expected output
        (mean) of decoding samples from the prior, or `None`.
      log_dir: The directory to write images (Python `str`).
      prefix: A specific label for the saved visualizations, which
        determines their filenames (Python `str`).
      viz_n: The number of images from each batch to visualize (Python `int`).


out = self.dense(inputs)  <EOL>out = self.output_layer(out)  <EOL>loc = out[..., :self.latent_size]<EOL>scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + <NUM_LIT>  <EOL>return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)<EOL>

Runs the model to generate a distribution `q(z_{1:T} | x_{1:T})`.

        Args:
          inputs: A batch of intermediate representations of image frames
            across all timesteps, of shape [..., batch_size, timesteps,
            hidden_size].

        Returns:
          A batch of MultivariateNormalDiag distributions with event shape
          [latent_size], batch shape [..., batch_size, timesteps], and
          sample shape [sample_shape, ..., batch_size, timesteps,
          latent_size].


def visualize_qualitative_analysis(inputs, model, samples=<NUM_LIT:1>, batch_size=<NUM_LIT:3>,<EOL>length=<NUM_LIT:8>):

average = lambda dist: tf.reduce_mean(<EOL>input_tensor=dist.mean(), axis=<NUM_LIT:0>)  <EOL>with tf.compat.v1.name_scope("<STR_LIT>"):<EOL><INDENT>reconstruct = functools.partial(model.reconstruct, inputs=inputs,<EOL>samples=samples)<EOL>visualize_reconstruction(inputs, average(reconstruct()))<EOL>visualize_reconstruction(inputs, average(reconstruct(sample_static=True)),<EOL>name="<STR_LIT>")<EOL>visualize_reconstruction(inputs, average(reconstruct(sample_dynamic=True)),<EOL>name="<STR_LIT>")<EOL>visualize_reconstruction(inputs, average(reconstruct(swap_static=True)),<EOL>name="<STR_LIT>")<EOL>visualize_reconstruction(inputs, average(reconstruct(swap_dynamic=True)),<EOL>name="<STR_LIT>")<EOL><DEDENT>with tf.compat.v1.name_scope("<STR_LIT>"):<EOL><INDENT>generate = functools.partial(model.generate, batch_size=batch_size,<EOL>length=length, samples=samples)<EOL>image_summary(average(generate(fix_static=True)), "<STR_LIT>")<EOL>image_summary(average(generate(fix_dynamic=True)), "<STR_LIT>")<EOL><DEDENT>

Visualizes a qualitative analysis of a given model.

    Args:
      inputs: A tensor of the original inputs, of shape [batch, timesteps,
        h, w, c].
      model: A DisentangledSequentialVAE model.
      samples: Number of samples to draw from the latent distributions.
      batch_size: Number of sequences to generate.
      length: Number of timesteps to generate for each sequence.


<EOL>features, static_sample = inputs<EOL>length = tf.shape(input=features)[-<NUM_LIT:2>]<EOL>static_sample = static_sample[..., tf.newaxis, :] + tf.zeros([length, <NUM_LIT:1>])<EOL>sample_shape_static = tf.shape(input=static_sample)[:-<NUM_LIT:3>]<EOL>sample_shape_inputs = tf.shape(input=features)[:-<NUM_LIT:3>]<EOL>broadcast_shape_inputs = tf.concat((sample_shape_static, [<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:1>]), <NUM_LIT:0>)<EOL>broadcast_shape_static = tf.concat((sample_shape_inputs, [<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:1>]), <NUM_LIT:0>)<EOL>features = features + tf.zeros(broadcast_shape_inputs)<EOL>static_sample = static_sample + tf.zeros(broadcast_shape_static)<EOL>combined = tf.concat((features, static_sample), axis=-<NUM_LIT:1>)<EOL>collapsed_shape = tf.concat(([-<NUM_LIT:1>], tf.shape(input=combined)[-<NUM_LIT:2>:]), axis=<NUM_LIT:0>)<EOL>out = tf.reshape(combined, collapsed_shape)<EOL>out = self.bilstm(out)  <EOL>out = self.rnn(out)  <EOL>expanded_shape = tf.concat(<EOL>(tf.shape(input=combined)[:-<NUM_LIT:2>], tf.shape(input=out)[<NUM_LIT:1>:]), axis=<NUM_LIT:0>)<EOL>out = tf.reshape(out, expanded_shape)  <EOL>out = self.output_layer(out)  <EOL>loc = out[..., :self.latent_size]<EOL>scale_diag = tf.nn.softplus(out[..., self.latent_size:]) + <NUM_LIT>  <EOL>return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)<EOL>

Runs the model to generate a distribution `q(z_{1:T} | x_{1:T}, f)`.

        This generates a list of batched MultivariateNormalDiag
        distributions using the output of the recurrent model at each
        timestep to parameterize each distribution.

        Args:
          inputs: A tuple of a batch of intermediate representations of
            image frames across all timesteps of shape [..., batch_size,
            timesteps, dimensions], and a sample of the static latent
            variable `f` of shape [..., batch_size, latent_size].

        Returns:
          A batch of MultivariateNormalDiag distributions with event shape
          [latent_size], batch shape [broadcasted_shape, batch_size,
          timesteps], and sample shape [sample_shape, broadcasted_shape,
          batch_size, timesteps, latent_size], where `broadcasted_shape` is
          the broadcasted sampled shape between the inputs and static
          sample.


return tf.nn.softplus(self._untransformed_stddev) + <NUM_LIT><EOL>

The diagonal standard deviation of the normal distribution.


del inputs  <EOL>with tf.compat.v1.name_scope(self._name):<EOL><INDENT>return tfd.MultivariateNormalDiag(self.loc, self.scale_diag)<EOL><DEDENT>

Runs the model to generate multivariate normal distribution.

        Args:
          inputs: Unused.

        Returns:
          A MultivariateNormalDiag distribution with event shape
          [dimensions], batch shape [], and sample shape [sample_shape,
          dimensions].


def __init__(self, dimensions, hidden_size):

super(LearnableMultivariateNormalDiagCell, self).__init__()<EOL>self.dimensions = dimensions<EOL>self.hidden_size = hidden_size<EOL>self.lstm_cell = tf.keras.layers.LSTMCell(hidden_size)<EOL>self.output_layer = tf.keras.layers.Dense(<NUM_LIT:2>*dimensions)<EOL>

Constructs a learnable multivariate diagonal normal cell.

        Args:
          dimensions: An integer corresponding to the dimensionality of the
            distribution.
          hidden_size: Dimensionality of the LSTM function parameters.