parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>logits=logits,<EOL>probs=probs,<EOL>validate_args=validate_args,<EOL>multidimensional=True,<EOL>name=name)<EOL>if validate_args:<EOL><INDENT>self._logits = distribution_util.embed_check_categorical_event_shape(<EOL>self._logits)<EOL><DEDENT>logits_shape_static = tensorshape_util.with_rank_at_least(<EOL>self._logits.shape, <NUM_LIT:1>)<EOL>if tensorshape_util.rank(logits_shape_static) is not None:<EOL><INDENT>self._batch_rank = tf.convert_to_tensor(<EOL>value=tensorshape_util.rank(logits_shape_static) - <NUM_LIT:1>,<EOL>dtype=tf.int32,<EOL>name="<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>self._batch_rank = tf.rank(self._logits) - <NUM_LIT:1><EOL><DEDENT><DEDENT>logits_shape = tf.shape(input=self._logits, name="<STR_LIT>")<EOL>num_categories = tf.compat.dimension_value(logits_shape_static[-<NUM_LIT:1>])<EOL>if num_categories is not None:<EOL><INDENT>self._num_categories = tf.convert_to_tensor(<EOL>value=num_categories, dtype=tf.int32, name="<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>self._num_categories = logits_shape[self._batch_rank]<EOL><DEDENT><DEDENT>if tensorshape_util.is_fully_defined(logits_shape_static[:-<NUM_LIT:1>]):<EOL><INDENT>self._batch_shape_val = tf.constant(<EOL>logits_shape_static[:-<NUM_LIT:1>].as_list(),<EOL>dtype=tf.int32,<EOL>name="<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>self._batch_shape_val = logits_shape[:-<NUM_LIT:1>]<EOL><DEDENT><DEDENT><DEDENT>super(Categorical, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._logits,<EOL>self._probs],<EOL>name=name)<EOL>

Initialize Categorical distributions using class log-probabilities.

        Args:
          logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities
            of a set of Categorical distributions. The first `N - 1` dimensions
            index into a batch of independent distributions and the last dimension
            represents a vector of logits for each class. Only one of `logits` or
            `probs` should be passed in.
          probs: An N-D `Tensor`, `N >= 1`, representing the probabilities
            of a set of Categorical distributions. The first `N - 1` dimensions
            index into a batch of independent distributions and the last dimension
            represents a vector of probabilities for each class. Only one of
            `logits` or `probs` should be passed in.
          dtype: The type of the event samples (default: int32).
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


return self._num_categories<EOL>

Scalar `int32` tensor: the number of categories.


return self._num_categories<EOL>


return self._logits<EOL>


return self._probs<EOL>


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

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


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

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

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

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


def _resolve_distribution_names(dist_fn_args, dist_names, leaf_name):

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

Uses arg names to resolve distribution names.


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

Returns the distribution's required args.


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

Calculate the KL divergence between two `JointDistributionSequential`s.

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

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

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


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

parameters = dict(locals())<EOL>with tf.name_scope(name or '<STR_LIT>') as name:<EOL><INDENT>self._model = model<EOL>self._build(model)<EOL>self._most_recently_built_distributions = [<EOL>None if a else d() for d, a<EOL>in zip(self._dist_fn_wrapped, self._dist_fn_args)]<EOL>self._always_use_specified_sample_shape = False<EOL>super(JointDistributionSequential, self).__init__(<EOL>dtype=None,  <EOL>reparameterization_type=None,  <EOL>validate_args=validate_args,<EOL>allow_nan_stats=False,<EOL>parameters=parameters,<EOL>graph_parents=[],<EOL>name=name)<EOL>self._unflatten(self._flatten(model))<EOL><DEDENT>

Construct the `JointDistributionSequential` distribution.

        Args:
          model: Python list of either tfd.Distribution instances and/or
            lambda functions which take the `k` previous distributions and returns a
            new tfd.Distribution instance.
          validate_args: Python `bool`.  Whether to validate input with asserts.
            If `validate_args` is `False`, and the inputs are invalid,
            correct behavior is not guaranteed.
            Default value: `False`.
          name: The name for ops managed by the distribution.
            Default value: `None` (i.e., `"JointDistributionSequential"`).


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

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


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

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

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

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

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

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

        #### Example

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


if any(self._dist_fn_args):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>')<EOL><DEDENT>return sum(joint_distribution_lib.maybe_check_wont_broadcast(<EOL>(d().entropy() for d in self._dist_fn_wrapped),<EOL>self.validate_args))<EOL>


def draw_sample(num_samples, num_classes, logits, num_trials, dtype, seed):

with tf.name_scope("<STR_LIT>"):<EOL><INDENT>num_trials = tf.ones_like(<EOL>logits[..., <NUM_LIT:0>], dtype=num_trials.dtype) * num_trials<EOL>logits = tf.ones_like(<EOL>num_trials[..., tf.newaxis], dtype=logits.dtype) * logits<EOL>flat_logits = tf.reshape(logits, [-<NUM_LIT:1>, num_classes])<EOL>flat_num_trials = num_samples * tf.reshape(num_trials, [-<NUM_LIT:1>])  <EOL>def _sample_one_batch_member(args):<EOL><INDENT>logits, num_cat_samples = args[<NUM_LIT:0>], args[<NUM_LIT:1>]  <EOL>x = tf.random.categorical(<EOL>logits[tf.newaxis, ...], num_cat_samples, seed=seed)<EOL>x = tf.reshape(x, shape=[num_samples, -<NUM_LIT:1>])  <EOL>x = tf.one_hot(<EOL>x, depth=num_classes)  <EOL>x = tf.reduce_sum(input_tensor=x, axis=-<NUM_LIT:2>)  <EOL>return tf.cast(x, dtype=dtype)<EOL><DEDENT>x = tf.map_fn(<EOL>_sample_one_batch_member, [flat_logits, flat_num_trials],<EOL>dtype=dtype)  <EOL>x = tf.transpose(a=x, perm=[<NUM_LIT:1>, <NUM_LIT:0>, <NUM_LIT:2>])<EOL>final_shape = tf.concat([[num_samples],<EOL>tf.shape(input=num_trials), [num_classes]],<EOL>axis=<NUM_LIT:0>)<EOL>x = tf.reshape(x, final_shape)<EOL>return x<EOL><DEDENT>

Sample a multinomial.

    The batch shape is given by broadcasting num_trials with
    remove_last_dimension(logits).

    Args:
      num_samples: Python int or singleton integer Tensor: number of multinomial
        samples to draw.
      num_classes: Python int or singleton integer Tensor: number of classes.
      logits: Floating Tensor with last dimension k, of (unnormalized) logit
        probabilities per class.
      num_trials: Tensor of number of categorical trials each multinomial consists
        of.  num_trials[..., tf.newaxis] must broadcast with logits.
      dtype: dtype at which to emit samples.
      seed: Random seed.

    Returns:
      samples: Tensor of given dtype and shape [n] + batch_shape + [k].


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

Initialize a batch of Multinomial distributions.

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


return self._total_count<EOL>

Number of trials used to construct a sample.


return self._probs<EOL>

Probability of drawing a `1` in that coordinate.


if not self.validate_args:<EOL><INDENT>return counts<EOL><DEDENT>counts = distribution_util.embed_check_nonnegative_integer_form(counts)<EOL>return distribution_util.with_dependencies([<EOL>assert_util.assert_equal(<EOL>self.total_count,<EOL>tf.reduce_sum(input_tensor=counts, axis=-<NUM_LIT:1>),<EOL>message="<STR_LIT>"),<EOL>], counts)<EOL>

Check counts for proper shape, values, then return tensor version.


def calculate_reshape(original_shape, new_shape, validate=False, name=None):

batch_shape_static = tensorshape_util.constant_value_as_shape(new_shape)<EOL>if tensorshape_util.is_fully_defined(batch_shape_static):<EOL><INDENT>return np.int32(batch_shape_static), batch_shape_static, []<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>original_size = tf.reduce_prod(input_tensor=original_shape)<EOL>implicit_dim = tf.equal(new_shape, -<NUM_LIT:1>)<EOL>size_implicit_dim = (<EOL>original_size // tf.maximum(<NUM_LIT:1>, -tf.reduce_prod(input_tensor=new_shape)))<EOL>new_ndims = tf.shape(input=new_shape)<EOL>expanded_new_shape = tf.where(  <EOL>implicit_dim, tf.fill(new_ndims, size_implicit_dim), new_shape)<EOL>validations = [] if not validate else [  <EOL>assert_util.assert_rank(<EOL>original_shape, <NUM_LIT:1>, message="<STR_LIT>"),<EOL>assert_util.assert_rank(<EOL>new_shape, <NUM_LIT:1>, message="<STR_LIT>"),<EOL>assert_util.assert_less_equal(<EOL>tf.math.count_nonzero(implicit_dim, dtype=tf.int32),<EOL><NUM_LIT:1>,<EOL>message="<STR_LIT>"),<EOL>assert_util.assert_positive(<EOL>expanded_new_shape, message="<STR_LIT>"),<EOL>assert_util.assert_equal(<EOL>tf.reduce_prod(input_tensor=expanded_new_shape),<EOL>original_size,<EOL>message="<STR_LIT>"),<EOL>]<EOL>return expanded_new_shape, batch_shape_static, validations<EOL><DEDENT>

Calculates the reshaped dimensions (replacing up to one -1 in reshape).


if tensorshape_util.rank(batch_shape.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(batch_shape.shape) != <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>tensorshape_util.rank(batch_shape.shape)))<EOL><DEDENT><DEDENT>batch_shape_static = tensorshape_util.constant_value_as_shape(batch_shape)<EOL>batch_size_static = tensorshape_util.num_elements(batch_shape_static)<EOL>dist_batch_size_static = tensorshape_util.num_elements(<EOL>distribution.batch_shape)<EOL>if batch_size_static is not None and dist_batch_size_static is not None:<EOL><INDENT>if batch_size_static != dist_batch_size_static:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>batch_size_static, dist_batch_size_static))<EOL><DEDENT><DEDENT>if tensorshape_util.dims(batch_shape_static) is not None:<EOL><INDENT>if any(<EOL>tf.compat.dimension_value(dim) is not None and<EOL>tf.compat.dimension_value(dim) < <NUM_LIT:1> for dim in batch_shape_static):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT><DEDENT>

