with tf.compat.v1.name_scope(<EOL>name, values=[bias_variance, slope_variance, shift, exponent]):<EOL><INDENT>dtype = util.maybe_get_common_dtype(<EOL>[bias_variance, slope_variance, shift, exponent])<EOL>if bias_variance is not None:<EOL><INDENT>bias_variance = tf.convert_to_tensor(<EOL>value=bias_variance, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._bias_variance = _validate_arg_if_not_none(<EOL>bias_variance, tf.compat.v1.assert_positive, validate_args)<EOL>if slope_variance is not None:<EOL><INDENT>slope_variance = tf.convert_to_tensor(<EOL>value=slope_variance, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._slope_variance = _validate_arg_if_not_none(<EOL>slope_variance, tf.compat.v1.assert_positive, validate_args)<EOL>if shift is not None:<EOL><INDENT>shift = tf.convert_to_tensor(value=shift, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._shift = shift<EOL>if exponent is not None:<EOL><INDENT>exponent = tf.convert_to_tensor(<EOL>value=exponent, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._exponent = _validate_arg_if_not_none(<EOL>exponent, tf.compat.v1.assert_positive, validate_args)<EOL><DEDENT>super(Polynomial, self).__init__(<EOL>feature_ndims, dtype=dtype, name=name)<EOL>

Construct a Polynomial kernel instance.

        Args:
          bias_variance: Positive floating point `Tensor` that controls the
            variance from the origin. If bias = 0, there is no variance and the
            fitted function goes through the origin.  Must be broadcastable with
            `slope_variance`, `shift`, `exponent`, and inputs to `apply` and
            `matrix` methods. A value of `None` is treated like 0.
            Default Value: `None`
          slope_variance: Positive floating point `Tensor` that controls the
            variance of the regression line slope that is the basis for the
            polynomial. Must be broadcastable with `bias_variance`, `shift`,
            `exponent`, and inputs to `apply` and `matrix` methods. A value of
            `None` is treated like 1.
            Default Value: `None`
          shift: Floating point `Tensor` that contols the intercept with the
            x-axis of the linear function to be exponentiated to get this
            polynomial. Must be broadcastable with `bias_variance`,
            `slope_variance`, `exponent` and inputs to `apply` and `matrix`
            methods. A value of `None` is treated like 0, which results in having
            the intercept at the origin.
            Default Value: `None`
          exponent: Positive floating point `Tensor` that controls the exponent
            (also known as the degree) of the polynomial function. Must be
            broadcastable with `bias_variance`, `slope_variance`, `shift`,
            and inputs to `apply` and `matrix` methods. A value of `None` is
            treated like 1, which results in a linear kernel.
            Default Value: `None`
          feature_ndims: Python `int` number of rightmost dims to include in
            kernel computation.
            Default Value: 1
          validate_args: If `True`, parameters are checked for validity despite
            possibly degrading runtime performance.
            Default Value: `False`
          name: Python `str` name prefixed to Ops created by this class.
            Default Value: `'Polynomial'`


return self._bias_variance<EOL>


return self._slope_variance<EOL>


return self._shift<EOL>

Shift of linear function that is exponentiated.


return self._exponent<EOL>


Construct a Linear kernel instance.

        Args:
          bias_variance: Positive floating point `Tensor` that controls the
            variance from the origin. If bias = 0, there is no variance and the
            fitted function goes through the origin (also known as the homogeneous
            linear kernel). Must be broadcastable with `slope_variance`,
            `shift` and inputs to `apply` and `matrix` methods. A value of
            `None` is treated like 0.
            Default Value: `None`
          slope_variance: Positive floating point `Tensor` that controls the
            variance of the regression line slope. Must be broadcastable with
            `bias_variance`, `shift`, and inputs to `apply` and `matrix`
            methods. A value of `None` is treated like 1.
            Default Value: `None`
          shift: Floating point `Tensor` that controls the intercept with the
            x-axis of the linear interpolation. Must be broadcastable with
            `bias_variance`, `slope_variance`, and inputs to `apply` and `matrix`
            methods. A value of `None` is treated like 0, which results in having
            the intercept at the origin.
          feature_ndims: Python `int` number of rightmost dims to include in
            kernel computation.
            Default Value: 1
          validate_args: If `True`, parameters are checked for validity despite
            possibly degrading runtime performance.
            Default Value: `False`
          name: Python `str` name prefixed to Ops created by this class.
            Default Value: `'Linear'`


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'

    ```


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


return _build_custom_rv(distribution=distribution,<EOL>sample_shape=sample_shape,<EOL>value=value,<EOL>name=_simple_name(distribution))<EOL>

Wrap an existing distribution as a traceable random variable.

    This enables the use of custom or user-provided distributions in
    Edward models. Unlike a bare `RandomVariable` object, this method
    wraps the constructor so it is included in the Edward trace and its
    values can be properly intercepted and overridden.

    Where possible, you should prefer the built-in constructors
    (`ed.Normal`, etc); these simultaneously construct a Distribution
    and a RandomVariable object so that the distribution parameters
    themselves may be intercepted and overridden. RVs constructed via
    `as_random_variable()` have a fixed distribution and may not support
    program transformations (e.g, conjugate marginalization) that rely
    on overriding distribution parameters.

    Args:
      distribution: tfd.Distribution governing the distribution of the random
        variable, such as sampling and log-probabilities.
      sample_shape: tf.TensorShape of samples to draw from the random variable.
        Default is `()` corresponding to a single sample.
      value: Fixed tf.Tensor to associate with random variable. Must have shape
        `sample_shape + distribution.batch_shape + distribution.event_shape`.
        Default is to sample from random variable according to `sample_shape`.

    Returns:
      rv: a `RandomVariable` wrapping the provided distribution.

    #### Example

    ```python
    from tensorflow_probability import distributions as tfd
    from tensorflow_probability import edward2 as ed

    def model():
      # equivalent to ed.Normal(0., 1., name='x')
      return ed.as_random_variable(tfd.Normal(0., 1., name='x'))

    log_joint = ed.make_log_joint_fn(model)
    output = log_joint(x=2.)
    ```


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


try:<EOL><INDENT>_interceptor_stack.stack.append(interceptor)<EOL>yield<EOL><DEDENT>finally:<EOL><INDENT>_interceptor_stack.stack.pop()<EOL><DEDENT>

Python context manager for interception.

    Upon entry, an interception context manager pushes an interceptor onto a
    thread-local stack. Upon exiting, it pops the interceptor from the stack.

    Args:
      interceptor: Function which takes a callable `f` and inputs `*args`,
        `**kwargs`.

    Yields:
      None.

    #### Examples

    Interception controls the execution of Edward programs. Below we illustrate
    how to set the value of a specific random variable within a program.

    ```python
    from tensorflow_probability import edward2 as ed

    def model():
      return ed.Poisson(rate=1.5, name="y")

    def interceptor(f, *args, **kwargs):
      if kwargs.get("name") == "y":
        kwargs["value"] = 42
      return interceptable(f)(*args, **kwargs)

    with ed.interception(interceptor):
      y = model()

    with tf.Session() as sess:
      assert sess.run(y.value) == 42
    ```

    Wrapping `f` as `interceptable` allows interceptors down the stack to
    additionally modify this operation. Since the operation `f()` is not wrapped
    by default, we could have called it directly. Refer also to the example in
    `get_next_interceptor()` for more details on nested interceptors.


try:<EOL><INDENT>interceptor = _interceptor_stack.stack.pop()<EOL>yield interceptor<EOL><DEDENT>finally:<EOL><INDENT>_interceptor_stack.stack.append(interceptor)<EOL><DEDENT>

Yields the top-most interceptor on the thread-local interceptor stack.

    Operations may be intercepted by multiple nested interceptors. Once reached,
    an operation can be forwarded through nested interceptors until resolved.
    To allow for nesting, implement interceptors by re-wrapping their first
    argument (`f`) as an `interceptable`. To avoid nesting, manipulate the
    computation without using `interceptable`.

    This function allows for nesting by manipulating the thread-local interceptor
    stack, so that operations are intercepted in the order of interceptor nesting.

    #### Examples

    ```python
    from tensorflow_probability import edward2 as ed

    def model():
      x = ed.Normal(loc=0., scale=1., name="x")
      y = ed.Normal(loc=x, scale=1., name="y")
      return x + y

    def double(f, *args, **kwargs):
      return 2. * interceptable(f)(*args, **kwargs)

    def set_y(f, *args, **kwargs):
      if kwargs.get("name") == "y":
        kwargs["value"] = 0.42
      return interceptable(f)(*args, **kwargs)

    with interception(double):
      with interception(set_y):
        z = model()
    ```

    This will firstly put `double` on the stack, and then `set_y`,
    resulting in the stack:
    (TOP) set_y -> double -> apply (BOTTOM)

    The execution of `model` is then (top lines are current stack state):
    1) (TOP) set_y -> double -> apply (BOTTOM);
    `ed.Normal(0., 1., "x")` is intercepted by `set_y`, and as the name is not "y"
    the operation is simply forwarded to the next interceptor on the stack.

    2) (TOP) double -> apply (BOTTOM);
    `ed.Normal(0., 1., "x")` is intercepted by `double`, to produce
    `2*ed.Normal(0., 1., "x")`, with the operation being forwarded down the stack.

    3) (TOP) apply (BOTTOM);
    `ed.Normal(0., 1., "x")` is intercepted by `apply`, which simply calls the
    constructor.

    (At this point, the nested calls to `get_next_interceptor()`, produced by
    forwarding operations, exit, and the current stack is again:
    (TOP) set_y -> double -> apply (BOTTOM))

    4) (TOP) set_y -> double -> apply (BOTTOM);
    `ed.Normal(0., 1., "y")` is intercepted by `set_y`,
    the value of `y` is set to 0.42 and the operation is forwarded down the stack.

    5) (TOP) double -> apply (BOTTOM);
    `ed.Normal(0., 1., "y")` is intercepted by `double`, to produce
    `2*ed.Normal(0., 1., "y")`, with the operation being forwarded down the stack.

    6) (TOP) apply (BOTTOM);
    `ed.Normal(0., 1., "y")` is intercepted by `apply`, which simply calls the
    constructor.

    The final values for `x` and `y` inside of `model()` are tensors where `x` is
    a random draw from Normal(0., 1.) doubled, and `y` is a constant 0.84, thus
    z = 2 * Normal(0., 1.) + 0.84.


@functools.wraps(func)<EOL>def func_wrapped(*args, **kwargs):<EOL><INDENT>with get_next_interceptor() as interceptor:<EOL><INDENT>return interceptor(func, *args, **kwargs)<EOL><DEDENT><DEDENT>return func_wrapped<EOL>

Decorator that wraps `func` so that its execution is intercepted.

    The wrapper passes `func` to the interceptor for the current thread.

    If there is no next interceptor, we perform an "immediate" call to `func`.
    That is, `func` terminates without forwarding its execution to another
    interceptor.

    Args:
      func: Function to wrap.

    Returns:
      The decorated function.


tape_data = collections.OrderedDict({})<EOL>def record(f, *args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>name = kwargs.get("<STR_LIT:name>")<EOL>output = interceptable(f)(*args, **kwargs)<EOL>if name:<EOL><INDENT>tape_data[name] = output<EOL><DEDENT>return output<EOL><DEDENT>with interception(record):<EOL><INDENT>yield tape_data<EOL><DEDENT>

Context manager for recording interceptable executions onto a tape.

    Similar to `tf.GradientTape`, operations are recorded if they are executed
    within this context manager. In addition, the operation must be registered
    (wrapped) as `ed.interceptable`.

    Yields:
      tape: OrderedDict where operations are recorded in sequence. Keys are
        the `name` keyword argument to the operation (typically, a random
        variable's `name`) and values are the corresponding output of the
        operation. If the operation has no name, it is not recorded.

    #### Examples

    ```python
    from tensorflow_probability import edward2 as ed

    def probabilistic_matrix_factorization():
      users = ed.Normal(0., 1., sample_shape=[5000, 128], name="users")
      items = ed.Normal(0., 1., sample_shape=[7500, 128], name="items")
      ratings = ed.Normal(loc=tf.matmul(users, items, transpose_b=True),
                          scale=0.1,
                          name="ratings")
      return ratings

    with ed.tape() as model_tape:
      ratings = probabilistic_matrix_factorization()

    assert model_tape["users"].shape == (5000, 128)
    assert model_tape["items"].shape == (7500, 128)
    assert model_tape["ratings"] == ratings
    ```


@functools.wraps(attr)<EOL>def func(a, *args):<EOL><INDENT>return attr(a.value, *args)<EOL><DEDENT>return func<EOL>

Defers an operator overload to `attr`.

    Args:
      attr: Operator attribute to use.

    Returns:
      Function calling operator attribute.


if tensor.dtype.is_numpy_compatible:<EOL><INDENT>text = repr(tensor.numpy()) if is_repr else str(tensor.numpy())<EOL><DEDENT>else:<EOL><INDENT>text = "<STR_LIT>"<EOL><DEDENT>if "<STR_LIT:\n>" in text:<EOL><INDENT>text = "<STR_LIT:\n>" + text<EOL><DEDENT>return text<EOL>

Human-readable representation of a tensor's numpy value.


self._distribution = distribution<EOL>self._sample_shape = sample_shape<EOL>if value is not None:<EOL><INDENT>value = tf.cast(value, self.distribution.dtype)<EOL>value_shape = value.shape<EOL>expected_value_shape = self.sample_shape.concatenate(<EOL>self.distribution.batch_shape).concatenate(<EOL>self.distribution.event_shape)<EOL>if not value_shape.is_compatible_with(expected_value_shape):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>" % (expected_value_shape, value_shape))<EOL><DEDENT><DEDENT>self._value = value<EOL>

Create a new random variable.

        Args:
          distribution: tfd.Distribution governing the distribution of the random
            variable, such as sampling and log-probabilities.
          sample_shape: tf.TensorShape of samples to draw from the random variable.
            Default is `()` corresponding to a single sample.
          value: Fixed tf.Tensor to associate with random variable. Must have shape
            `sample_shape + distribution.batch_shape + distribution.event_shape`.
            Default is to sample from random variable according to `sample_shape`.

        Raises:
          ValueError: `value` has incompatible shape with
            `sample_shape + distribution.batch_shape + distribution.event_shape`.


return self._distribution<EOL>


return self.value.dtype<EOL>

`Dtype` of elements in this random variable.


if isinstance(self._sample_shape, tf.Tensor):<EOL><INDENT>return tf.TensorShape(tf.get_static_value(self._sample_shape))<EOL><DEDENT>return tf.TensorShape(self._sample_shape)<EOL>

Sample shape of random variable as a `TensorShape`.


with tf.compat.v1.name_scope(name):<EOL><INDENT>if isinstance(self._sample_shape, tf.Tensor):<EOL><INDENT>return self._sample_shape<EOL><DEDENT>return tf.convert_to_tensor(<EOL>value=self.sample_shape.as_list(), dtype=tf.int32)<EOL><DEDENT>

Sample shape of random variable as a 1-D `Tensor`.

        Args:
          name: name to give to the op

        Returns:
          sample_shape: `Tensor`.


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.


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


if not isinstance(self.value, ops.EagerTensor):<EOL><INDENT>raise NotImplementedError("<STR_LIT>")<EOL><DEDENT>return self.value.numpy()<EOL>

Value as NumPy array, only available for TF Eager.


return self.shape<EOL>


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


def log_joint_fn(*args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>log_probs = []<EOL>def interceptor(rv_constructor, *rv_args, **rv_kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>rv_name = rv_kwargs.get("<STR_LIT:name>")<EOL>if rv_name is None:<EOL><INDENT>raise KeyError("<STR_LIT>"<EOL>"<STR_LIT>".format(rv_constructor.__name__))<EOL><DEDENT>previously_specified_value = rv_kwargs.get("<STR_LIT:value>")<EOL>value = kwargs.get(rv_name, previously_specified_value)<EOL>if value is None:<EOL><INDENT>raise LookupError("<STR_LIT>"<EOL>"<STR_LIT>".format(rv_name))<EOL><DEDENT>rv_kwargs["<STR_LIT:value>"] = value<EOL>rv = rv_constructor(*rv_args, **rv_kwargs)<EOL>log_prob = tf.reduce_sum(input_tensor=rv.distribution.log_prob(rv.value))<EOL>log_probs.append(log_prob)<EOL>return rv<EOL><DEDENT>model_kwargs = _get_function_inputs(model, kwargs)<EOL>with interception(interceptor):<EOL><INDENT>model(*args, **model_kwargs)<EOL><DEDENT>log_prob = sum(log_probs)<EOL>return log_prob<EOL><DEDENT>return log_joint_fn<EOL>

Takes Edward probabilistic program and returns its log joint function.

    Args:
      model: Python callable which executes the generative process of a
        computable probability distribution using `ed.RandomVariable`s.

    Returns:
      A log-joint probability function. Its inputs are `model`'s original inputs
      and random variables which appear during the program execution. Its output
      is a scalar tf.Tensor.

    #### Examples

    Below we define Bayesian logistic regression as an Edward program,
    representing the model's generative process. We apply `make_log_joint_fn` in
    order to represent the model in terms of its joint probability function.

    ```python
    from tensorflow_probability import edward2 as ed

    def logistic_regression(features):
      coeffs = ed.Normal(loc=0., scale=1.,
                         sample_shape=features.shape[1], name="coeffs")
      outcomes = ed.Bernoulli(logits=tf.tensordot(features, coeffs, [[1], [0]]),
                              name="outcomes")
      return outcomes

    log_joint = ed.make_log_joint_fn(logistic_regression)

    features = tf.random_normal([3, 2])
    coeffs_value = tf.random_normal([2])
    outcomes_value = tf.round(tf.random_uniform([3]))
    output = log_joint(features, coeffs=coeffs_value, outcomes=outcomes_value)
    ```


if hasattr(f, "<STR_LIT>"):  <EOL><INDENT>f = f._func  <EOL><DEDENT>try:  <EOL><INDENT>argspec = inspect.getfullargspec(f)<EOL><DEDENT>except AttributeError:<EOL><INDENT>argspec = inspect.getargspec(f)<EOL><DEDENT>fkwargs = {k: v for k, v in six.iteritems(src_kwargs) if k in argspec.args}<EOL>return fkwargs<EOL>

Filters inputs to be compatible with function `f`'s signature.

    Args:
      f: Function according to whose input signature we filter arguments.
      src_kwargs: Keyword arguments to filter according to `f`.

    Returns:
      kwargs: Dict of key-value pairs in `src_kwargs` which exist in `f`'s
        signature.


def _fn(dtype, shape, name, trainable, add_variable_fn):<EOL><INDENT>"""<STR_LIT>"""<EOL>loc = add_variable_fn(<EOL>name=name + '<STR_LIT>',<EOL>shape=shape,<EOL>initializer=loc_initializer,<EOL>regularizer=loc_regularizer,<EOL>constraint=loc_constraint,<EOL>dtype=dtype,<EOL>trainable=trainable)<EOL>if is_singular:<EOL><INDENT>return loc, None<EOL><DEDENT>untransformed_scale = add_variable_fn(<EOL>name=name + '<STR_LIT>',<EOL>shape=shape,<EOL>initializer=untransformed_scale_initializer,<EOL>regularizer=untransformed_scale_regularizer,<EOL>constraint=untransformed_scale_constraint,<EOL>dtype=dtype,<EOL>trainable=trainable)<EOL>scale = (np.finfo(dtype.as_numpy_dtype).eps +<EOL>tf.nn.softplus(untransformed_scale))<EOL>return loc, scale<EOL><DEDENT>return _fn<EOL>

Makes closure which creates `loc`, `scale` params from `tf.get_variable`.

    This function produces a closure which produces `loc`, `scale` using
    `tf.get_variable`. The closure accepts the following arguments:

      dtype: Type of parameter's event.
      shape: Python `list`-like representing the parameter's event shape.
      name: Python `str` name prepended to any created (or existing)
        `tf.Variable`s.
      trainable: Python `bool` indicating all created `tf.Variable`s should be
        added to the graph collection `GraphKeys.TRAINABLE_VARIABLES`.
      add_variable_fn: `tf.get_variable`-like `callable` used to create (or
        access existing) `tf.Variable`s.

    Args:
      is_singular: Python `bool` indicating if `scale is None`. Default: `False`.
      loc_initializer: Initializer function for the `loc` parameters.
        The default is `tf.random_normal_initializer(mean=0., stddev=0.1)`.
      untransformed_scale_initializer: Initializer function for the `scale`
        parameters. Default value: `tf.random_normal_initializer(mean=-3.,
        stddev=0.1)`. This implies the softplus transformed result is initialized
        near `0`. It allows a `Normal` distribution with `scale` parameter set to
        this value to approximately act like a point mass.
      loc_regularizer: Regularizer function for the `loc` parameters.
        The default (`None`) is to use the `tf.get_variable` default.
      untransformed_scale_regularizer: Regularizer function for the `scale`
        parameters. The default (`None`) is to use the `tf.get_variable` default.
      loc_constraint: An optional projection function to be applied to the
        loc after being updated by an `Optimizer`. The function must take as input
        the unprojected variable and must return the projected variable (which
        must have the same shape). Constraints are not safe to use when doing
        asynchronous distributed training.
        The default (`None`) is to use the `tf.get_variable` default.
      untransformed_scale_constraint: An optional projection function to be
        applied to the `scale` parameters after being updated by an `Optimizer`
        (e.g. used to implement norm constraints or value constraints). The
        function must take as input the unprojected variable and must return the
        projected variable (which must have the same shape). Constraints are not
        safe to use when doing asynchronous distributed training. The default
        (`None`) is to use the `tf.get_variable` default.

    Returns:
      default_loc_scale_fn: Python `callable` which instantiates `loc`, `scale`
      parameters from args: `dtype, shape, name, trainable, add_variable_fn`.


loc_scale_fn = default_loc_scale_fn(<EOL>is_singular=is_singular,<EOL>loc_initializer=loc_initializer,<EOL>untransformed_scale_initializer=untransformed_scale_initializer,<EOL>loc_regularizer=loc_regularizer,<EOL>untransformed_scale_regularizer=untransformed_scale_regularizer,<EOL>loc_constraint=loc_constraint,<EOL>untransformed_scale_constraint=untransformed_scale_constraint)<EOL>def _fn(dtype, shape, name, trainable, add_variable_fn):<EOL><INDENT>"""<STR_LIT>"""<EOL>loc, scale = loc_scale_fn(dtype, shape, name, trainable, add_variable_fn)<EOL>if scale is None:<EOL><INDENT>dist = tfd.Deterministic(loc=loc)<EOL><DEDENT>else:<EOL><INDENT>dist = tfd.Normal(loc=loc, scale=scale)<EOL><DEDENT>batch_ndims = tf.size(input=dist.batch_shape_tensor())<EOL>return tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)<EOL><DEDENT>return _fn<EOL>

Creates a function to build Normal distributions with trainable params.

    This function produces a closure which produces `tfd.Normal`
    parameterized by a loc` and `scale` each created using `tf.get_variable`.

    Args:
      is_singular: Python `bool` if `True`, forces the special case limit of
        `scale->0`, i.e., a `Deterministic` distribution.
      loc_initializer: Initializer function for the `loc` parameters.
        The default is `tf.random_normal_initializer(mean=0., stddev=0.1)`.
      untransformed_scale_initializer: Initializer function for the `scale`
        parameters. Default value: `tf.random_normal_initializer(mean=-3.,
        stddev=0.1)`. This implies the softplus transformed result is initialized
        near `0`. It allows a `Normal` distribution with `scale` parameter set to
        this value to approximately act like a point mass.
      loc_regularizer: Regularizer function for the `loc` parameters.
      untransformed_scale_regularizer: Regularizer function for the `scale`
        parameters.
      loc_constraint: An optional projection function to be applied to the
        loc after being updated by an `Optimizer`. The function must take as input
        the unprojected variable and must return the projected variable (which
        must have the same shape). Constraints are not safe to use when doing
        asynchronous distributed training.
      untransformed_scale_constraint: An optional projection function to be
        applied to the `scale` parameters after being updated by an `Optimizer`
        (e.g. used to implement norm constraints or value constraints). The
        function must take as input the unprojected variable and must return the
        projected variable (which must have the same shape). Constraints are not
        safe to use when doing asynchronous distributed training.

    Returns:
      make_normal_fn: Python `callable` which creates a `tfd.Normal`
        using from args: `dtype, shape, name, trainable, add_variable_fn`.


def default_multivariate_normal_fn(dtype, shape, name, trainable,<EOL>add_variable_fn):

del name, trainable, add_variable_fn   <EOL>dist = tfd.Normal(loc=tf.zeros(shape, dtype), scale=dtype.as_numpy_dtype(<NUM_LIT:1>))<EOL>batch_ndims = tf.size(input=dist.batch_shape_tensor())<EOL>return tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)<EOL>

Creates multivariate standard `Normal` distribution.

    Args:
      dtype: Type of parameter's event.
      shape: Python `list`-like representing the parameter's event shape.
      name: Python `str` name prepended to any created (or existing)
        `tf.Variable`s.
      trainable: Python `bool` indicating all created `tf.Variable`s should be
        added to the graph collection `GraphKeys.TRAINABLE_VARIABLES`.
      add_variable_fn: `tf.get_variable`-like `callable` used to create (or
        access existing) `tf.Variable`s.

    Returns:
      Multivariate standard `Normal` distribution.


if function_type == '<STR_LIT>':<EOL><INDENT>function = tf.keras.utils.deserialize_keras_object(serial)<EOL><DEDENT>elif function_type == '<STR_LIT>':<EOL><INDENT>function = generic_utils.func_load(serial)<EOL><DEDENT>else:<EOL><INDENT>raise TypeError('<STR_LIT>', function_type)<EOL><DEDENT>return function<EOL>

Deserializes the Keras-serialized function.

    (De)serializing Python functions from/to bytecode is unsafe. Therefore we
    also use the function's type as an anonymous function ('lambda') or named
    function in the Python environment ('function'). In the latter case, this lets
    us use the Python scope to obtain the function rather than reload it from
    bytecode. (Note that both cases are brittle!)

    Keras-deserialized functions do not perform lexical scoping. Any modules that
    the function requires must be imported within the function itself.

    This serialization mimicks the implementation in `tf.keras.layers.Lambda`.

    Args:
      serial: Serialized Keras object: typically a dict, string, or bytecode.
      function_type: Python string denoting 'function' or 'lambda'.

    Returns:
      function: Function the serialized Keras object represents.

    #### Examples

    ```python
    serial, function_type = serialize_function(lambda x: x)
    function = deserialize_function(serial, function_type)
    assert function(2.3) == 2.3  # function is identity
    ```


if isinstance(func, types.LambdaType):<EOL><INDENT>return generic_utils.func_dump(func), '<STR_LIT>'<EOL><DEDENT>return func.__name__, '<STR_LIT>'<EOL>

Serializes function for Keras.

    (De)serializing Python functions from/to bytecode is unsafe. Therefore we
    return the function's type as an anonymous function ('lambda') or named
    function in the Python environment ('function'). In the latter case, this lets
    us use the Python scope to obtain the function rather than reload it from
    bytecode. (Note that both cases are brittle!)

    This serialization mimicks the implementation in `tf.keras.layers.Lambda`.

    Args:
      func: Python function to serialize.

    Returns:
      (serial, function_type): Serialized object, which is a tuple of its
      bytecode (if function is anonymous) or name (if function is named), and its
      function type.


return shape[:<NUM_LIT:1>] + shape[-<NUM_LIT:1>:] + shape[<NUM_LIT:1>:-<NUM_LIT:1>]<EOL>

Given a shape with channels last, returns a shape with channels first.

    Assumes we're dealing with the standard shapes as used by
    [TF's convolution operators][1].

    [1]: https://www.tensorflow.org/api_docs/python/tf/nn/convolution

    Args:
      shape: A sequence with each element corresponding to an axis in a shape.

    Returns:
      The transposed shape.


return shape[:<NUM_LIT:1>] + shape[<NUM_LIT:2>:] + shape[<NUM_LIT:1>:<NUM_LIT:2>]<EOL>

Given a shape with channels first, returns a shape with channels last.

    Assumes we're dealing with the standard shapes as used by
    [TF's convolution operators][1].

    [1]: https://www.tensorflow.org/api_docs/python/tf/nn/convolution

    Args:
      shape: A sequence with each element corresponding to an axis in a shape.

    Returns:
      The transposed shape.


def __init__(self, input_shape, data_format=None, **kwargs):

self._standardize_to_channels_last = (<EOL>data_format is not None and data_format.startswith('<STR_LIT>'))<EOL>self._rank = input_shape.rank<EOL>if self._standardize_to_channels_last:<EOL><INDENT>data_format = channels_first_to_last(data_format)<EOL>input_shape = tf.TensorShape(<EOL>channels_first_to_last(input_shape.as_list()))<EOL><DEDENT>super(CPUConvolution, self).__init__(<EOL>input_shape, data_format=data_format, **kwargs)<EOL>

Creates the convolution operator.

        Args:
          input_shape: A `TensorShape` input shape.
          data_format: A string describing the data format, see
            https://www.tensorflow.org/api_docs/python/tf/nn/convolution.
          **kwargs: Passed to `nn_ops.Convolution`.


if (isinstance(x, (Normal, tfd.Normal)) and<EOL>not isinstance(x, dtc._TensorCoercible) and<EOL>not isinstance(self, dtc._TensorCoercible)):<EOL><INDENT>return Normal(loc=self.loc + x.loc,<EOL>scale=tf.sqrt(self.scale**<NUM_LIT> + x.scale**<NUM_LIT>))<EOL><DEDENT>return super(Normal, self).__add__(x)<EOL>

Distributional arithmetic unless both have value semantics.


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 _value(self, dtype=None, name=None, as_ref=False):

<EOL>if as_ref:<EOL><INDENT>raise NotImplementedError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if self._concrete_value is None:<EOL><INDENT>if self._convert_to_tensor_fn is None:<EOL><INDENT>raise NotImplementedError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(type(self), self))<EOL><DEDENT>with self._name_scope('<STR_LIT:value>'):<EOL><INDENT>self._concrete_value = (self._convert_to_tensor_fn(self)<EOL>if callable(self._convert_to_tensor_fn)<EOL>else self._convert_to_tensor_fn)<EOL>if not tf.is_tensor(self._concrete_value):<EOL><INDENT>self._concrete_value = tfd._convert_to_tensor(<EOL>value=self._concrete_value,<EOL>name=name or '<STR_LIT>',<EOL>dtype=dtype,<EOL>dtype_hint=self.dtype)<EOL><DEDENT><DEDENT><DEDENT>return self._concrete_value<EOL>

Get the value returned by `tf.convert_to_tensor(distribution)`.

    Note: this function may mutate the distribution instance state by caching
    the concretized `Tensor` value.

    Args:
      dtype: Must return a `Tensor` with the given `dtype` if specified.
      name: If the conversion function creates a new `Tensor`, it should use the
        given `name` if specified.
      as_ref: `as_ref` is true, the function must return a `Tensor` reference,
        such as a `Variable`.
    Returns:
      concretized_distribution_value: `Tensor` identical to
      `tf.convert_to_tensor(distribution)`.

    #### Examples

    ```python
    tfd = tfp.distributions
    x = tfd.Normal(0.5, 1).set_tensor_conversion(tfd.Distribution.mean)

    x._value()
    # ==> tf.convert_to_tensor(x) ==> 0.5

    x._value() + 2
    # ==> tf.convert_to_tensor(x) + 2. ==> 2.5

    x + 2
    # ==> tf.convert_to_tensor(x) + 2. ==> 2.5
    ```


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.


dtype = tf.as_dtype(dtype or tf.keras.backend.floatx())<EOL>if isinstance(shape, tf.TensorShape):<EOL><INDENT>shape_dtype = tf.int32<EOL>shape_ = np.int32(shape)<EOL><DEDENT>else:<EOL><INDENT>if not tf.is_tensor(shape):<EOL><INDENT>shape = tf.convert_to_tensor(<EOL>value=shape, dtype_hint=tf.int32, name='<STR_LIT>')<EOL><DEDENT>shape_dtype = shape.dtype.base_dtype<EOL>shape_ = tf.get_static_value(shape, partial=True)<EOL><DEDENT>sizes_ = tf.get_static_value(self.sizes)<EOL>if sizes_ is not None:<EOL><INDENT>sizes_ = np.array(sizes_, shape_dtype.as_numpy_dtype)<EOL><DEDENT>assertions = []<EOL>message = '<STR_LIT>'<EOL>n = shape[-<NUM_LIT:1>] if shape_ is None or shape_[-<NUM_LIT:1>] is None else shape_[-<NUM_LIT:1>]<EOL>if sizes_ is not None and not tf.is_tensor(n):<EOL><INDENT>if sum(sizes_) != n:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif self.validate_args:<EOL><INDENT>assertions.append(tf.compat.v1.assert_equal(<EOL>shape[-<NUM_LIT:1>], tf.reduce_sum(input_tensor=self.sizes), message=message))<EOL><DEDENT>s = (shape[:-<NUM_LIT:1>]<EOL>if shape_ is None or any(s is None for s in shape_[:-<NUM_LIT:1>])<EOL>else shape_[:-<NUM_LIT:1>])<EOL>if sizes_ is not None and isinstance(s, (np.ndarray, np.generic)):<EOL><INDENT>return tf.concat([<EOL>tf.keras.initializers.get(init)(np.concatenate([<EOL>s, np.array([e], shape_dtype.as_numpy_dtype)], axis=-<NUM_LIT:1>), dtype)<EOL>for init, e in zip(self.initializers, sizes_.tolist())<EOL>], axis=-<NUM_LIT:1>)<EOL><DEDENT>sizes = tf.split(self.sizes, len(self.initializers))<EOL>return tf.concat([<EOL>tf.keras.initializers.get(init)(tf.concat([s, e], axis=-<NUM_LIT:1>), dtype)<EOL>for init, e in zip(self.initializers, sizes)<EOL>], axis=-<NUM_LIT:1>)<EOL>

Returns a tensor object initialized as specified by the initializer.

        Args:
          shape: Shape of the tensor.
          dtype: Optional dtype of the tensor. If not provided will return tensor
           of `tf.float32`.


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.


return cls(**{<EOL>'<STR_LIT>': [tf.compat.v2.initializers.deserialize(init)<EOL>for init in config.get('<STR_LIT>', [])],<EOL>'<STR_LIT>': config.get('<STR_LIT>', []),<EOL>'<STR_LIT>': config.get('<STR_LIT>', False),<EOL>})<EOL>

Instantiates an initializer from a configuration dictionary.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [event_shape]):<EOL><INDENT>event_shape = tf.convert_to_tensor(<EOL>value=event_shape, dtype=tf.int32, name='<STR_LIT>')<EOL>event_shape_const = tf.get_static_value(event_shape)<EOL>if event_shape_const is not None:<EOL><INDENT>return np.prod(event_shape_const)<EOL><DEDENT>else:<EOL><INDENT>return tf.reduce_prod(input_tensor=event_shape)<EOL><DEDENT><DEDENT>

Computes the number of elements in a tensor with shape `event_shape`.

    Args:
      event_shape: A tensor shape.
      name: The name to use for the tensor op to compute the number of elements
        (if such an op needs to be created).

    Returns:
      event_size: The number of elements in `tensor_shape`.  Returns a numpy int
      when the number of elements can be computed immediately.  Otherwise, returns
      a scalar tensor.


def _eval_all_one_hot(fn, dist, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>'):<EOL><INDENT>event_size = dist.event_shape_tensor()[-<NUM_LIT:1>]<EOL>batch_ndims = tf.size(input=dist.batch_shape_tensor())<EOL>x = tf.reshape(<EOL>tf.eye(event_size, dtype=dist.dtype),<EOL>shape=tf.pad(<EOL>tensor=tf.ones(batch_ndims, tf.int32),<EOL>paddings=[[<NUM_LIT:1>, <NUM_LIT:1>]],<EOL>constant_values=event_size))<EOL>perm = tf.pad(tensor=tf.range(<NUM_LIT:1>, batch_ndims + <NUM_LIT:1>), paddings=[[<NUM_LIT:0>, <NUM_LIT:1>]])<EOL>return tf.transpose(a=fn(dist, x), perm=perm)<EOL><DEDENT>

OneHotCategorical helper computing probs, cdf, etc over its support.


if identifier is None:<EOL><INDENT>return None<EOL><DEDENT>if isinstance(identifier, six.string_types):<EOL><INDENT>identifier = str(identifier)<EOL>return _deserialize(identifier)<EOL><DEDENT>if isinstance(identifier, dict):<EOL><INDENT>return _deserialize(identifier)<EOL><DEDENT>if isinstance(identifier, property):<EOL><INDENT>identifier = identifier.fget<EOL><DEDENT>if callable(identifier):<EOL><INDENT>return identifier<EOL><DEDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>', identifier)<EOL>

Return a convert-to-tensor func, given a name, config, callable, etc.


<EOL>if isinstance(make_distribution_fn, six.string_types):<EOL><INDENT>make_distribution_fn = _deserialize_function(make_distribution_fn)<EOL><DEDENT>convert_to_tensor_fn = _get_convert_to_tensor_fn(convert_to_tensor_fn)<EOL>kwargs.pop('<STR_LIT>', None)<EOL>def _fn(*fargs, **fkwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>d = make_distribution_fn(*fargs, **fkwargs)<EOL>value_is_seq = isinstance(d.dtype, collections.Sequence)<EOL>maybe_composite_convert_to_tensor_fn = (<EOL>(lambda d: tensor_tuple.TensorTuple(convert_to_tensor_fn(d)))<EOL>if value_is_seq else convert_to_tensor_fn)<EOL>distribution = dtc._TensorCoercible(  <EOL>distribution=d,<EOL>convert_to_tensor_fn=maybe_composite_convert_to_tensor_fn)<EOL>value = distribution._value()  <EOL>value._tfp_distribution = distribution  <EOL>if value_is_seq:<EOL><INDENT>value.shape = value[-<NUM_LIT:1>].shape<EOL>value.get_shape = value[-<NUM_LIT:1>].get_shape<EOL>value.dtype = value[-<NUM_LIT:1>].dtype<EOL>distribution.shape = value[-<NUM_LIT:1>].shape<EOL>distribution.get_shape = value[-<NUM_LIT:1>].get_shape<EOL><DEDENT>else:<EOL><INDENT>distribution.shape = value.shape<EOL>distribution.get_shape = value.get_shape<EOL><DEDENT>return distribution, value<EOL><DEDENT>super(DistributionLambda, self).__init__(_fn, **kwargs)<EOL>self._make_distribution_fn = make_distribution_fn<EOL>self._convert_to_tensor_fn = convert_to_tensor_fn<EOL>self._enter_dunder_call = False<EOL>

Create a `DistributionLambda` Keras layer.

        Args:
          make_distribution_fn: Python `callable` that takes previous layer outputs
            and returns a `tfd.Distribution` instance.
          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`.
          **kwargs: Additional keyword arguments passed to `tf.keras.Layer`.


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

Returns the config of this layer.

        This Layer's `make_distribution_fn` is serialized via a library built on
        Python pickle.  This serialization of Python functions is provided for
        convenience, but:

          1. The use of this format for long-term storage of models is discouraged.
             In particular, it may not be possible to deserialize in a different
             version of Python.

          2. While serialization is generally supported for lambdas, local
             functions, and static methods (and closures over these constructs),
             complex functions may fail to serialize.

          3. `Tensor` objects (and functions referencing `Tensor` objects) can only
             be serialized when the tensor value is statically known.  (Such Tensors
             are serialized as numpy arrays.)

        Instead of relying on `DistributionLambda.get_config`, consider subclassing
        `DistributionLambda` and directly implementing Keras serialization via
        `get_config` / `from_config`.

        NOTE: At the moment, `DistributionLambda` can only be serialized if the
        `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"`)


def _variable_creator(self, next_creator, **kwargs):

return next_creator(**kwargs)<EOL>

Overrides `tfkl.Lambda` method causing error in weight initialization.


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

Initialize the `MultivariateNormalTriL` layer.

        Args:
          event_size: Scalar `int` representing the size 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`.
          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.


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[event_size]):<EOL><INDENT>return event_size + event_size * (event_size + <NUM_LIT:1>) // <NUM_LIT:2><EOL><DEDENT>

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


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

Initialize the `OneHotCategorical` layer.

        Args:
          event_size: Scalar `int` representing the size 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>


return event_size<EOL>


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, num_components,<EOL>dtype=None, validate_args=False, name=None):<DEDENT>

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[params, event_size, num_components]):<EOL><INDENT>dist = MixtureSameFamily.new(<EOL>params,<EOL>num_components,<EOL>OneHotCategorical(<EOL>event_size,<EOL>validate_args=False,  <EOL>name=name),<EOL>validate_args=validate_args,<EOL>name=name)<EOL>dist._mean = functools.partial(<EOL>_eval_all_one_hot, tfd.Distribution.prob, dist)<EOL>dist.log_mean = functools.partial(<EOL>_eval_all_one_hot, tfd.Distribution.log_prob, dist)<EOL>return dist<EOL><DEDENT>


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

with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[event_size, num_components]):<EOL><INDENT>return MixtureSameFamily.params_size(<EOL>num_components,<EOL>OneHotCategorical.params_size(event_size, name=name),<EOL>name=name)<EOL><DEDENT>


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_shape=(), dtype=None, 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>new_shape = tf.concat([<EOL>tf.shape(input=params)[:-<NUM_LIT:1>],<EOL>event_shape,<EOL>], axis=<NUM_LIT:0>)<EOL>dist = tfd.Independent(<EOL>tfd.Bernoulli(<EOL>logits=tf.reshape(params, new_shape),<EOL>dtype=dtype or params.dtype.base_dtype,<EOL>validate_args=validate_args),<EOL>reinterpreted_batch_ndims=tf.size(input=event_shape),<EOL>validate_args=validate_args)<EOL>dist._logits = dist.distribution._logits  <EOL>dist._probs = dist.distribution._probs  <EOL>dist.logits = tfd.Bernoulli.logits<EOL>dist.probs = tfd.Bernoulli.probs<EOL>return dist<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 _event_size(<EOL>event_shape, name=name or '<STR_LIT>')<EOL><DEDENT>


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


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


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(IndependentLogistic, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>


convert_to_tensor_fn = _get_convert_to_tensor_fn(convert_to_tensor_fn)<EOL>kwargs.pop('<STR_LIT>', None)<EOL>super(IndependentNormal, self).__init__(<EOL>lambda t: IndependentNormal.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 `IndependentNormal` 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>loc_params, scale_params = tf.split(params, <NUM_LIT:2>, axis=-<NUM_LIT:1>)<EOL>return tfd.Independent(<EOL>tfd.Normal(<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>


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(IndependentNormal, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<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`.


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>


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>


super(MixtureSameFamily, self).__init__(<EOL>lambda t: MixtureSameFamily.new(  <EOL>t, num_components, component_layer, validate_args),<EOL>convert_to_tensor_fn,<EOL>**kwargs)<EOL>

Initialize the `MixtureSameFamily` distribution layer.

        Args:
          num_components: Number of component distributions in the mixture
            distribution.
          component_layer: Python `callable` that, given a tensor of shape
            `batch_shape + [num_components, component_params_size]`, returns a
            `tfd.Distribution`-like instance that implements the component
            distribution (with batch shape `batch_shape + [num_components]`) --
            e.g., a TFP distribution layer.
          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`.


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


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

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[num_components, component_params_size]):<EOL><INDENT>num_components = tf.convert_to_tensor(<EOL>value=num_components, name='<STR_LIT>', dtype_hint=tf.int32)<EOL>component_params_size = tf.convert_to_tensor(<EOL>value=component_params_size, name='<STR_LIT>')<EOL>num_components = dist_util.prefer_static_value(num_components)<EOL>component_params_size = dist_util.prefer_static_value(<EOL>component_params_size)<EOL>return num_components + num_components * component_params_size<EOL><DEDENT>

Number of `params` needed to create a `MixtureSameFamily` distribution.

        Arguments:
          num_components: Number of component distributions in the mixture
            distribution.
          component_params_size: Number of parameters needed to create a single
            component distribution.
          name: The name to use for the op to compute the number of parameters
            (if such an op needs to be created).

        Returns:
         params_size: The number of parameters needed to create the mixture
           distribution.


convert_to_tensor_fn = _get_convert_to_tensor_fn(convert_to_tensor_fn)<EOL>kwargs.pop('<STR_LIT>', None)<EOL>super(MixtureNormal, self).__init__(<EOL>lambda t: MixtureNormal.new(  <EOL>t, num_components, event_shape, validate_args),<EOL>convert_to_tensor_fn,<EOL>**kwargs)<EOL>self._num_components = num_components<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 `MixtureNormal` distribution layer.

        Args:
          num_components: Number of component distributions in the mixture
            distribution.
          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`.


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

return MixtureSameFamily.new(<EOL>params,<EOL>num_components,<EOL>IndependentNormal(event_shape, validate_args=validate_args, name=name),<EOL>validate_args=validate_args,<EOL>name=name)<EOL>


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


config = {<EOL>'<STR_LIT>': self._num_components,<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(MixtureNormal, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>


convert_to_tensor_fn = _get_convert_to_tensor_fn(convert_to_tensor_fn)<EOL>kwargs.pop('<STR_LIT>', None)<EOL>super(MixtureLogistic, self).__init__(<EOL>lambda t: MixtureLogistic.new(  <EOL>t, num_components, event_shape, validate_args),<EOL>convert_to_tensor_fn,<EOL>**kwargs)<EOL>self._num_components = num_components<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 `MixtureLogistic` distribution layer.

        Args:
          num_components: Number of component distributions in the mixture
            distribution.
          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`.


return MixtureSameFamily.new(<EOL>params,<EOL>num_components,<EOL>IndependentLogistic(<EOL>event_shape, validate_args=validate_args, name=name),<EOL>validate_args=validate_args,<EOL>name=name)<EOL>


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


config = {<EOL>'<STR_LIT>': self._num_components,<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(MixtureLogistic, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>


super(VariationalGaussianProcess, self).__init__(<EOL>lambda x: VariationalGaussianProcess.new(  <EOL>x,<EOL>kernel_provider=self._kernel_provider,<EOL>event_shape=self._event_shape,<EOL>inducing_index_points=self._inducing_index_points,<EOL>variational_inducing_observations_loc=(<EOL>self._variational_inducing_observations_loc),<EOL>variational_inducing_observations_scale=(<EOL>self._variational_inducing_observations_scale),<EOL>mean_fn=self._mean_fn,<EOL>observation_noise_variance=tf.nn.softplus(<EOL>self._unconstrained_observation_noise_variance),<EOL>jitter=self._jitter))<EOL>tmp_kernel = kernel_provider.kernel<EOL>self._dtype = tmp_kernel.dtype.as_numpy_dtype<EOL>self._feature_ndims = tmp_kernel.feature_ndims<EOL>self._num_inducing_points = num_inducing_points<EOL>self._event_shape = tf.TensorShape(event_shape)<EOL>self._mean_fn = mean_fn<EOL>self._jitter = jitter<EOL>self._inducing_index_points_initializer = inducing_index_points_initializer<EOL>self._unconstrained_observation_noise_variance_initializer = (<EOL>unconstrained_observation_noise_variance_initializer)<EOL>self._kernel_provider = kernel_provider<EOL>

Construct a VariationalGaussianProcess Layer.

        Args:
          num_inducing_points: number of inducing points in the
            VariationalGaussianProcess distribution.
          kernel_provider: a `Layer` instance equipped with an @property, which
            yields a `PositiveSemidefiniteKernel` instance. The latter is used to
            parameterize the constructed VariationalGaussianProcess distribution
            returned by calling the layer.
          event_shape: the shape of the output of the layer. This translates to a
            batch of underlying VariationalGaussianProcess distribtuions. For
            example, `event_shape = [3]` means we are modeling a batch of 3
            distributions over functions. We can think of this as a distrbution over
            3-dimensional vector-valued functions.
          inducing_index_points_initializer: a `tf.keras.initializer.Initializer`
            used to initialize the trainable `inducing_index_points` variables.
            Training VGP's is pretty sensitive to choice of initial inducing index
            point locations. A reasonable heuristic is to scatter them near the
            data, not too close to each other.
          unconstrained_observation_noise_variance_initializer: a
            `tf.keras.initializer.Initializer` used to initialize the unconstrained
            observation noise variable. The observation noise variance is computed
            from this variable via the `tf.nn.softplus` function.
          mean_fn: a callable that maps layer inputs to mean function values. Passed
            to the mean_fn parameter of VariationalGaussianProcess distribution. If
            omitted, defaults to a constant function with trainable variable value.
          jitter: a small term added to the diagonal of various kernel matrices for
            numerical stability.
          name: name to give to this layer and the scope of ops and variables it
            contains.


<EOL>super(_ConvVariational, self).__init__(<EOL>activity_regularizer=activity_regularizer,<EOL>**kwargs)<EOL>self.rank = rank<EOL>self.filters = filters<EOL>self.kernel_size = tf_layers_util.normalize_tuple(<EOL>kernel_size, rank, '<STR_LIT>')<EOL>self.strides = tf_layers_util.normalize_tuple(strides, rank, '<STR_LIT>')<EOL>self.padding = tf_layers_util.normalize_padding(padding)<EOL>self.data_format = tf_layers_util.normalize_data_format(data_format)<EOL>self.dilation_rate = tf_layers_util.normalize_tuple(<EOL>dilation_rate, rank, '<STR_LIT>')<EOL>self.activation = tf.keras.activations.get(activation)<EOL>self.input_spec = tf.keras.layers.InputSpec(ndim=self.rank + <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:
          rank: An integer, the rank of the convolution, e.g. "2" for 2D
            convolution.
          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 n integers, specifying the
            length of the convolution window.
          strides: An integer or tuple/list of n integers,
            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, ...,
            channels)` while `channels_first` corresponds to inputs with shape
            `(batch, channels, ...)`.
          dilation_rate: An integer or tuple/list of n integers, 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}


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.


config = {<EOL>'<STR_LIT>': self.filters,<EOL>'<STR_LIT>': self.kernel_size,<EOL>'<STR_LIT>': self.strides,<EOL>'<STR_LIT>': self.padding,<EOL>'<STR_LIT>': self.data_format,<EOL>'<STR_LIT>': self.dilation_rate,<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(_ConvVariational, 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.


config = config.copy()<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>serial = config[function_key]<EOL>function_type = config.pop(function_key + '<STR_LIT>')<EOL>if serial is not None:<EOL><INDENT>config[function_key] = tfp_layers_util.deserialize_function(<EOL>serial,<EOL>function_type=function_type)<EOL><DEDENT><DEDENT>return cls(**config)<EOL>

Creates a layer from its config.

        This method is the reverse of `get_config`, capable of instantiating the
        same layer from the config dictionary.

        Args:
          config: A Python dictionary, typically the output of `get_config`.

        Returns:
          layer: A layer instance.


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}


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 2 integers, specifying the
            height and width of the 2D convolution window.
            Can be a single integer to specify the same value for
            all spatial dimensions.
          strides: An integer or tuple/list of 2 integers,
            specifying the strides of the convolution along the 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, height,
            width, channels)` while `channels_first` corresponds to inputs with
            shape `(batch, channels, height, width)`.
          dilation_rate: An integer or tuple/list of 2 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}


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}


Construct layer.

        Args:
          rank: An integer, the rank of the convolution, e.g. "2" for 2D
            convolution.
          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 n integers, specifying the
            length of the convolution window.
          strides: An integer or tuple/list of n integers,
            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, ...,
            channels)` while `channels_first` corresponds to inputs with shape
            `(batch, channels, ...)`.
          dilation_rate: An integer or tuple/list of n integers, 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).


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


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 2 integers, specifying the
            height and width of the 2D convolution window.
            Can be a single integer to specify the same value for
            all spatial dimensions.
          strides: An integer or tuple/list of 2 integers,
            specifying the strides of the convolution along the 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, height,
            width, channels)` while `channels_first` corresponds to inputs with
            shape `(batch, channels, height, width)`.
          dilation_rate: An integer or tuple/list of 2 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}
          seed: Python scalar `int` which initializes the random number
            generator. Default value: `None` (i.e., use global seed).


log_avg_prob = (<EOL>tf.reduce_logsumexp(input_tensor=-tf.nn.softplus(-t), axis=<NUM_LIT:0>) -<EOL>tf.math.log(tf.cast(tf.shape(input=t)[<NUM_LIT:0>], t.dtype)))<EOL>return log_avg_prob - tf.math.log1p(-tf.exp(log_avg_prob))<EOL>

Computes `logit(mean(expit(t)))` in a numerically stable manner.


paddings = tf.concat(<EOL>[tf.zeros([tf.rank(x) - <NUM_LIT:1>, <NUM_LIT:2>], dtype=tf.int32), [[<NUM_LIT:1>, <NUM_LIT:0>]]], axis=<NUM_LIT:0>)<EOL>return tf.pad(tensor=x, paddings=paddings, constant_values=value)<EOL>


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.


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 _distribution_to_params(self, distribution, batch_shape):

raise NotImplementedError<EOL>

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


raise NotImplementedError<EOL>


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>


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