Helper to __init__ which makes or raises assertions.


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

Construct BatchReshape distribution.

        Args:
          distribution: The base distribution instance to reshape. Typically an
            instance of `Distribution`.
          batch_shape: Positive `int`-like vector-shaped `Tensor` representing
            the new shape of the batch dimensions. Up to one dimension may contain
            `-1`, meaning the remainder of the batch size.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: The name to give Ops created by the initializer.
            Default value: `"BatchReshape" + distribution.name`.

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


x_ndims = (<EOL>tf.rank(x) if tensorshape_util.rank(x.shape) is None else<EOL>tensorshape_util.rank(x.shape))<EOL>event_ndims = (<EOL>tf.size(input=self.event_shape_tensor())<EOL>if tensorshape_util.rank(self.event_shape) is None else<EOL>tensorshape_util.rank(self.event_shape))<EOL>batch_ndims = (<EOL>tf.size(input=self._batch_shape_unexpanded)<EOL>if tensorshape_util.rank(self.batch_shape) is None else<EOL>tensorshape_util.rank(self.batch_shape))<EOL>sample_ndims = x_ndims - batch_ndims - event_ndims<EOL>if isinstance(sample_ndims, int):<EOL><INDENT>static_sample_shape = x.shape[:sample_ndims]<EOL><DEDENT>else:<EOL><INDENT>static_sample_shape = tf.TensorShape(None)<EOL><DEDENT>if tensorshape_util.is_fully_defined(static_sample_shape):<EOL><INDENT>sample_shape = np.int32(static_sample_shape)<EOL><DEDENT>else:<EOL><INDENT>sample_shape = tf.shape(input=x)[:sample_ndims]<EOL><DEDENT>return sample_shape, static_sample_shape<EOL>

Computes graph and static `sample_shape`.


def _call_reshape_input_output(self, fn, x, extra_kwargs=None):

<EOL>with tf.control_dependencies(self._runtime_assertions +<EOL>self._validate_sample_arg(x)):<EOL><INDENT>sample_shape, static_sample_shape = self._sample_shape(x)<EOL>old_shape = tf.concat(<EOL>[<EOL>sample_shape,<EOL>self.distribution.batch_shape_tensor(),<EOL>self.event_shape_tensor(),<EOL>],<EOL>axis=<NUM_LIT:0>)<EOL>x_reshape = tf.reshape(x, old_shape)<EOL>result = fn(x_reshape, **extra_kwargs) if extra_kwargs else fn(x_reshape)<EOL>new_shape = tf.concat(<EOL>[<EOL>sample_shape,<EOL>self._batch_shape_unexpanded,<EOL>], axis=<NUM_LIT:0>)<EOL>result = tf.reshape(result, new_shape)<EOL>if (tensorshape_util.rank(static_sample_shape) is not None and<EOL>tensorshape_util.rank(self.batch_shape) is not None):<EOL><INDENT>new_shape = tensorshape_util.concatenate(static_sample_shape,<EOL>self.batch_shape)<EOL>tensorshape_util.set_shape(result, new_shape)<EOL><DEDENT>return result<EOL><DEDENT>

Calls `fn`, appropriately reshaping its input `x` and output.


<EOL>with tf.control_dependencies(self._runtime_assertions):<EOL><INDENT>if event_shape_list is None:<EOL><INDENT>event_shape_list = [self._event_shape_tensor()]<EOL><DEDENT>if static_event_shape_list is None:<EOL><INDENT>static_event_shape_list = [self.event_shape]<EOL><DEDENT>new_shape = tf.concat(<EOL>[self._batch_shape_unexpanded] + event_shape_list, axis=<NUM_LIT:0>)<EOL>result = tf.reshape(fn(**extra_kwargs) if extra_kwargs else fn(),<EOL>new_shape)<EOL>if (tensorshape_util.rank(self.batch_shape) is not None and<EOL>tensorshape_util.rank(self.event_shape) is not None):<EOL><INDENT>event_shape = tf.TensorShape([])<EOL>for rss in static_event_shape_list:<EOL><INDENT>event_shape = tensorshape_util.concatenate(event_shape, rss)<EOL><DEDENT>static_shape = tensorshape_util.concatenate(<EOL>self.batch_shape, event_shape)<EOL>tensorshape_util.set_shape(result, static_shape)<EOL><DEDENT>return result<EOL><DEDENT>

Calls `fn` and appropriately reshapes its output.


with tf.name_scope("<STR_LIT>"):<EOL><INDENT>x_ndims = (<EOL>tf.rank(x) if tensorshape_util.rank(x.shape) is None else<EOL>tensorshape_util.rank(x.shape))<EOL>event_ndims = (<EOL>tf.size(input=self.event_shape_tensor())<EOL>if tensorshape_util.rank(self.event_shape) is None else<EOL>tensorshape_util.rank(self.event_shape))<EOL>batch_ndims = (<EOL>tf.size(input=self._batch_shape_unexpanded)<EOL>if tensorshape_util.rank(self.batch_shape) is None else<EOL>tensorshape_util.rank(self.batch_shape))<EOL>expected_batch_event_ndims = batch_ndims + event_ndims<EOL>if (isinstance(x_ndims, int) and<EOL>isinstance(expected_batch_event_ndims, int)):<EOL><INDENT>if x_ndims < expected_batch_event_ndims:<EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>expected_batch_event_ndims, x_ndims))<EOL><DEDENT>ndims_assertion = []<EOL><DEDENT>elif self.validate_args:<EOL><INDENT>ndims_assertion = [<EOL>assert_util.assert_greater_equal(<EOL>x_ndims,<EOL>expected_batch_event_ndims,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>"),<EOL>name="<STR_LIT>"),<EOL>]<EOL><DEDENT>if (tensorshape_util.is_fully_defined(self.batch_shape) and<EOL>tensorshape_util.is_fully_defined(self.event_shape)):<EOL><INDENT>expected_batch_event_shape = np.int32(<EOL>tensorshape_util.concatenate(self.batch_shape, self.event_shape))<EOL><DEDENT>else:<EOL><INDENT>expected_batch_event_shape = tf.concat(<EOL>[<EOL>self.batch_shape_tensor(),<EOL>self.event_shape_tensor(),<EOL>], axis=<NUM_LIT:0>)<EOL><DEDENT>sample_ndims = x_ndims - expected_batch_event_ndims<EOL>if isinstance(sample_ndims, int):<EOL><INDENT>sample_ndims = max(sample_ndims, <NUM_LIT:0>)<EOL><DEDENT>if (isinstance(sample_ndims, int) and<EOL>tensorshape_util.is_fully_defined(x.shape[sample_ndims:])):<EOL><INDENT>actual_batch_event_shape = np.int32(x.shape[sample_ndims:])<EOL><DEDENT>else:<EOL><INDENT>sample_ndims = tf.maximum(sample_ndims, <NUM_LIT:0>)<EOL>actual_batch_event_shape = tf.shape(input=x)[sample_ndims:]<EOL><DEDENT>if (isinstance(expected_batch_event_shape, np.ndarray) and<EOL>isinstance(actual_batch_event_shape, np.ndarray)):<EOL><INDENT>if any(expected_batch_event_shape != actual_batch_event_shape):<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>expected_batch_event_shape,<EOL>actual_batch_event_shape))<EOL><DEDENT>runtime_assertions = ndims_assertion<EOL><DEDENT>elif self.validate_args:<EOL><INDENT>with tf.control_dependencies(ndims_assertion):<EOL><INDENT>shape_assertion = assert_util.assert_equal(<EOL>expected_batch_event_shape,<EOL>actual_batch_event_shape,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>"),<EOL>name="<STR_LIT>")<EOL><DEDENT>runtime_assertions = [shape_assertion]<EOL><DEDENT>else:<EOL><INDENT>runtime_assertions = []<EOL><DEDENT>return runtime_assertions<EOL><DEDENT>

Helper which validates sample arg, e.g., input to `log_prob`.


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

Initialize a batch of Zipf distributions.

        Args:
          power: `Float` like `Tensor` representing the power parameter. Must be
            strictly greater than `1`.
          dtype: The `dtype` of `Tensor` returned by `sample`.
            Default value: `tf.int32`.
          interpolate_nondiscrete: Python `bool`. When `False`, `log_prob` returns
            `-inf` (and `prob` returns `0`) for non-integer inputs. When `True`,
            `log_prob` evaluates the continuous function `-power log(k) -
            log(zeta(power))` , which matches the Zipf pmf at integer arguments `k`
            (note that this function is not itself a normalized probability
            log-density).
            Default value: `True`.
          sample_maximum_iterations: Maximum number of iterations of allowable
            iterations in `sample`. When `validate_args=True`, samples which fail to
            reach convergence (subject to this cap) are masked out with
            `self.dtype.min` or `nan` depending on `self.dtype.is_integer`.
            Default value: `100`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or more
            of the statistic's batch members are undefined.
            Default value: `False`.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: `'Zipf'`.

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


return self._power<EOL>


return self._interpolate_nondiscrete<EOL>

Interpolate (log) probs on non-integer inputs.


return self._sample_maximum_iterations<EOL>

Maximum number of allowable iterations in `sample`.


x = tf.cast(x, self.power.dtype)<EOL>t = self.power - <NUM_LIT:1.><EOL>return tf.exp((-t) * tf.math.log1p(x) - tf.math.log(t))<EOL>

Integral of the `hat` function, used for sampling.

        We choose a `hat` function, h(x) = x^(-power), which is a continuous
        (unnormalized) density touching each positive integer at the (unnormalized)
        pmf. This function implements `hat` integral: H(x) = int_x^inf h(t) dt;
        which is needed for sampling purposes.

        Arguments:
          x: A Tensor of points x at which to evaluate H(x).

        Returns:
          A Tensor containing evaluation H(x) at x.


x = tf.cast(x, self.power.dtype)<EOL>t = self.power - <NUM_LIT:1.><EOL>return tf.math.expm1(-(tf.math.log(t) + tf.math.log(x)) / t)<EOL>


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

Construct Vector Laplace distribution on `R^k`.

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

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

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

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

        where:

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

        Additional leading dimensions (if any) will index batches.

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

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

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


parameters = dict(locals())<EOL>with tf.name_scope(<EOL>name or '<STR_LIT>') as name:<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[kernel,<EOL>index_points,<EOL>inducing_index_points,<EOL>variational_inducing_observations_loc,<EOL>variational_inducing_observations_scale,<EOL>observation_noise_variance,<EOL>predictive_noise_variance,<EOL>jitter], tf.float32)<EOL>index_points = tf.convert_to_tensor(<EOL>value=index_points, dtype=dtype, name='<STR_LIT>')<EOL>inducing_index_points = tf.convert_to_tensor(<EOL>value=inducing_index_points, dtype=dtype,<EOL>name='<STR_LIT>')<EOL>variational_inducing_observations_loc = tf.convert_to_tensor(<EOL>value=variational_inducing_observations_loc, dtype=dtype,<EOL>name='<STR_LIT>')<EOL>variational_inducing_observations_scale = tf.convert_to_tensor(<EOL>value=variational_inducing_observations_scale, dtype=dtype,<EOL>name='<STR_LIT>')<EOL>observation_noise_variance = tf.convert_to_tensor(<EOL>value=observation_noise_variance,<EOL>dtype=dtype,<EOL>name='<STR_LIT>')<EOL>if predictive_noise_variance is None:<EOL><INDENT>predictive_noise_variance = observation_noise_variance<EOL><DEDENT>else:<EOL><INDENT>predictive_noise_variance = tf.convert_to_tensor(<EOL>value=predictive_noise_variance, dtype=dtype,<EOL>name='<STR_LIT>')<EOL><DEDENT>jitter = tf.convert_to_tensor(<EOL>value=jitter, dtype=dtype, name='<STR_LIT>')<EOL>self._kernel = kernel<EOL>self._index_points = index_points<EOL>self._inducing_index_points = inducing_index_points<EOL>self._variational_inducing_observations_posterior = (<EOL>mvn_linear_operator.MultivariateNormalLinearOperator(<EOL>loc=variational_inducing_observations_loc,<EOL>scale=tf.linalg.LinearOperatorFullMatrix(<EOL>variational_inducing_observations_scale),<EOL>name='<STR_LIT>'))<EOL>if mean_fn is None:<EOL><INDENT>mean_fn = lambda x: tf.zeros([<NUM_LIT:1>], dtype=dtype)<EOL><DEDENT>else:<EOL><INDENT>if not callable(mean_fn):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT><DEDENT>self._mean_fn = mean_fn<EOL>self._observation_noise_variance = observation_noise_variance<EOL>self._predictive_noise_variance = predictive_noise_variance<EOL>self._jitter = jitter<EOL>with tf.name_scope('<STR_LIT>'):<EOL><INDENT>kzz = _add_diagonal_shift(<EOL>kernel.matrix(inducing_index_points, inducing_index_points),<EOL>jitter)<EOL>self._chol_kzz = tf.linalg.cholesky(kzz)<EOL>self._kzz_inv_varloc = _solve_cholesky_factored_system_vec(<EOL>self._chol_kzz,<EOL>(variational_inducing_observations_loc -<EOL>mean_fn(inducing_index_points)),<EOL>name='<STR_LIT>')<EOL>loc, scale = self._compute_posterior_predictive_params()<EOL>super(VariationalGaussianProcess, self).__init__(<EOL>loc=loc,<EOL>scale=scale,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL>self._parameters = parameters<EOL>self._graph_parents = [<EOL>index_points,<EOL>inducing_index_points,<EOL>variational_inducing_observations_loc,<EOL>variational_inducing_observations_scale,<EOL>observation_noise_variance,<EOL>jitter]<EOL><DEDENT><DEDENT>

Instantiate a VariationalGaussianProcess Distribution.

        Args:
          kernel: `PositiveSemidefiniteKernel`-like instance representing the
            GP's covariance function.
          index_points: `float` `Tensor` representing finite (batch of) vector(s) of
            points in the index set over which the VGP is defined. Shape has the
            form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
            dimensions and must equal `kernel.feature_ndims` and `e1` is the number
            (size) of index points in each batch (we denote it `e1` to distinguish
            it from the numer of inducing index points, denoted `e2` below).
            Ultimately the VariationalGaussianProcess distribution corresponds to an
            `e1`-dimensional multivariate normal. The batch shape must be
            broadcastable with `kernel.batch_shape`, the batch shape of
            `inducing_index_points`, and any batch dims yielded by `mean_fn`.
          inducing_index_points: `float` `Tensor` of locations of inducing points in
            the index set. Shape has the form `[b1, ..., bB, e2, f1, ..., fF]`, just
            like `index_points`. The batch shape components needn't be identical to
            those of `index_points`, but must be broadcast compatible with them.
          variational_inducing_observations_loc: `float` `Tensor`; the mean of the
            (full-rank Gaussian) variational posterior over function values at the
            inducing points, conditional on observed data. Shape has the form `[b1,
            ..., bB, e2]`, where `b1, ..., bB` is broadcast compatible with other
            parameters' batch shapes, and `e2` is the number of inducing points.
          variational_inducing_observations_scale: `float` `Tensor`; the scale
            matrix of the (full-rank Gaussian) variational posterior over function
            values at the inducing points, conditional on observed data. Shape has
            the form `[b1, ..., bB, e2, e2]`, where `b1, ..., bB` is broadcast
            compatible with other parameters and `e2` is the number of inducing
            points.
          mean_fn: Python `callable` that acts on index points to produce a (batch
            of) vector(s) of mean values at those index points. Takes a `Tensor` of
            shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
            (broadcastable with) `[b1, ..., bB]`. Default value: `None` implies
            constant zero function.
          observation_noise_variance: `float` `Tensor` representing the variance
            of the noise in the Normal likelihood distribution of the model. May be
            batched, in which case the batch shape must be broadcastable with the
            shapes of all other batched parameters (`kernel.batch_shape`,
            `index_points`, etc.).
            Default value: `0.`
          predictive_noise_variance: `float` `Tensor` representing additional
            variance in the posterior predictive model. If `None`, we simply re-use
            `observation_noise_variance` for the posterior predictive noise. If set
            explicitly, however, we use the given value. This allows us, for
            example, to omit predictive noise variance (by setting this to zero) to
            obtain noiseless posterior predictions of function values, conditioned
            on noisy observations.
          jitter: `float` scalar `Tensor` added to the diagonal of the covariance
            matrix to ensure positive definiteness of the covariance matrix.
            Default value: `1e-6`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
            Default value: `False`.
          allow_nan_stats: Python `bool`, default `True`. When `True`,
            statistics (e.g., mean, mode, variance) use the value "`NaN`" to
            indicate the result is undefined. When `False`, an exception is raised
            if one or more of the statistic's batch members are undefined.
            Default value: `False`.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: "VariationalGaussianProcess".

        Raises:
          ValueError: if `mean_fn` is not `None` and is not callable.


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

Variational loss for the VGP.

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

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

        #### References

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


with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[inducing_index_points,<EOL>observation_index_points,<EOL>observations,<EOL>observation_noise_variance,<EOL>jitter], tf.float32)<EOL>inducing_index_points = tf.convert_to_tensor(<EOL>value=inducing_index_points,<EOL>dtype=dtype, name='<STR_LIT>')<EOL>observation_index_points = tf.convert_to_tensor(<EOL>value=observation_index_points, dtype=dtype,<EOL>name='<STR_LIT>')<EOL>observations = tf.convert_to_tensor(<EOL>value=observations, dtype=dtype, name='<STR_LIT>')<EOL>observation_noise_variance = tf.convert_to_tensor(<EOL>value=observation_noise_variance,<EOL>dtype=dtype,<EOL>name='<STR_LIT>')<EOL>jitter = tf.convert_to_tensor(<EOL>value=jitter, dtype=dtype, name='<STR_LIT>')<EOL>if mean_fn is None:<EOL><INDENT>mean_fn = lambda x: tf.zeros([<NUM_LIT:1>], dtype=dtype)<EOL><DEDENT>else:<EOL><INDENT>if not callable(mean_fn):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT><DEDENT>kzz = kernel.matrix(inducing_index_points, inducing_index_points)<EOL>kzx = kernel.matrix(inducing_index_points, observation_index_points)<EOL>noise_var_inv = tf.math.reciprocal(observation_noise_variance)<EOL>sigma_inv = _add_diagonal_shift(<EOL>kzz + noise_var_inv * tf.matmul(kzx, kzx, adjoint_b=True),<EOL>jitter)<EOL>chol_sigma_inv = tf.linalg.cholesky(sigma_inv)<EOL>kzx_lin_op = tf.linalg.LinearOperatorFullMatrix(kzx)<EOL>kzx_obs = kzx_lin_op.matvec(<EOL>observations - mean_fn(observation_index_points))<EOL>kzz_lin_op = tf.linalg.LinearOperatorFullMatrix(kzz)<EOL>loc = (mean_fn(inducing_index_points) +<EOL>noise_var_inv * kzz_lin_op.matvec(<EOL>_solve_cholesky_factored_system_vec(chol_sigma_inv, kzx_obs)))<EOL>chol_sigma_inv_lin_op = tf.linalg.LinearOperatorLowerTriangular(<EOL>chol_sigma_inv)<EOL>scale = chol_sigma_inv_lin_op.solve(kzz)<EOL>return loc, scale<EOL><DEDENT>

Model selection for optimal variational hyperparameters.

        Given the full training set (parameterized by `observations` and
        `observation_index_points`), compute the optimal variational
        location and scale for the VGP. This is based of the method suggested
        in [Titsias, 2009][1].

        Args:
          kernel: `PositiveSemidefiniteKernel`-like instance representing the
            GP's covariance function.
          inducing_index_points: `float` `Tensor` of locations of inducing points in
            the index set. Shape has the form `[b1, ..., bB, e2, f1, ..., fF]`, just
            like `observation_index_points`. The batch shape components needn't be
            identical to those of `observation_index_points`, but must be broadcast
            compatible with them.
          observation_index_points: `float` `Tensor` representing finite (batch of)
            vector(s) of points where observations are defined. Shape has the
            form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
            dimensions and must equal `kernel.feature_ndims` and `e1` is the number
            (size) of index points in each batch (we denote it `e1` to distinguish
            it from the numer of inducing index points, denoted `e2` below).
          observations: `float` `Tensor` representing collection, or batch of
            collections, of observations corresponding to
            `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
            must be brodcastable with the batch and example shapes of
            `observation_index_points`. The batch shape `[b1, ..., bB]` must be
            broadcastable with the shapes of all other batched parameters
            (`kernel.batch_shape`, `observation_index_points`, etc.).
          observation_noise_variance: `float` `Tensor` representing the variance
            of the noise in the Normal likelihood distribution of the model. May be
            batched, in which case the batch shape must be broadcastable with the
            shapes of all other batched parameters (`kernel.batch_shape`,
            `index_points`, etc.).
            Default value: `0.`
          mean_fn: Python `callable` that acts on index points to produce a (batch
            of) vector(s) of mean values at those index points. Takes a `Tensor` of
            shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
            (broadcastable with) `[b1, ..., bB]`. Default value: `None` implies
            constant zero function.
          jitter: `float` scalar `Tensor` added to the diagonal of the covariance
            matrix to ensure positive definiteness of the covariance matrix.
            Default value: `1e-6`.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: "optimal_variational_posterior".
        Returns:
          loc, scale: Tuple representing the variational location and scale.
        Raises:
          ValueError: if `mean_fn` is not `None` and is not callable.

        #### References

        [1]: Titsias, M. "Variational Model Selection for Sparse Gaussian Process
             Regression", 2009.
             http://proceedings.mlr.press/v5/titsias09a/titsias09a.pdf


parameters = dict(locals())<EOL>if scale is None:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if not dtype_util.is_floating(scale.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>with tf.name_scope(name):<EOL><INDENT>loc = loc if loc is None else tf.convert_to_tensor(<EOL>value=loc, name="<STR_LIT>", dtype=scale.dtype)<EOL>batch_shape, event_shape = distribution_util.shapes_from_loc_and_scale(<EOL>loc, scale)<EOL>super(VectorLaplaceLinearOperator, self).__init__(<EOL>distribution=laplace.Laplace(<EOL>loc=tf.zeros([], dtype=scale.dtype),<EOL>scale=tf.ones([], dtype=scale.dtype)),<EOL>bijector=affine_linear_operator_bijector.AffineLinearOperator(<EOL>shift=loc, scale=scale, validate_args=validate_args),<EOL>batch_shape=batch_shape,<EOL>event_shape=event_shape,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>self._parameters = parameters<EOL><DEDENT>

Construct Vector Laplace distribution on `R^k`.

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

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

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

        Additional leading dimensions (if any) will index batches.

        Args:
          loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
            implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
            `b >= 0` and `k` is the event size.
          scale: Instance of `LinearOperator` with same `dtype` as `loc` and shape
            `[B1, ..., Bb, k, k]`.
          validate_args: Python `bool`, default `False`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are
            invalid, correct behavior is not guaranteed.
          allow_nan_stats: Python `bool`, default `True`. If `False`, raise an
            exception if a statistic (e.g. mean/mode/etc...) is undefined for any
            batch member If `True`, batch members with valid parameters leading to
            undefined statistics will return NaN for this statistic.
          name: The name to give Ops created by the initializer.

        Raises:
          ValueError: if `scale` is unspecified.
          TypeError: if not `scale.dtype.is_floating`


return self.bijector.shift<EOL>

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


return self.bijector.scale<EOL>

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


<EOL>parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([total_count, concentration], tf.float32)<EOL>self._total_count = tf.convert_to_tensor(<EOL>value=total_count, name="<STR_LIT>", dtype=dtype)<EOL>if validate_args:<EOL><INDENT>self._total_count = (<EOL>distribution_util.embed_check_nonnegative_integer_form(<EOL>self._total_count))<EOL><DEDENT>self._concentration = self._maybe_assert_valid_concentration(<EOL>tf.convert_to_tensor(<EOL>value=concentration, name="<STR_LIT>", dtype=dtype),<EOL>validate_args)<EOL>self._total_concentration = tf.reduce_sum(<EOL>input_tensor=self._concentration, axis=-<NUM_LIT:1>)<EOL>self._broadcasted_concentration = tf.ones_like(<EOL>self._total_count[..., tf.newaxis]) * self._concentration<EOL><DEDENT>super(DirichletMultinomial, self).__init__(<EOL>dtype=dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>parameters=parameters,<EOL>graph_parents=[self._total_count, self._concentration],<EOL>name=name)<EOL>

Initialize a batch of DirichletMultinomial distributions.

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


return self._total_count<EOL>

Number of trials used to construct a sample.


return self._concentration<EOL>

Concentration parameter; expected prior counts for that coordinate.


return self._total_concentration<EOL>

Sum of last dim of concentration parameter.


<EOL>c0 = self.total_concentration[..., tf.newaxis]<EOL>return tf.sqrt((<NUM_LIT:1.> + c0 / self.total_count[..., tf.newaxis]) / (<NUM_LIT:1.> + c0))<EOL>

Helper to `_covariance` and `_variance` which computes a shared scale.


def _maybe_assert_valid_concentration(self, concentration, validate_args):

if not validate_args:<EOL><INDENT>return concentration<EOL><DEDENT>concentration = distribution_util.embed_check_categorical_event_shape(<EOL>concentration)<EOL>return distribution_util.with_dependencies([<EOL>assert_util.assert_positive(<EOL>concentration, message="<STR_LIT>"),<EOL>], concentration)<EOL>

Checks the validity of the concentration parameter.


if not self.validate_args:<EOL><INDENT>return counts<EOL><DEDENT>counts = distribution_util.embed_check_nonnegative_integer_form(counts)<EOL>return distribution_util.with_dependencies([<EOL>assert_util.assert_equal(<EOL>self.total_count,<EOL>tf.reduce_sum(input_tensor=counts, axis=-<NUM_LIT:1>),<EOL>message="<STR_LIT>"),<EOL>], counts)<EOL>

Check counts for proper shape, values, then return tensor version.


<EOL>if tf.executing_eagerly():<EOL><INDENT>tf.compat.v1.set_random_seed(seed)<EOL>return None<EOL><DEDENT>return seed<EOL>

Helper which uses graph seed if using TFE.


def _mixture_stddev_np(pi_vector, mu_vector, sigma_vector):

pi_vector = np.expand_dims(pi_vector, axis=<NUM_LIT:1>)<EOL>mean_wa = np.matmul(pi_vector, np.expand_dims(mu_vector, axis=<NUM_LIT:2>))<EOL>var_wa = np.matmul(pi_vector, np.expand_dims(sigma_vector**<NUM_LIT:2>, axis=<NUM_LIT:2>))<EOL>mid_term = np.matmul(pi_vector, np.expand_dims(mu_vector**<NUM_LIT:2>, axis=<NUM_LIT:2>))<EOL>mixture_variance = (<EOL>np.squeeze(var_wa) + np.squeeze(mid_term) - np.squeeze(mean_wa**<NUM_LIT:2>))<EOL>return np.sqrt(mixture_variance)<EOL>

Computes the standard deviation of a univariate mixture distribution.

    Acts upon `np.array`s (not `tf.Tensor`s).

    Args:
      pi_vector: A `np.array` of mixture weights. Shape `[batch, components]`.
      mu_vector: A `np.array` of means. Shape `[batch, components]`
      sigma_vector: A `np.array` of stddevs. Shape `[batch, components]`.

    Returns:
      A `np.array` containing the batch of standard deviations.


flat_xs = tuple(flat_xs)  <EOL>if not validate_args:<EOL><INDENT>return flat_xs<EOL><DEDENT>msg = '<STR_LIT>'<EOL>s = tuple(x.shape for x in flat_xs)<EOL>if all(tensorshape_util.is_fully_defined(s_) for s_ in s):<EOL><INDENT>if not all(a == b for a, b in zip(s[<NUM_LIT:1>:], s[:-<NUM_LIT:1>])):<EOL><INDENT>raise ValueError(msg)<EOL><DEDENT>return flat_xs<EOL><DEDENT>assertions = [assert_util.assert_equal(a, b, message=msg)<EOL>for a, b in zip(s[<NUM_LIT:1>:], s[:-<NUM_LIT:1>])]<EOL>with tf.control_dependencies(assertions):<EOL><INDENT>return tuple(tf.identity(x) for x in flat_xs)<EOL><DEDENT>


if self._most_recently_built_distributions is None:<EOL><INDENT>return super(JointDistribution, self).dtype<EOL><DEDENT>return self._unflatten(<EOL>None if d is None else d.dtype<EOL>for d in self._most_recently_built_distributions)<EOL>

The `DType` of `Tensor`s handled by this `Distribution`.


if self._most_recently_built_distributions is None:<EOL><INDENT>return None<EOL><DEDENT>return self._unflatten(<EOL>None if d is None else d.reparameterization_type<EOL>for d in self._most_recently_built_distributions)<EOL>

Describes how samples from the distribution are reparameterized.

        Currently this is one of the static instances
        `tfd.FULLY_REPARAMETERIZED` or `tfd.NOT_REPARAMETERIZED`.

        Returns:
          reparameterization_type: `ReparameterizationType` of each distribution in
            `model`.


<EOL>if self._most_recently_built_distributions is None:<EOL><INDENT>return None<EOL><DEDENT>return self._unflatten(<EOL>tf.TensorShape(None) if d is None else d.batch_shape<EOL>for d in self._most_recently_built_distributions)<EOL>

Shape of a single sample from a single event index as a `TensorShape`.

        May be partially defined or unknown.

        The batch dimensions are indexes into independent, non-identical
        parameterizations of this distribution.

        Returns:
          batch_shape: `tuple` of `TensorShape`s representing the `batch_shape` for
            each distribution in `model`.


with self._name_scope(name):<EOL><INDENT>return self._unflatten(self._map_attr_over_dists('<STR_LIT>')[<NUM_LIT:1>])<EOL><DEDENT>

Shape of a single sample from a single event index as a 1-D `Tensor`.

        The batch dimensions are indexes into independent, non-identical
        parameterizations of this distribution.

        Args:
          name: name to give to the op

        Returns:
          batch_shape: `Tensor` representing batch shape of each distribution in
            `model`.


<EOL>if self._most_recently_built_distributions is None:<EOL><INDENT>return None<EOL><DEDENT>return self._unflatten(<EOL>tf.TensorShape(None) if d is None else d.event_shape<EOL>for d in self._most_recently_built_distributions)<EOL>

Shape of a single sample from a single batch as a `TensorShape`.

        May be partially defined or unknown.

        Returns:
          event_shape: `tuple` of `TensorShape`s representing the `event_shape` for
            each distribution in `model`.


Shape of a single sample from a single batch as a 1-D int32 `Tensor`.

        Args:
          name: name to give to the op

        Returns:
          event_shape: `tuple` of `Tensor`s representing the `event_shape` for each
            distribution in `model`.


def sample_distributions(self, sample_shape=(), seed=None, value=None,<EOL>name='<STR_LIT>'):

with self._name_scope(name):<EOL><INDENT>ds, xs = self._call_flat_sample_distributions(sample_shape, seed, value)<EOL>return self._unflatten(ds), self._unflatten(xs)<EOL><DEDENT>

Generate samples and the (random) distributions.

        Note that a call to `sample()` without arguments will generate a single
        sample.

        Args:
          sample_shape: 0D or 1D `int32` `Tensor`. Shape of the generated samples.
          seed: Python integer seed for generating random numbers.
          value: `list` of `Tensor`s in `distribution_fn` order to use to
            parameterize other ("downstream") distribution makers.
            Default value: `None` (i.e., draw a sample from each distribution).
          name: name prepended to ops created by this function.
            Default value: `"sample_distributions"`.

        Returns:
          distributions: a `tuple` of `Distribution` instances for each of
            `distribution_fn`.
          samples: a `tuple` of `Tensor`s with prepended dimensions `sample_shape`
            for each of `distribution_fn`.


def log_prob_parts(self, value, name='<STR_LIT>'):

with self._name_scope(name):<EOL><INDENT>_, xs = self._map_measure_over_dists('<STR_LIT>', value)<EOL>return self._unflatten(xs)<EOL><DEDENT>

Log probability density/mass function.

        Args:
          value: `list` of `Tensor`s in `distribution_fn` order for which we compute
            the `log_prob_parts` and to parameterize other ("downstream")
            distributions.
          name: name prepended to ops created by this function.
            Default value: `"log_prob_parts"`.

        Returns:
          log_prob_parts: a `tuple` of `Tensor`s representing the `log_prob` for
            each `distribution_fn` evaluated at each corresponding `value`.


def prob_parts(self, value, name='<STR_LIT>'):

Log probability density/mass function.

        Args:
          value: `list` of `Tensor`s in `distribution_fn` order for which we compute
            the `prob_parts` and to parameterize other ("downstream") distributions.
          name: name prepended to ops created by this function.
            Default value: `"prob_parts"`.

        Returns:
          prob_parts: a `tuple` of `Tensor`s representing the `prob` for
            each `distribution_fn` evaluated at each corresponding `value`.


Indicates that `event_shape == []`.

        Args:
          name: Python `str` prepended to names of ops created by this function.

        Returns:
          is_scalar_event: `bool` scalar `Tensor` for each distribution in `model`.


Indicates that `batch_shape == []`.

        Args:
          name: Python `str` prepended to names of ops created by this function.

        Returns:
          is_scalar_batch: `bool` scalar `Tensor` for each distribution in `model`.


parameters = dict(locals())<EOL>self._input_output_cholesky = input_output_cholesky<EOL>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>if not dtype_util.is_floating(scale_operator.dtype):<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>" %<EOL>scale_operator.dtype)<EOL><DEDENT>if not scale_operator.is_square:<EOL><INDENT>print(scale_operator.to_dense().eval())<EOL>raise ValueError("<STR_LIT>")<EOL><DEDENT>self._scale_operator = scale_operator<EOL>self._df = tf.convert_to_tensor(<EOL>value=df, dtype=scale_operator.dtype, name="<STR_LIT>")<EOL>dtype_util.assert_same_float_dtype([self._df, self._scale_operator])<EOL>if tf.compat.dimension_value(self._scale_operator.shape[-<NUM_LIT:1>]) is None:<EOL><INDENT>self._dimension = tf.cast(<EOL>self._scale_operator.domain_dimension_tensor(),<EOL>dtype=self._scale_operator.dtype,<EOL>name="<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>self._dimension = tf.convert_to_tensor(<EOL>value=tf.compat.dimension_value(self._scale_operator.shape[-<NUM_LIT:1>]),<EOL>dtype=self._scale_operator.dtype,<EOL>name="<STR_LIT>")<EOL><DEDENT>df_val = tf.get_static_value(self._df)<EOL>dim_val = tf.get_static_value(self._dimension)<EOL>if df_val is not None and dim_val is not None:<EOL><INDENT>df_val = np.asarray(df_val)<EOL>if not df_val.shape:<EOL><INDENT>df_val = [df_val]<EOL><DEDENT>if np.any(df_val < dim_val):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>"<EOL>% (df_val, dim_val))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions = assert_util.assert_less_equal(<EOL>self._dimension,<EOL>self._df,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>" % (self._dimension, self._df)))<EOL>self._df = distribution_util.with_dependencies(<EOL>[assertions], self._df)<EOL><DEDENT><DEDENT><DEDENT>super(_WishartLinearOperator, self).__init__(<EOL>dtype=self._scale_operator.dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>parameters=parameters,<EOL>graph_parents=(<EOL>[self._df, self._dimension] + self._scale_operator.graph_parents),<EOL>name=name)<EOL>

Construct Wishart distributions.

        Args:
          df: `float` or `double` tensor, the degrees of freedom of the
            distribution(s). `df` must be greater than or equal to `k`.
          scale_operator: `float` or `double` instance of `LinearOperator`.
          input_output_cholesky: Python `bool`. If `True`, functions whose input or
            output have the semantics of samples assume inputs are in Cholesky form
            and return outputs in Cholesky form. In particular, if this flag is
            `True`, input to `log_prob` is presumed of Cholesky form and output from
            `sample`, `mean`, and `mode` are of Cholesky form.  Setting this
            argument to `True` is purely a computational optimization and does not
            change the underlying distribution; for instance, `mean` returns the
            Cholesky of the mean, not the mean of Cholesky factors. The `variance`
            and `stddev` methods are unaffected by this flag.
            Default value: `False` (i.e., input/output does not have Cholesky
            semantics).
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          TypeError: if scale is not floating-type
          TypeError: if scale.dtype != df.dtype
          ValueError: if df < k, where scale operator event shape is
            `(k, k)`


return self._df<EOL>

Wishart distribution degree(s) of freedom.


if self._input_output_cholesky:<EOL><INDENT>return self.scale_operator.to_dense()<EOL><DEDENT>else:<EOL><INDENT>return self._square_scale_operator()<EOL><DEDENT>


return self._scale_operator<EOL>

Wishart distribution scale matrix as an Linear Operator.


return self._input_output_cholesky<EOL>

Boolean indicating if `Tensor` input/outputs are Cholesky factorized.


return self._dimension<EOL>

Dimension of underlying vector space. The `p` in `R^(p*p)`.


with self._name_scope(name):<EOL><INDENT>return (self._multi_digamma(<NUM_LIT:0.5> * self.df, self.dimension) +<EOL>self.dimension * math.log(<NUM_LIT>) +<EOL><NUM_LIT:2> * self.scale_operator.log_abs_determinant())<EOL><DEDENT>

Computes E[log(det(X))] under this Wishart distribution.


with self._name_scope(name):<EOL><INDENT>return (self.df * self.scale_operator.log_abs_determinant() +<EOL><NUM_LIT:0.5> * self.df * self.dimension * math.log(<NUM_LIT>) +<EOL>self._multi_lgamma(<NUM_LIT:0.5> * self.df, self.dimension))<EOL><DEDENT>

Computes the log normalizing constant, log(Z).


def _multi_gamma_sequence(self, a, p, name="<STR_LIT>"):

with self._name_scope(name):<EOL><INDENT>seq = tf.linspace(<EOL>tf.constant(<NUM_LIT:0.>, dtype=self.dtype), <NUM_LIT:0.5> - <NUM_LIT:0.5> * p, tf.cast(<EOL>p, tf.int32))<EOL>return seq + tf.expand_dims(a, [-<NUM_LIT:1>])<EOL><DEDENT>

Creates sequence used in multivariate (di)gamma; shape = shape(a)+[p].


def _multi_lgamma(self, a, p, name="<STR_LIT>"):

with self._name_scope(name):<EOL><INDENT>seq = self._multi_gamma_sequence(a, p)<EOL>return (<NUM_LIT> * p * (p - <NUM_LIT:1.>) * math.log(math.pi) +<EOL>tf.reduce_sum(input_tensor=tf.math.lgamma(seq), axis=[-<NUM_LIT:1>]))<EOL><DEDENT>

Computes the log multivariate gamma function; log(Gamma_p(a)).


def _multi_digamma(self, a, p, name="<STR_LIT>"):

with self._name_scope(name):<EOL><INDENT>seq = self._multi_gamma_sequence(a, p)<EOL>return tf.reduce_sum(input_tensor=tf.math.digamma(seq), axis=[-<NUM_LIT:1>])<EOL><DEDENT>

Computes the multivariate digamma function; Psi_p(a).


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>if (scale is None) == (scale_tril is None):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>dtype = dtype_util.common_dtype([df, scale, scale_tril], tf.float32)<EOL>df = tf.convert_to_tensor(value=df, name="<STR_LIT>", dtype=dtype)<EOL>if scale is not None:<EOL><INDENT>scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>", dtype=dtype)<EOL>if validate_args:<EOL><INDENT>scale = distribution_util.assert_symmetric(scale)<EOL><DEDENT>scale_tril = tf.linalg.cholesky(scale)<EOL><DEDENT>else:  <EOL><INDENT>scale_tril = tf.convert_to_tensor(<EOL>value=scale_tril, name="<STR_LIT>", dtype=dtype)<EOL>if validate_args:<EOL><INDENT>scale_tril = distribution_util.with_dependencies([<EOL>assert_util.assert_positive(<EOL>tf.linalg.diag_part(scale_tril),<EOL>message="<STR_LIT>"),<EOL>assert_util.assert_equal(<EOL>tf.shape(input=scale_tril)[-<NUM_LIT:1>],<EOL>tf.shape(input=scale_tril)[-<NUM_LIT:2>],<EOL>message="<STR_LIT>")<EOL>], scale_tril)<EOL><DEDENT><DEDENT><DEDENT>super(Wishart, self).__init__(<EOL>df=df,<EOL>scale_operator=tf.linalg.LinearOperatorLowerTriangular(<EOL>tril=scale_tril,<EOL>is_non_singular=True,<EOL>is_positive_definite=True,<EOL>is_square=True),<EOL>input_output_cholesky=input_output_cholesky,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL><DEDENT>self._parameters = parameters<EOL>

Construct Wishart distributions.

        Args:
          df: `float` or `double` `Tensor`. Degrees of freedom, must be greater than
            or equal to dimension of the scale matrix.
          scale: `float` or `double` `Tensor`. The symmetric positive definite
            scale matrix of the distribution. Exactly one of `scale` and
            'scale_tril` must be passed.
          scale_tril: `float` or `double` `Tensor`. The Cholesky factorization
            of the symmetric positive definite scale matrix of the distribution.
            Exactly one of `scale` and 'scale_tril` must be passed.
          input_output_cholesky: Python `bool`. If `True`, functions whose input or
            output have the semantics of samples assume inputs are in Cholesky form
            and return outputs in Cholesky form. In particular, if this flag is
            `True`, input to `log_prob` is presumed of Cholesky form and output from
            `sample`, `mean`, and `mode` are of Cholesky form.  Setting this
            argument to `True` is purely a computational optimization and does not
            change the underlying distribution; for instance, `mean` returns the
            Cholesky of the mean, not the mean of Cholesky factors. The `variance`
            and `stddev` methods are unaffected by this flag.
            Default value: `False` (i.e., input/output does not have Cholesky
            semantics).
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.
        Raises:
          ValueError: if zero or both of 'scale' and 'scale_tril' are passed in.


value = tf.cast(value, self.dtype)<EOL>return tf.compat.v1.placeholder_with_default(<EOL>value, shape=value.shape if self.use_static_shape else None)<EOL>

Helper to create inputs with varied dtypes an static shapes.


@functools.wraps(f)<EOL>def _check_arg_and_apply_f(*args, **kwargs):<EOL><INDENT>dist = args[<NUM_LIT:0>]<EOL>x = args[<NUM_LIT:1>]<EOL>with tf.control_dependencies([<EOL>assert_util.assert_greater_equal(<EOL>x, dist.loc, message="<STR_LIT>")<EOL>] if dist.validate_args else []):<EOL><INDENT>return f(*args, **kwargs)<EOL><DEDENT><DEDENT>return _check_arg_and_apply_f<EOL>

Decorator function for argument bounds checking.

    This decorator is meant to be used with methods that require the first
    argument to be in the support of the distribution. If `validate_args` is
    `True`, the method is wrapped with an assertion that the first argument is
    greater than or equal to `loc`, since the support of the half-Cauchy
    distribution is given by `[loc, infinity)`.

    Args:
      f: method to be decorated.

    Returns:
      Returns a decorated method that, when `validate_args` attribute of the class
      is `True`, will assert that all elements in the first argument are within
      the support of the distribution before executing the original method.


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

Construct a half-Cauchy distribution with `loc` and `scale`.

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

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


return self._loc<EOL>

Distribution parameter for the location.


return self._scale<EOL>


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


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

Reconstruct input `x` from a its normalized version.


def _extend_support_with_default_value(self, x, f, default_value):

with tf.name_scope("<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, dtype=self.dtype, name="<STR_LIT:x>")<EOL>loc = self.loc + tf.zeros_like(self.scale) + tf.zeros_like(x)<EOL>x = x + tf.zeros_like(loc)<EOL>y = f(tf.where(x < loc, self._inv_z(<NUM_LIT:0.5>) + tf.zeros_like(x), x))<EOL>if default_value == <NUM_LIT:0.>:<EOL><INDENT>default_value = tf.zeros_like(y)<EOL><DEDENT>elif default_value == <NUM_LIT:1.>:<EOL><INDENT>default_value = tf.ones_like(y)<EOL><DEDENT>else:<EOL><INDENT>default_value = tf.fill(<EOL>dims=tf.shape(input=y),<EOL>value=dtype_util.as_numpy_dtype(self.dtype)(default_value))<EOL><DEDENT>return tf.where(x < loc, default_value, y)<EOL><DEDENT>

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

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

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


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

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

Calculate the batched KL divergence KL(a || b) with a, b OneHotCategorical.

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

    Returns:
      Batchwise KL(a || b)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>name=name, logits=logits, probs=probs, validate_args=validate_args,<EOL>multidimensional=True)<EOL>logits_shape_static = tensorshape_util.with_rank_at_least(<EOL>self._logits.shape, <NUM_LIT:1>)<EOL>if tensorshape_util.rank(logits_shape_static) is not None:<EOL><INDENT>self._batch_rank = tf.convert_to_tensor(<EOL>value=tensorshape_util.rank(logits_shape_static) - <NUM_LIT:1>,<EOL>dtype=tf.int32,<EOL>name="<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>self._batch_rank = tf.rank(self._logits) - <NUM_LIT:1><EOL><DEDENT><DEDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>self._event_size = tf.shape(input=self._logits)[-<NUM_LIT:1>]<EOL><DEDENT><DEDENT>super(OneHotCategorical, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._logits, self._probs],<EOL>name=name)<EOL>

Initialize OneHotCategorical distributions using class log-probabilities.

        Args:
          logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities of a
            set of Categorical distributions. The first `N - 1` dimensions index
            into a batch of independent distributions and the last dimension
            represents a vector of logits for each class. Only one of `logits` or
            `probs` should be passed in.
          probs: An N-D `Tensor`, `N >= 1`, representing the probabilities of a set
            of Categorical distributions. The first `N - 1` dimensions index into a
            batch of independent distributions and the last dimension represents a
            vector of probabilities for each class. Only one of `logits` or `probs`
            should be passed in.
          dtype: The type of the event samples (default: int32).
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.


return self._event_size<EOL>

Scalar `int32` tensor: the number of classes.


return self._logits<EOL>


return self._probs<EOL>


def _explore(v):<EOL><INDENT>if v.depth < <NUM_LIT:0>:<EOL><INDENT>v.depth = ((<NUM_LIT:1> + max([-<NUM_LIT:1>] + [_explore(annotated_graph[u])<EOL>for u in v.parents]))<EOL>if v.parents else <NUM_LIT:0>)<EOL><DEDENT>return v.depth<EOL><DEDENT>annotated_graph = {k: _Node(k, v) for k, v in g.items()}<EOL>for v in annotated_graph.values():<EOL><INDENT>_explore(v)<EOL><DEDENT>return annotated_graph<EOL>

Computes the number of edges on longest path from node to root.


def _explore(u):<EOL><INDENT>"""<STR_LIT>"""<EOL>if u.depth < <NUM_LIT:0>:<EOL><INDENT>return  <EOL><DEDENT>if not u.parents:<EOL><INDENT>result.append((u.name, u.parents))<EOL>u.depth = -<NUM_LIT:1>  <EOL>return<EOL><DEDENT>b = (u.name, [])<EOL>result.append(b)<EOL>u.depth = -<NUM_LIT:1>  <EOL>d = <NUM_LIT:0><EOL>for v in sorted((g.get(p) for p in u.parents), key=lambda v: v.depth):<EOL><INDENT>n0 = len(result)<EOL>_explore(v)<EOL>n1 = len(result)<EOL>b[<NUM_LIT:1>].extend(['<STR_LIT:_>']*d + [v.name])<EOL>d = n1 - n0 - <NUM_LIT:1><EOL><DEDENT><DEDENT>g = _depth(g)<EOL>result = []<EOL>for u in sorted(g.values(), key=lambda v: v.depth, reverse=True):<EOL><INDENT>_explore(u)<EOL><DEDENT>return tuple(reversed(result))<EOL>

Creates tuple of str tuple-str pairs representing resolved & sorted DAG.


def _make(dist_fn, args):<EOL><INDENT>if args is None:<EOL><INDENT>return lambda *_: dist_fn<EOL><DEDENT>if not args:<EOL><INDENT>return lambda *_: dist_fn()<EOL><DEDENT>def _fn(*xs):<EOL><INDENT>kwargs = dict(zip(args, reversed(xs[-len(args):])))<EOL>kwargs.pop('<STR_LIT:_>', None)<EOL>return dist_fn(**kwargs)<EOL><DEDENT>return _fn<EOL><DEDENT>named_makers = _convert_to_dict(named_makers)<EOL>g = {k: (None if distribution_util.is_distribution_instance(v)<EOL>else joint_distribution_sequential._get_required_args(v))  <EOL>for k, v in named_makers.items()}<EOL>g = _best_order(g)<EOL>dist_fn_name, dist_fn_args = zip(*g)<EOL>dist_fn_args = tuple(None if a is None else tuple(a) for a in dist_fn_args)<EOL>dist_fn_wrapped = tuple(_make(named_makers[name], parents)<EOL>for (name, parents) in g)<EOL>dist_fn = tuple(named_makers.get(n) for n in dist_fn_name)<EOL>return dist_fn, dist_fn_wrapped, dist_fn_args, dist_fn_name<EOL>

Creates lists of callables suitable for JDSeq.


if hasattr(x, '<STR_LIT>'):<EOL><INDENT>return True<EOL><DEDENT>return isinstance(x, collections.Mapping)<EOL>

Returns `True` if input is convertible to `dict`, `False` otherwise.


if hasattr(x, '<STR_LIT>'):<EOL><INDENT>return x._asdict()<EOL><DEDENT>return dict(x)<EOL>


super(JointDistributionNamed, self).__init__(<EOL>model, validate_args, name or '<STR_LIT>')<EOL>

Construct the `JointDistributionNamed` distribution.

        Args:
          model: Python `dict` or `namedtuple` of distribution-making functions each
            with required args corresponding only to other keys.
          validate_args: Python `bool`.  Whether to validate input with asserts.
            If `validate_args` is `False`, and the inputs are invalid,
            correct behavior is not guaranteed.
            Default value: `False`.
          name: The name for ops managed by the distribution.
            Default value: `None` (i.e., `"JointDistributionNamed"`).


if not _is_dict_like(model):<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>type(model).__name__))<EOL><DEDENT>[<EOL>self._dist_fn,<EOL>self._dist_fn_wrapped,<EOL>self._dist_fn_args,<EOL>self._dist_fn_name,  <EOL>] = _prob_chain_rule_flatten(model)<EOL>

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


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

Constructs inverse Gaussian distribution with `loc` and `concentration`.

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


return self._loc<EOL>


return self._concentration<EOL>


def _uniform_unit_norm(dimension, shape, dtype, seed):

<EOL>raw = normal.Normal(<EOL>loc=dtype_util.as_numpy_dtype(dtype)(<NUM_LIT:0>),<EOL>scale=dtype_util.as_numpy_dtype(dtype)(<NUM_LIT:1>)).sample(<EOL>tf.concat([shape, [dimension]], axis=<NUM_LIT:0>), seed=seed())<EOL>unit_norm = raw / tf.norm(tensor=raw, ord=<NUM_LIT:2>, axis=-<NUM_LIT:1>)[..., tf.newaxis]<EOL>return unit_norm<EOL>

Returns a batch of points chosen uniformly from the unit hypersphere.


<EOL>multiples = tf.concat([[n], tf.ones_like(tensor.shape)], axis=<NUM_LIT:0>)<EOL>return tf.tile(tf.expand_dims(tensor, axis=<NUM_LIT:0>), multiples)<EOL>

Replicate the input tensor n times along a new (major) dimension.


if dimension < <NUM_LIT:0>:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>')<EOL><DEDENT>parameters = dict(locals())<EOL>self._input_output_cholesky = input_output_cholesky<EOL>with tf.name_scope(name):<EOL><INDENT>concentration = tf.convert_to_tensor(<EOL>value=concentration,<EOL>name='<STR_LIT>',<EOL>dtype=dtype_util.common_dtype([concentration],<EOL>preferred_dtype=tf.float32))<EOL>with tf.control_dependencies([<EOL>assert_util.assert_non_negative(concentration - <NUM_LIT:1.>),<EOL>] if validate_args else []):<EOL><INDENT>self._dimension = dimension<EOL>self._concentration = tf.identity(concentration, name='<STR_LIT>')<EOL><DEDENT><DEDENT>super(LKJ, self).__init__(<EOL>dtype=self._concentration.dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>parameters=parameters,<EOL>graph_parents=[self._concentration],<EOL>name=name)<EOL>

Construct LKJ distributions.

        Args:
          dimension: Python `int`. The dimension of the correlation matrices
            to sample.
          concentration: `float` or `double` `Tensor`. The positive concentration
            parameter of the LKJ distributions. The pdf of a sample matrix `X` is
            proportional to `det(X) ** (concentration - 1)`.
          input_output_cholesky: Python `bool`. If `True`, functions whose input or
            output have the semantics of samples assume inputs are in Cholesky form
            and return outputs in Cholesky form. In particular, if this flag is
            `True`, input to `log_prob` is presumed of Cholesky form and output from
            `sample` is of Cholesky form.  Setting this argument to `True` is purely
            a computational optimization and does not change the underlying
            distribution. Additionally, validation checks which are only defined on
            the multiplied-out form are omitted, even if `validate_args` is `True`.
            Default value: `False` (i.e., input/output does not have Cholesky
            semantics).
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value `NaN` to indicate the
            result is undefined. When `False`, an exception is raised if one or
            more of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          ValueError: If `dimension` is negative.


return self._dimension<EOL>

Dimension of returned correlation matrices.


return self._input_output_cholesky<EOL>

Boolean indicating if `Tensor` input/outputs are Cholesky factorized.


def _sample_n(self, num_samples, seed=None, name=None):

if self.dimension < <NUM_LIT:0>:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>')<EOL><DEDENT>seed = seed_stream.SeedStream(seed, '<STR_LIT>')<EOL>with tf.name_scope('<STR_LIT>' or name):<EOL><INDENT>if not dtype_util.is_floating(self.concentration.dtype):<EOL><INDENT>raise TypeError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT:{}>'.format(dtype_util.name(self.concentration.dtype)))<EOL><DEDENT>concentration = _replicate(num_samples, self.concentration)<EOL>concentration_shape = tf.shape(input=concentration)<EOL>if self.dimension <= <NUM_LIT:1>:<EOL><INDENT>shape = tf.concat([<EOL>concentration_shape, [self.dimension, self.dimension]], axis=<NUM_LIT:0>)<EOL>return tf.ones(shape=shape, dtype=self.concentration.dtype)<EOL><DEDENT>beta_conc = concentration + (self.dimension - <NUM_LIT>) / <NUM_LIT><EOL>beta_dist = beta.Beta(concentration1=beta_conc, concentration0=beta_conc)<EOL>corr12 = <NUM_LIT> * beta_dist.sample(seed=seed()) - <NUM_LIT:1.><EOL>first_row = tf.concat([<EOL>tf.ones_like(corr12)[..., tf.newaxis],<EOL>tf.zeros_like(corr12)[..., tf.newaxis]], axis=-<NUM_LIT:1>)<EOL>second_row = tf.concat([<EOL>corr12[..., tf.newaxis],<EOL>tf.sqrt(<NUM_LIT:1> - corr12**<NUM_LIT:2>)[..., tf.newaxis]], axis=-<NUM_LIT:1>)<EOL>chol_result = tf.concat([<EOL>first_row[..., tf.newaxis, :],<EOL>second_row[..., tf.newaxis, :]], axis=-<NUM_LIT:2>)<EOL>for n in range(<NUM_LIT:2>, self.dimension):<EOL><INDENT>beta_conc -= <NUM_LIT:0.5><EOL>norm = beta.Beta(<EOL>concentration1=n/<NUM_LIT>,<EOL>concentration0=beta_conc<EOL>).sample(seed=seed())<EOL>distance = tf.sqrt(norm)[..., tf.newaxis]<EOL>direction = _uniform_unit_norm(<EOL>n, concentration_shape, self.concentration.dtype, seed)<EOL>raw_correlation = distance * direction  <EOL>new_row = tf.concat(<EOL>[raw_correlation, tf.sqrt(<NUM_LIT:1.> - norm[..., tf.newaxis])], axis=-<NUM_LIT:1>)<EOL>chol_result = tf.concat([<EOL>chol_result,<EOL>tf.zeros_like(chol_result[..., <NUM_LIT:0>][..., tf.newaxis])], axis=-<NUM_LIT:1>)<EOL>chol_result = tf.concat(<EOL>[chol_result, new_row[..., tf.newaxis, :]], axis=-<NUM_LIT:2>)<EOL><DEDENT>if self.input_output_cholesky:<EOL><INDENT>return chol_result<EOL><DEDENT>result = tf.matmul(chol_result, chol_result, transpose_b=True)<EOL>result = tf.linalg.set_diag(<EOL>result,<EOL>tf.ones(shape=tf.shape(input=result)[:-<NUM_LIT:1>], dtype=result.dtype))<EOL>return result<EOL><DEDENT>

Returns a Tensor of samples from an LKJ distribution.

        Args:
          num_samples: Python `int`. The number of samples to draw.
          seed: Python integer seed for RNG
          name: Python `str` name prefixed to Ops created by this function.

        Returns:
          samples: A Tensor of correlation matrices with shape `[n, B, D, D]`,
            where `B` is the shape of the `concentration` parameter, and `D`
            is the `dimension`.

        Raises:
          ValueError: If `dimension` is negative.


def _log_unnorm_prob(self, x, name=None):

with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if self.input_output_cholesky:<EOL><INDENT>logdet = <NUM_LIT> * tf.reduce_sum(<EOL>input_tensor=tf.math.log(tf.linalg.diag_part(x)), axis=[-<NUM_LIT:1>])<EOL><DEDENT>else:<EOL><INDENT>_, logdet = tf.linalg.slogdet(x)<EOL><DEDENT>answer = (self.concentration - <NUM_LIT:1.>) * logdet<EOL>return answer<EOL><DEDENT>

Returns the unnormalized log density of an LKJ distribution.

        Args:
          x: `float` or `double` `Tensor` of correlation matrices.  The shape of `x`
            must be `B + [D, D]`, where `B` broadcasts with the shape of
            `concentration`.
          name: Python `str` name prefixed to Ops created by this function.

        Returns:
          log_p: A Tensor of the unnormalized log density of each matrix element of
            `x`, with respect to an LKJ distribution with parameter the
            corresponding element of `concentration`.


<EOL>with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>logpi = np.log(np.pi)<EOL>ans = tf.zeros_like(self.concentration)<EOL>for k in range(<NUM_LIT:1>, self.dimension):<EOL><INDENT>ans += logpi * (k / <NUM_LIT>)<EOL>ans += tf.math.lgamma(self.concentration +<EOL>(self.dimension - <NUM_LIT:1> - k) / <NUM_LIT>)<EOL>ans -= tf.math.lgamma(self.concentration + (self.dimension - <NUM_LIT:1>) / <NUM_LIT>)<EOL><DEDENT>return ans<EOL><DEDENT>

Returns the log normalization of an LKJ distribution.

        Args:
          name: Python `str` name prefixed to Ops created by this function.

        Returns:
          log_z: A Tensor of the same shape and dtype as `concentration`, containing
            the corresponding log normalizers.


@kullback_leibler.RegisterKL(GaussianProcess, normal.Normal)<EOL>def _kl_gp_normal(gp, n, name=None):

return _kl_gp_compatible(gp, n, name or '<STR_LIT>')<EOL>

Calculate the batched KL divergence KL(gp || n).

    Args:
      gp: instance of a GaussianProcess distribution object.
      n: instance of a Normal distribution object.
      name: (optional) Name to use for created operations.
        default is 'kl_gp_normal'.

    Returns:
      Batchwise KL(gp || n)


@kullback_leibler.RegisterKL(<EOL>GaussianProcess, mvn_linear_operator.MultivariateNormalLinearOperator)<EOL>def _kl_gp_mvn(gp, mvn, name=None):

return _kl_gp_compatible(gp, mvn, name or '<STR_LIT>')<EOL>

Calculate the batched KL divergence KL(gp || mvn).

    Args:
      gp: instance of a GaussianProcess distribution object.
      mvn: instance of a multivariate Normal distribution object (any subclass of
        MultivariateNormalLinearOperator)
      name: (optional) Name to use for created operations.
        default is 'kl_gp_mvn'.

    Returns:
      Batchwise KL(gp || mvn)


@kullback_leibler.RegisterKL(normal.Normal, GaussianProcess)<EOL>def _kl_normal_gp(n, gp, name=None):

return _kl_compatible_gp(n, gp, name or '<STR_LIT>')<EOL>

Calculate the batched KL divergence KL(gp || n).

    Args:
      n: instance of a Normal distribution object.
      gp: instance of a GaussianProcess distribution object.
      name: (optional) Name to use for created operations.
        default is 'kl_normal_gp'.

    Returns:
      Batchwise KL(n || gp)