parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._mixture_distribution = mixture_distribution<EOL>self._components_distribution = components_distribution<EOL>self._runtime_assertions = []<EOL>s = components_distribution.event_shape_tensor()<EOL>self._event_ndims = tf.compat.dimension_value(s.shape[<NUM_LIT:0>])<EOL>if self._event_ndims is None:<EOL><INDENT>self._event_ndims = tf.size(input=s)<EOL><DEDENT>self._event_size = tf.reduce_prod(input_tensor=s)<EOL>if not dtype_util.is_integer(mixture_distribution.dtype):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>".format(<EOL>dtype_util.name(mixture_distribution.dtype)))<EOL><DEDENT>if (tensorshape_util.rank(mixture_distribution.event_shape) is not None<EOL>and tensorshape_util.rank(mixture_distribution.event_shape) != <NUM_LIT:0>):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>elif validate_args:<EOL><INDENT>self._runtime_assertions += [<EOL>assert_util.assert_equal(<EOL>tf.size(input=mixture_distribution.event_shape_tensor()),<EOL><NUM_LIT:0>,<EOL>message="<STR_LIT>"),<EOL>]<EOL><DEDENT>mdbs = mixture_distribution.batch_shape<EOL>cdbs = tensorshape_util.with_rank_at_least(<EOL>components_distribution.batch_shape, <NUM_LIT:1>)[:-<NUM_LIT:1>]<EOL>if tensorshape_util.is_fully_defined(<EOL>mdbs) and tensorshape_util.is_fully_defined(cdbs):<EOL><INDENT>if tensorshape_util.rank(mdbs) != <NUM_LIT:0> and mdbs != cdbs:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>tensorshape_util.as_list(mdbs),<EOL>tensorshape_util.as_list(cdbs)))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>mdbs = mixture_distribution.batch_shape_tensor()<EOL>cdbs = components_distribution.batch_shape_tensor()[:-<NUM_LIT:1>]<EOL>self._runtime_assertions += [<EOL>assert_util.assert_equal(<EOL>distribution_utils.pick_vector(<EOL>mixture_distribution.is_scalar_batch(), cdbs, mdbs),<EOL>cdbs,<EOL>message=(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>"))<EOL>]<EOL><DEDENT>km = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(mixture_distribution.logits.shape,<EOL><NUM_LIT:1>)[-<NUM_LIT:1>])<EOL>kc = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(<EOL>components_distribution.batch_shape, <NUM_LIT:1>)[-<NUM_LIT:1>])<EOL>if km is not None and kc is not None and km != kc:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(km, kc))<EOL><DEDENT>elif validate_args:<EOL><INDENT>km = tf.shape(input=mixture_distribution.logits)[-<NUM_LIT:1>]<EOL>kc = components_distribution.batch_shape_tensor()[-<NUM_LIT:1>]<EOL>self._runtime_assertions += [<EOL>assert_util.assert_equal(<EOL>km,<EOL>kc,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>")),<EOL>]<EOL><DEDENT>elif km is None:<EOL><INDENT>km = tf.shape(input=mixture_distribution.logits)[-<NUM_LIT:1>]<EOL><DEDENT>self._num_components = km<EOL>self._reparameterize = reparameterize<EOL>if reparameterize:<EOL><INDENT>if (self._components_distribution.reparameterization_type !=<EOL>reparameterization.FULLY_REPARAMETERIZED):<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>reparameterization_type = reparameterization.FULLY_REPARAMETERIZED<EOL><DEDENT>else:<EOL><INDENT>reparameterization_type = reparameterization.NOT_REPARAMETERIZED<EOL><DEDENT>super(MixtureSameFamily, self).__init__(<EOL>dtype=self._components_distribution.dtype,<EOL>reparameterization_type=reparameterization_type,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=(<EOL>self._mixture_distribution._graph_parents  <EOL>+ self._components_distribution._graph_parents),  <EOL>name=name)<EOL><DEDENT>

Construct a `MixtureSameFamily` distribution.

        Args:
          mixture_distribution: `tfp.distributions.Categorical`-like instance.
            Manages the probability of selecting components. The number of
            categories must match the rightmost batch dimension of the
            `components_distribution`. Must have either scalar `batch_shape` or
            `batch_shape` matching `components_distribution.batch_shape[:-1]`.
          components_distribution: `tfp.distributions.Distribution`-like instance.
            Right-most batch dimension indexes components.
          reparameterize: Python `bool`, default `False`. Whether to reparameterize
            samples of the distribution using implicit reparameterization gradients
            [(Figurnov et al., 2018)][1]. The gradients for the mixture logits are
            equivalent to the ones described by [(Graves, 2016)][2]. The gradients
            for the components parameters are also computed using implicit
            reparameterization (as opposed to ancestral sampling), meaning that
            all components are updated every step.
            Only works when:
              (1) components_distribution is fully reparameterized;
              (2) components_distribution is either a scalar distribution or
              fully factorized (tfd.Independent applied to a scalar distribution);
              (3) batch shape has a known rank.
            Experimental, may be slow and produce infs/NaNs.
          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 not dtype_util.is_integer(mixture_distribution.dtype)`.
          ValueError: if mixture_distribution does not have scalar `event_shape`.
          ValueError: if `mixture_distribution.batch_shape` and
            `components_distribution.batch_shape[:-1]` are both fully defined and
            the former is neither scalar nor equal to the latter.
          ValueError: if `mixture_distribution` categories does not equal
            `components_distribution` rightmost batch shape.

        #### References

        [1]: Michael Figurnov, Shakir Mohamed and Andriy Mnih. Implicit
             reparameterization gradients. In _Neural Information Processing
             Systems_, 2018. https://arxiv.org/abs/1805.08498

        [2]: Alex Graves. Stochastic Backpropagation through Mixture Density
             Distributions. _arXiv_, 2016. https://arxiv.org/abs/1607.05690


<EOL>x = tf.stop_gradient(x)<EOL>x_2d_shape = [-<NUM_LIT:1>, self._event_size]  <EOL>def reshaped_distributional_transform(x_2d):<EOL><INDENT>return tf.reshape(<EOL>self._distributional_transform(tf.reshape(x_2d, tf.shape(input=x))),<EOL>x_2d_shape)<EOL><DEDENT>transform_2d, jacobian = _value_and_batch_jacobian(<EOL>reshaped_distributional_transform, tf.reshape(x, x_2d_shape))<EOL>transform_2d = _prevent_2nd_derivative(transform_2d)<EOL>surrogate_x_2d = -tf.linalg.triangular_solve(<EOL>tf.stop_gradient(jacobian), tf.expand_dims(transform_2d, axis=-<NUM_LIT:1>),<EOL>lower=True)  <EOL>surrogate_x = tf.reshape(surrogate_x_2d, tf.shape(input=x))<EOL>return x + (surrogate_x - tf.stop_gradient(surrogate_x))<EOL>

Adds reparameterization (pathwise) gradients to samples of the mixture.

        Implicit reparameterization gradients are
           dx/dphi = -(d transform(x, phi) / dx)^-1 * d transform(x, phi) / dphi,
        where transform(x, phi) is distributional transform that removes all
        parameters from samples x.

        We implement them by replacing x with
          -stop_gradient(d transform(x, phi) / dx)^-1 * transform(x, phi)]
        for the backward pass (gradient computation).
        The derivative of this quantity w.r.t. phi is then the implicit
        reparameterization gradient.
        Note that this replaces the gradients w.r.t. both the mixture
        distribution parameters and components distributions parameters.

        Limitations:
          1. Fundamental: components must be fully reparameterized.
          2. Distributional transform is currently only implemented for
            factorized components.
          3. Distributional transform currently only works for known rank of the
            batch tensor.

        Arguments:
          x: Sample of mixture distribution

        Returns:
          Tensor with same value as x, but with reparameterization gradients


if tensorshape_util.rank(x.shape) is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>if isinstance(self._components_distribution, independent.Independent):<EOL><INDENT>univariate_components = self._components_distribution.distribution<EOL><DEDENT>else:<EOL><INDENT>univariate_components = self._components_distribution<EOL><DEDENT>with tf.control_dependencies([<EOL>assert_util.assert_equal(<EOL>univariate_components.is_scalar_event(),<EOL>True,<EOL>message="<STR_LIT>")<EOL>]):<EOL><INDENT>x_padded = self._pad_sample_dims(x)  <EOL>log_prob_x = univariate_components.log_prob(x_padded)  <EOL>cdf_x = univariate_components.cdf(x_padded)  <EOL>cumsum_log_prob_x = tf.reshape(<EOL>tf.math.cumsum(<EOL>tf.reshape(log_prob_x, [-<NUM_LIT:1>, self._event_size]),<EOL>exclusive=True,<EOL>axis=-<NUM_LIT:1>),<EOL>tf.shape(input=log_prob_x))  <EOL>logits_mix_prob = distribution_utils.pad_mixture_dimensions(<EOL>self.mixture_distribution.logits, self, self.mixture_distribution,<EOL>self._event_ndims)  <EOL>log_posterior_weights_x = logits_mix_prob + cumsum_log_prob_x<EOL>component_axis = tensorshape_util.rank(x.shape) - self._event_ndims<EOL>posterior_weights_x = tf.nn.softmax(log_posterior_weights_x,<EOL>axis=component_axis)<EOL>return tf.reduce_sum(<EOL>input_tensor=posterior_weights_x * cdf_x, axis=component_axis)<EOL><DEDENT>

Performs distributional transform of the mixture samples.

        Distributional transform removes the parameters from samples of a
        multivariate distribution by applying conditional CDFs:
          (F(x_1), F(x_2 | x1_), ..., F(x_d | x_1, ..., x_d-1))
        (the indexing is over the "flattened" event dimensions).
        The result is a sample of product of Uniform[0, 1] distributions.

        We assume that the components are factorized, so the conditional CDFs become
          F(x_i | x_1, ..., x_i-1) = sum_k w_i^k F_k (x_i),
        where w_i^k is the posterior mixture weight: for i > 0
          w_i^k = w_k prob_k(x_1, ..., x_i-1) / sum_k' w_k' prob_k'(x_1, ..., x_i-1)
        and w_0^k = w_k is the mixture probability of the k-th component.

        Arguments:
          x: Sample of mixture distribution

        Returns:
          Result of the distributional transform


def _scale_mle(self, samples, scale_candidates):

dist = tfd.Horseshoe(scale=scale_candidates)<EOL>dims = tf.shape(input=scale_candidates)<EOL>num_candidates = dims[-<NUM_LIT:1>]<EOL>original_batch_shape = dims[:-<NUM_LIT:1>]<EOL>log_likelihood = tf.reduce_sum(<EOL>input_tensor=dist.log_prob(<EOL>tf.reshape(samples,<EOL>tf.concat([[-<NUM_LIT:1>], original_batch_shape, [<NUM_LIT:1>]], axis=<NUM_LIT:0>))),<EOL>axis=<NUM_LIT:0>)<EOL>mask = tf.one_hot(<EOL>tf.argmax(input=log_likelihood, axis=-<NUM_LIT:1>),<EOL>depth=num_candidates,<EOL>dtype=self.dtype)<EOL>return tf.reduce_sum(input_tensor=scale_candidates * mask, axis=-<NUM_LIT:1>)<EOL>

Max log-likelihood estimate for scale.

        Args:
          samples: Observed data points.
          scale_candidates: A simple grid of candiates for
            scale, with shape original_batch_shape + [num_candidates],
            where different candidates for a single scalar parameter are at the
            inner most dimension (axis -1).


def quadrature_scheme_lognormal_gauss_hermite(<EOL>loc, scale, quadrature_size,<EOL>validate_args=False, name=None):

with tf.name_scope(<EOL>name or "<STR_LIT>"):<EOL><INDENT>grid, probs = np.polynomial.hermite.hermgauss(deg=quadrature_size)<EOL>npdt = dtype_util.as_numpy_dtype(loc.dtype)<EOL>grid = grid.astype(npdt)<EOL>probs = probs.astype(npdt)<EOL>probs /= np.linalg.norm(probs, ord=<NUM_LIT:1>, keepdims=True)<EOL>probs = tf.convert_to_tensor(value=probs, name="<STR_LIT>", dtype=loc.dtype)<EOL>grid = (loc[..., tf.newaxis] + np.sqrt(<NUM_LIT>) * scale[..., tf.newaxis] * grid)<EOL>return grid, probs<EOL><DEDENT>

Use Gauss-Hermite quadrature to form quadrature on positive-reals.

    Note: for a given `quadrature_size`, this method is generally less accurate
    than `quadrature_scheme_lognormal_quantiles`.

    Args:
      loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
        the LogNormal prior.
      scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
        the LogNormal prior.
      quadrature_size: Python `int` scalar representing the number of quadrature
        points.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      name: Python `str` name prefixed to Ops created by this class.

    Returns:
      grid: (Batch of) length-`quadrature_size` vectors representing the
        `log_rate` parameters of a `Poisson`.
      probs: (Batch of) length-`quadrature_size` vectors representing the
        weight associate with each `grid` value.


def quadrature_scheme_lognormal_quantiles(<EOL>loc, scale, quadrature_size,<EOL>validate_args=False, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>dist = transformed_distribution.TransformedDistribution(<EOL>distribution=normal.Normal(loc=loc, scale=scale),<EOL>bijector=exp_bijector.Exp(),<EOL>validate_args=validate_args)<EOL>batch_ndims = tensorshape_util.rank(dist.batch_shape)<EOL>if batch_ndims is None:<EOL><INDENT>batch_ndims = tf.shape(input=dist.batch_shape_tensor())[<NUM_LIT:0>]<EOL><DEDENT>def _compute_quantiles():<EOL><INDENT>"""<STR_LIT>"""<EOL>zero = tf.zeros([], dtype=dist.dtype)<EOL>edges = tf.linspace(zero, <NUM_LIT:1.>, quadrature_size + <NUM_LIT:3>)[<NUM_LIT:1>:-<NUM_LIT:1>]<EOL>edges = tf.reshape(<EOL>edges,<EOL>shape=tf.concat(<EOL>[[-<NUM_LIT:1>], tf.ones([batch_ndims], dtype=tf.int32)], axis=<NUM_LIT:0>))<EOL>quantiles = dist.quantile(edges)<EOL>perm = tf.concat([tf.range(<NUM_LIT:1>, <NUM_LIT:1> + batch_ndims), [<NUM_LIT:0>]], axis=<NUM_LIT:0>)<EOL>quantiles = tf.transpose(a=quantiles, perm=perm)<EOL>return quantiles<EOL><DEDENT>quantiles = _compute_quantiles()<EOL>grid = (quantiles[..., :-<NUM_LIT:1>] + quantiles[..., <NUM_LIT:1>:]) / <NUM_LIT><EOL>new_shape = tensorshape_util.concatenate(dist.batch_shape,<EOL>[quadrature_size])<EOL>tensorshape_util.set_shape(grid, new_shape)<EOL>probs = tf.fill(<EOL>dims=[quadrature_size], value=<NUM_LIT:1.> / tf.cast(quadrature_size, dist.dtype))<EOL>return grid, probs<EOL><DEDENT>

Use LogNormal quantiles to form quadrature on positive-reals.

    Args:
      loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
        the LogNormal prior.
      scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
        the LogNormal prior.
      quadrature_size: Python `int` scalar representing the number of quadrature
        points.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      name: Python `str` name prefixed to Ops created by this class.

    Returns:
      grid: (Batch of) length-`quadrature_size` vectors representing the
        `log_rate` parameters of a `Poisson`.
      probs: (Batch of) length-`quadrature_size` vectors representing the
        weight associate with each `grid` value.


args_ = [tf.get_static_value(x) for x in args]<EOL>if any(vec is None for vec in args_):<EOL><INDENT>return tf.concat(args, axis=<NUM_LIT:0>)<EOL><DEDENT>return [val for vec in args_ for val in vec]<EOL>

Concatenates input vectors, statically if possible.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([loc, scale], tf.float32)<EOL>if loc is not None:<EOL><INDENT>loc = tf.convert_to_tensor(value=loc, name="<STR_LIT>", dtype=dtype)<EOL><DEDENT>if scale is not None:<EOL><INDENT>scale = tf.convert_to_tensor(value=scale, dtype=dtype, name="<STR_LIT>")<EOL><DEDENT>self._quadrature_grid, self._quadrature_probs = tuple(quadrature_fn(<EOL>loc, scale, quadrature_size, validate_args))<EOL>dt = self._quadrature_grid.dtype<EOL>if not dtype_util.base_equal(dt, self._quadrature_probs.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>dtype_util.name(dt),<EOL>dtype_util.name(self._quadrature_probs.dtype)))<EOL><DEDENT>self._distribution = poisson.Poisson(<EOL>log_rate=self._quadrature_grid,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats)<EOL>self._mixture_distribution = categorical.Categorical(<EOL>logits=tf.math.log(self._quadrature_probs),<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats)<EOL>self._loc = loc<EOL>self._scale = scale<EOL>self._quadrature_size = quadrature_size<EOL>super(PoissonLogNormalQuadratureCompound, self).__init__(<EOL>dtype=dt,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[loc, scale],<EOL>name=name)<EOL><DEDENT>

Constructs the PoissonLogNormalQuadratureCompound`.

        Note: `probs` returned by (optional) `quadrature_fn` are presumed to be
        either a length-`quadrature_size` vector or a batch of vectors in 1-to-1
        correspondence with the returned `grid`. (I.e., broadcasting is only
        partially supported.)

        Args:
          loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
            the LogNormal prior.
          scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
            the LogNormal prior.
          quadrature_size: Python `int` scalar representing the number of quadrature
            points.
          quadrature_fn: Python callable taking `loc`, `scale`,
            `quadrature_size`, `validate_args` and returning `tuple(grid, probs)`
            representing the LogNormal grid and corresponding normalized weight.
            normalized) weight.
            Default value: `quadrature_scheme_lognormal_quantiles`.
          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 `quadrature_grid` and `quadrature_probs` have different base
            `dtype`.


return self._mixture_distribution<EOL>

Distribution which randomly selects a Poisson with quadrature param.


return self._distribution<EOL>

Base Poisson parameterized by a quadrature grid.


return self._loc<EOL>

Location parameter of the LogNormal prior.


return self._scale<EOL>


parameters = dict(locals())<EOL>if not isinstance(cat, categorical.Categorical):<EOL><INDENT>raise TypeError("<STR_LIT>" %<EOL>cat)<EOL><DEDENT>if not components:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if not isinstance(components, (list, tuple)):<EOL><INDENT>raise TypeError("<STR_LIT>" %<EOL>components)<EOL><DEDENT>if not all(isinstance(c, distribution.Distribution) for c in components):<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>" % components)<EOL><DEDENT>dtype = components[<NUM_LIT:0>].dtype<EOL>if not all(d.dtype == dtype for d in components):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>" % [(d.name, d.dtype) for d in components])<EOL><DEDENT>static_event_shape = components[<NUM_LIT:0>].event_shape<EOL>static_batch_shape = cat.batch_shape<EOL>for di, d in enumerate(components):<EOL><INDENT>if not tensorshape_util.is_compatible_with(static_batch_shape,<EOL>d.batch_shape):<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(di))<EOL><DEDENT>static_event_shape = tensorshape_util.merge_with(<EOL>static_event_shape, d.event_shape)<EOL>static_batch_shape = tensorshape_util.merge_with(<EOL>static_batch_shape, d.batch_shape)<EOL><DEDENT>if tensorshape_util.rank(static_event_shape) is None:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>with tf.name_scope(name) as name:<EOL><INDENT>num_components = cat.event_size<EOL>static_num_components = tf.get_static_value(num_components)<EOL>if static_num_components is None:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>static_num_components = int(static_num_components)<EOL>if static_num_components != len(components):<EOL><INDENT>raise ValueError("<STR_LIT>" %<EOL>(static_num_components, len(components)))<EOL><DEDENT>cat_batch_shape = cat.batch_shape_tensor()<EOL>cat_batch_rank = tf.size(input=cat_batch_shape)<EOL>if validate_args:<EOL><INDENT>batch_shapes = [d.batch_shape_tensor() for d in components]<EOL>batch_ranks = [tf.size(input=bs) for bs in batch_shapes]<EOL>check_message = ("<STR_LIT>"<EOL>"<STR_LIT>")<EOL>self._assertions = [<EOL>assert_util.assert_equal(<EOL>cat_batch_rank, batch_ranks[di], message=check_message % di)<EOL>for di in range(len(components))<EOL>]<EOL>self._assertions += [<EOL>assert_util.assert_equal(<EOL>cat_batch_shape, batch_shapes[di], message=check_message % di)<EOL>for di in range(len(components))<EOL>]<EOL><DEDENT>else:<EOL><INDENT>self._assertions = []<EOL><DEDENT>self._cat = cat<EOL>self._components = list(components)<EOL>self._num_components = static_num_components<EOL>self._static_event_shape = static_event_shape<EOL>self._static_batch_shape = static_batch_shape<EOL>self._use_static_graph = use_static_graph<EOL>if use_static_graph and static_num_components is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT><DEDENT>graph_parents = self._cat._graph_parents  <EOL>for c in self._components:<EOL><INDENT>graph_parents += c._graph_parents  <EOL><DEDENT>super(Mixture, 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=graph_parents,<EOL>name=name)<EOL>

Initialize a Mixture distribution.

        A `Mixture` is defined by a `Categorical` (`cat`, representing the
        mixture probabilities) and a list of `Distribution` objects
        all having matching dtype, batch shape, event shape, and continuity
        properties (the components).

        The `num_classes` of `cat` must be possible to infer at graph construction
        time and match `len(components)`.

        Args:
          cat: A `Categorical` distribution instance, representing the probabilities
              of `distributions`.
          components: A list or tuple of `Distribution` instances.
            Each instance must have the same type, be defined on the same domain,
            and have matching `event_shape` and `batch_shape`.
          validate_args: Python `bool`, default `False`. If `True`, raise a runtime
            error if batch or event ranks are inconsistent between cat and any of
            the distributions. This is only checked if the ranks cannot be
            determined statically at graph construction time.
          allow_nan_stats: Boolean, 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.
          use_static_graph: Calls to `sample` will not rely on dynamic tensor
            indexing, allowing for some static graph compilation optimizations, but
            at the expense of sampling all underlying distributions in the mixture.
            (Possibly useful when running on TPUs).
            Default value: `False` (i.e., use dynamic indexing).
          name: A name for this distribution (optional).

        Raises:
          TypeError: If cat is not a `Categorical`, or `components` is not
            a list or tuple, or the elements of `components` are not
            instances of `Distribution`, or do not have matching `dtype`.
          ValueError: If `components` is an empty list or tuple, or its
            elements do not have a statically known event rank.
            If `cat.num_classes` cannot be inferred at graph creation time,
            or the constant value of `cat.num_classes` is not equal to
            `len(components)`, or all `components` and `cat` do not have
            matching static batch shapes, or all components do not
            have matching static event shapes.


expanded_x = x<EOL>for _ in range(tensorshape_util.rank(self.event_shape)):<EOL><INDENT>expanded_x = tf.expand_dims(expanded_x, -<NUM_LIT:1>)<EOL><DEDENT>return expanded_x<EOL>

Expand the rank of x up to static_event_rank times for broadcasting.

        The static event rank was checked to not be None at construction time.

        Args:
          x: A tensor to expand.
        Returns:
          The expanded tensor.


with self._name_scope(name):<EOL><INDENT>with tf.control_dependencies(self._assertions):<EOL><INDENT>distribution_entropies = [d.entropy() for d in self.components]<EOL>cat_probs = self._cat_probs(log_probs=False)<EOL>partial_entropies = [<EOL>c_p * m for (c_p, m) in zip(cat_probs, distribution_entropies)<EOL>]<EOL>return tf.add_n(partial_entropies)<EOL><DEDENT><DEDENT>

r"""A lower bound on the entropy of this mixture model.

        The bound below is not always very tight, and its usefulness depends
        on the mixture probabilities and the components in use.

        A lower bound is useful for ELBO when the `Mixture` is the variational
        distribution:

        \\(
        \log p(x) >= ELBO = \int q(z) \log p(x, z) dz + H[q]
        \\)

        where \\( p \\) is the prior distribution, \\( q \\) is the variational,
        and \\( H[q] \\) is the entropy of \\( q \\). If there is a lower bound
        \\( G[q] \\) such that \\( H[q] \geq G[q] \\) then it can be used in
        place of \\( H[q] \\).

        For a mixture of distributions \\( q(Z) = \sum_i c_i q_i(Z) \\) with
        \\( \sum_i c_i = 1 \\), by the concavity of \\( f(x) = -x \log x \\), a
        simple lower bound is:

        \\(
        \begin{align}
        H[q] & = - \int q(z) \log q(z) dz \\\
           & = - \int (\sum_i c_i q_i(z)) \log(\sum_i c_i q_i(z)) dz \\\
           & \geq - \sum_i c_i \int q_i(z) \log q_i(z) dz \\\
           & = \sum_i c_i H[q_i]
        \end{align}
        \\)

        This is the term we calculate below for \\( G[q] \\).

        Args:
          name: A name for this operation (optional).

        Returns:
          A lower bound on the Mixture's entropy.


which_softmax = tf.nn.log_softmax if log_probs else tf.nn.softmax<EOL>cat_probs = which_softmax(self.cat.logits)<EOL>cat_probs = tf.unstack(cat_probs, num=self.num_components, axis=-<NUM_LIT:1>)<EOL>return cat_probs<EOL>

Get a list of num_components batchwise probabilities.


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) as 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(VectorExponentialLinearOperator, self).__init__(<EOL>distribution=exponential.Exponential(<EOL>rate=tf.ones([], dtype=scale.dtype),<EOL>allow_nan_stats=allow_nan_stats),<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 Exponential distribution supported on a subset of `R^k`.

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

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

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

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


shape = tensorshape_util.concatenate(self.batch_shape, self.event_shape)<EOL>has_static_shape = tensorshape_util.is_fully_defined(shape)<EOL>if not has_static_shape:<EOL><INDENT>shape = tf.concat([<EOL>self.batch_shape_tensor(),<EOL>self.event_shape_tensor(),<EOL>], <NUM_LIT:0>)<EOL><DEDENT>return shape<EOL>


@kullback_leibler.RegisterKL(VonMises, VonMises)<EOL>def _kl_von_mises_von_mises(d1, d2, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>i0e_concentration1 = tf.math.bessel_i0e(d1.concentration)<EOL>i1e_concentration1 = tf.math.bessel_i1e(d1.concentration)<EOL>i0e_concentration2 = tf.math.bessel_i0e(d2.concentration)<EOL>return ((d2.concentration - d1.concentration) +<EOL>tf.math.log(i0e_concentration2 / i0e_concentration1) +<EOL>(d1.concentration - d2.concentration * tf.cos(d1.loc - d2.loc)) *<EOL>(i1e_concentration1 / i0e_concentration1))<EOL><DEDENT>

Batchwise KL divergence KL(d1 || d2) with d1 and d2 von Mises.

    Args:
      d1: instance of a von Mises distribution object.
      d2: instance of a a von Mises distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_von_mises_von_mises".

    Returns:
      Batchwise KL(d1 || d2)


x = tf.convert_to_tensor(value=x)<EOL>concentration = tf.convert_to_tensor(value=concentration)<EOL>dtype = x.dtype<EOL>num_periods = tf.round(x / (<NUM_LIT> * np.pi))<EOL>x -= (<NUM_LIT> * np.pi) * num_periods<EOL>ck = <NUM_LIT><EOL>num_terms = <NUM_LIT:20><EOL>cdf_series, dcdf_dconcentration_series = _von_mises_cdf_series(<EOL>x, concentration, num_terms, dtype)<EOL>cdf_normal, dcdf_dconcentration_normal = _von_mises_cdf_normal(<EOL>x, concentration, dtype)<EOL>use_series = concentration < ck<EOL>cdf = tf.where(use_series, cdf_series, cdf_normal)<EOL>cdf += num_periods<EOL>dcdf_dconcentration = tf.where(use_series, dcdf_dconcentration_series,<EOL>dcdf_dconcentration_normal)<EOL>def grad(dy):<EOL><INDENT>prob = tf.exp(concentration * (tf.cos(x) - <NUM_LIT:1.>)) / (<EOL>(<NUM_LIT> * np.pi) * tf.math.bessel_i0e(concentration))<EOL>return dy * prob, dy * dcdf_dconcentration<EOL><DEDENT>return cdf, grad<EOL>

Computes the cumulative density function (CDF) of von Mises distribution.

    Denote the density of vonMises(loc=0, concentration=concentration) by p(t).
    Note that p(t) is periodic, p(t) = p(t + 2 pi).
    The CDF at the point x is defined as int_{-pi}^x p(t) dt.
    Thus, when x in [-pi, pi], the CDF is in [0, 1]; when x is in [pi, 3pi], the
    CDF is in [1, 2], etc.

    The CDF is not available in closed form. Instead, we use the method [1]
    which uses either a series expansion or a Normal approximation, depending on
    the value of concentration.

    We also compute the derivative of the CDF w.r.t. both x and concentration.
    The derivative w.r.t. x is p(x), while the derivative w.r.t. concentration is
    computed
    using automatic differentiation. We use forward mode for the series case
    (which allows to save memory) and backward mode for the Normal approximation.

    Arguments:
      x: The point at which to evaluate the CDF.
      concentration: The concentration parameter of the von Mises distribution.

    Returns:
      The value of the CDF computed elementwise.

    References:
      [1] G. Hill "Algorithm 518: Incomplete Bessel Function I_0. The Von Mises
      Distribution." ACM Transactions on Mathematical Software, 1977


def _von_mises_cdf_series(x, concentration, num_terms, dtype):

<EOL>num_terms = tf.cast(num_terms, dtype=dtype)<EOL>def loop_body(n, rn, drn_dconcentration, vn, dvn_dconcentration):<EOL><INDENT>"""<STR_LIT>"""<EOL>denominator = <NUM_LIT> * n / concentration + rn<EOL>ddenominator_dk = -<NUM_LIT> * n / concentration ** <NUM_LIT:2> + drn_dconcentration<EOL>rn = <NUM_LIT:1.> / denominator<EOL>drn_dconcentration = -ddenominator_dk / denominator ** <NUM_LIT:2><EOL>multiplier = tf.sin(n * x) / n + vn<EOL>vn = rn * multiplier<EOL>dvn_dconcentration = (drn_dconcentration * multiplier +<EOL>rn * dvn_dconcentration)<EOL>n -= <NUM_LIT:1.><EOL>return n, rn, drn_dconcentration, vn, dvn_dconcentration<EOL><DEDENT>(_, _, _, vn, dvn_dconcentration) = tf.while_loop(<EOL>cond=lambda n, *_: n > <NUM_LIT:0.>,<EOL>body=loop_body,<EOL>loop_vars=(<EOL>num_terms,  <EOL>tf.zeros_like(x, name="<STR_LIT>"),<EOL>tf.zeros_like(x, name="<STR_LIT>"),<EOL>tf.zeros_like(x, name="<STR_LIT>"),<EOL>tf.zeros_like(x, name="<STR_LIT>"),<EOL>),<EOL>)<EOL>cdf = <NUM_LIT> + x / (<NUM_LIT> * np.pi) + vn / np.pi<EOL>dcdf_dconcentration = dvn_dconcentration / np.pi<EOL>cdf_clipped = tf.clip_by_value(cdf, <NUM_LIT:0.>, <NUM_LIT:1.>)<EOL>dcdf_dconcentration *= tf.cast((cdf >= <NUM_LIT:0.>) & (cdf <= <NUM_LIT:1.>), dtype)<EOL>return cdf_clipped, dcdf_dconcentration<EOL>

Computes the von Mises CDF and its derivative via series expansion.


def _von_mises_cdf_normal(x, concentration, dtype):

def cdf_func(concentration):<EOL><INDENT>"""<STR_LIT>"""<EOL>z = ((np.sqrt(<NUM_LIT> / np.pi) / tf.math.bessel_i0e(concentration)) *<EOL>tf.sin(<NUM_LIT> * x))<EOL>z2 = z ** <NUM_LIT:2><EOL>z3 = z2 * z<EOL>z4 = z2 ** <NUM_LIT:2><EOL>c = <NUM_LIT> * concentration<EOL>c1 = <NUM_LIT><EOL>xi = z - z3 / ((c - <NUM_LIT> * z2 - <NUM_LIT>) / <NUM_LIT> -<EOL>(z4 + (<NUM_LIT> / <NUM_LIT>) * z2 + <NUM_LIT> / <NUM_LIT>) / (c - c1 - z2 + <NUM_LIT>)) ** <NUM_LIT:2><EOL>distrib = normal.Normal(tf.cast(<NUM_LIT:0.>, dtype), tf.cast(<NUM_LIT:1.>, dtype))<EOL>return distrib.cdf(xi)<EOL><DEDENT>return value_and_gradient(cdf_func, concentration)<EOL>

Computes the von Mises CDF and its derivative via Normal approximation.


def random_von_mises(shape, concentration, dtype=tf.float32, seed=None):

seed = SeedStream(seed, salt="<STR_LIT>")<EOL>concentration = tf.convert_to_tensor(<EOL>value=concentration, dtype=dtype, name="<STR_LIT>")<EOL>@tf.custom_gradient<EOL>def rejection_sample_with_gradient(concentration):<EOL><INDENT>"""<STR_LIT>"""<EOL>r = <NUM_LIT:1.> + tf.sqrt(<NUM_LIT:1.> + <NUM_LIT> * concentration ** <NUM_LIT:2>)<EOL>rho = (r - tf.sqrt(<NUM_LIT> * r)) / (<NUM_LIT> * concentration)<EOL>s_exact = (<NUM_LIT:1.> + rho ** <NUM_LIT:2>) / (<NUM_LIT> * rho)<EOL>s_approximate = <NUM_LIT:1.> / concentration<EOL>s_concentration_cutoff_dict = {<EOL>tf.float16: <NUM_LIT>,<EOL>tf.float32: <NUM_LIT>,<EOL>tf.float64: <NUM_LIT>,<EOL>}<EOL>s_concentration_cutoff = s_concentration_cutoff_dict[dtype]<EOL>s = tf.where(concentration > s_concentration_cutoff, s_exact, s_approximate)<EOL>def loop_body(done, u, w):<EOL><INDENT>"""<STR_LIT>"""<EOL>u = tf.random.uniform(<EOL>shape, minval=-<NUM_LIT:1.>, maxval=<NUM_LIT:1.>, dtype=dtype, seed=seed())<EOL>z = tf.cos(np.pi * u)<EOL>w = tf.where(done, w, (<NUM_LIT:1.> + s * z) / (s + z))<EOL>y = concentration * (s - w)<EOL>v = tf.random.uniform(<EOL>shape, minval=<NUM_LIT:0.>, maxval=<NUM_LIT:1.>, dtype=dtype, seed=seed())<EOL>accept = (y * (<NUM_LIT> - y) >= v) | (tf.math.log(y / v) + <NUM_LIT:1.> >= y)<EOL>return done | accept, u, w<EOL><DEDENT>_, u, w = tf.while_loop(<EOL>cond=lambda done, *_: ~tf.reduce_all(input_tensor=done),<EOL>body=loop_body,<EOL>loop_vars=(<EOL>tf.zeros(shape, dtype=tf.bool, name="<STR_LIT>"),<EOL>tf.zeros(shape, dtype=dtype, name="<STR_LIT:u>"),<EOL>tf.zeros(shape, dtype=dtype, name="<STR_LIT:w>"),<EOL>),<EOL>maximum_iterations=<NUM_LIT:100>,<EOL>parallel_iterations=<NUM_LIT:1> if seed.original_seed is None else <NUM_LIT:10>,<EOL>)<EOL>x = tf.sign(u) * tf.math.acos(w)<EOL>def grad(dy):<EOL><INDENT>"""<STR_LIT>"""<EOL>broadcast_concentration = concentration + tf.zeros_like(x)<EOL>_, dcdf_dconcentration = value_and_gradient(<EOL>lambda conc: von_mises_cdf(x, conc), broadcast_concentration)<EOL>inv_prob = tf.exp(-broadcast_concentration * (tf.cos(x) - <NUM_LIT:1.>)) * (<EOL>(<NUM_LIT> * np.pi) * tf.math.bessel_i0e(broadcast_concentration))<EOL>ret = dy * (-inv_prob * dcdf_dconcentration)<EOL>num_sample_dimensions = (tf.rank(broadcast_concentration) -<EOL>tf.rank(concentration))<EOL>return tf.reduce_sum(<EOL>input_tensor=ret, axis=tf.range(num_sample_dimensions))<EOL><DEDENT>return x, grad<EOL><DEDENT>return rejection_sample_with_gradient(concentration)<EOL>

Samples from the standardized von Mises distribution.

    The distribution is vonMises(loc=0, concentration=concentration), so the mean
    is zero.
    The location can then be changed by adding it to the samples.

    The sampling algorithm is rejection sampling with wrapped Cauchy proposal [1].
    The samples are pathwise differentiable using the approach of [2].

    Arguments:
      shape: The output sample shape.
      concentration: The concentration parameter of the von Mises distribution.
      dtype: The data type of concentration and the outputs.
      seed: (optional) The random seed.

    Returns:
      Differentiable samples of standardized von Mises.

    References:
      [1] Luc Devroye "Non-Uniform Random Variate Generation", Springer-Verlag,
      1986; Chapter 9, p. 473-476.
      http://www.nrbook.com/devroye/Devroye_files/chapter_nine.pdf
      + corrections http://www.nrbook.com/devroye/Devroye_files/errors.pdf
      [2] Michael Figurnov, Shakir Mohamed, Andriy Mnih. "Implicit
      Reparameterization Gradients", 2018.


parameters = dict(locals())<EOL>with tf.name_scope(name) as 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_non_negative(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>dtype_util.assert_same_float_dtype([self._loc, self._concentration])<EOL><DEDENT><DEDENT>super(VonMises, self).__init__(<EOL>dtype=self._concentration.dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._loc, self._concentration],<EOL>name=name)<EOL>

Construct von Mises distributions with given location and concentration.

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

        Args:
          loc: Floating point tensor, the circular means of the distribution(s).
          concentration: Floating point tensor, the level of concentration of the
            distribution(s) around `loc`. Must take non-negative values.
            `concentration = 0` defines a Uniform distribution, while
            `concentration = +inf` indicates a Deterministic distribution at `loc`.
          validate_args: Python `bool`, default `False`. When `True` distribution
            parameters are checked for validity despite possibly degrading runtime
            performance. When `False` invalid inputs may silently render incorrect
            outputs.
          allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
            (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
            result is undefined. When `False`, an exception is raised if one or more
            of the statistic's batch members are undefined.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          TypeError: if loc and concentration are different dtypes.


return self._loc<EOL>

Distribution parameter for the circular mean (loc).


return self._concentration<EOL>

Distribution parameter for the concentration.


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

Standardize input `x` to a zero-loc von Mises.


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>dtype_util.assert_same_float_dtype([self._loc, self._scale])<EOL><DEDENT><DEDENT>super(Logistic, 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 Logistic distributions with mean and scale `loc` and `scale`.

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

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

        Raises:
          TypeError: if loc and scale are different dtypes.


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>

Standardize input `x` to a unit logistic.


parameters = dict(locals())<EOL>def _convert_to_tensor(x, name, dtype=None):<EOL><INDENT>return None if x is None else tf.convert_to_tensor(<EOL>value=x, name=name, dtype=dtype)<EOL><DEDENT>with tf.name_scope(name) as name:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>dtype = dtype_util.common_dtype([<EOL>loc, scale_diag, scale_identity_multiplier, scale_perturb_factor,<EOL>scale_perturb_diag<EOL>], tf.float32)<EOL>has_low_rank = (scale_perturb_factor is not None or<EOL>scale_perturb_diag is not None)<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=validate_args,<EOL>assert_positive=has_low_rank,<EOL>dtype=dtype)<EOL>scale_perturb_factor = _convert_to_tensor(<EOL>scale_perturb_factor, name="<STR_LIT>", dtype=dtype)<EOL>scale_perturb_diag = _convert_to_tensor(<EOL>scale_perturb_diag, name="<STR_LIT>", dtype=dtype)<EOL>if has_low_rank:<EOL><INDENT>scale = tf.linalg.LinearOperatorLowRankUpdate(<EOL>scale,<EOL>u=scale_perturb_factor,<EOL>diag_update=scale_perturb_diag,<EOL>is_diag_update_positive=scale_perturb_diag is None,<EOL>is_non_singular=True,  <EOL>is_self_adjoint=True,<EOL>is_positive_definite=True,<EOL>is_square=True)<EOL><DEDENT><DEDENT><DEDENT>super(MultivariateNormalDiagPlusLowRank, self).__init__(<EOL>loc=loc,<EOL>scale=scale,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL>self._parameters = parameters<EOL>

Construct Multivariate Normal distribution on `R^k`.

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

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

        Recall that `covariance = scale @ scale.T`. A (non-batch) `scale` matrix is:

        ```none
        scale = diag(scale_diag + scale_identity_multiplier ones(k)) +
            scale_perturb_factor @ diag(scale_perturb_diag) @ scale_perturb_factor.T
        ```

        where:

        * `scale_diag.shape = [k]`,
        * `scale_identity_multiplier.shape = []`,
        * `scale_perturb_factor.shape = [k, r]`, typically `k >> r`, and,
        * `scale_perturb_diag.shape = [r]`.

        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`.
          scale_perturb_factor: Floating-point `Tensor` representing a rank-`r`
            perturbation added to `scale`. May have shape `[B1, ..., Bb, k, r]`,
            `b >= 0`, and characterizes `b`-batches of rank-`r` updates to `scale`.
            When `None`, no rank-`r` update is added to `scale`.
          scale_perturb_diag: Floating-point `Tensor` representing a diagonal matrix
            inside the rank-`r` perturbation added to `scale`. May have shape
            `[B1, ..., Bb, r]`, `b >= 0`, and characterizes `b`-batches of `r x r`
            diagonal matrices inside the perturbation added to `scale`. When
            `None`, an identity matrix is used inside the perturbation. Can only be
            specified if `scale_perturb_factor` is also specified.
          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.


@kullback_leibler.RegisterKL(Independent, Independent)<EOL>def _kl_independent(a, b, name="<STR_LIT>"):

p = a.distribution<EOL>q = b.distribution<EOL>if (tensorshape_util.is_fully_defined(a.event_shape) and<EOL>tensorshape_util.is_fully_defined(b.event_shape)):<EOL><INDENT>if a.event_shape == b.event_shape:<EOL><INDENT>if p.event_shape == q.event_shape:<EOL><INDENT>num_reduce_dims = (tensorshape_util.rank(a.event_shape) -<EOL>tensorshape_util.rank(p.event_shape))<EOL>reduce_dims = [-i - <NUM_LIT:1> for i in range(<NUM_LIT:0>, num_reduce_dims)]<EOL>return tf.reduce_sum(<EOL>input_tensor=kullback_leibler.kl_divergence(p, q, name=name),<EOL>axis=reduce_dims)<EOL><DEDENT>else:<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT><DEDENT>else:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT><DEDENT>else:<EOL><INDENT>with tf.control_dependencies(<EOL>[<EOL>assert_util.assert_equal(a.event_shape_tensor(),<EOL>b.event_shape_tensor()),<EOL>assert_util.assert_equal(p.event_shape_tensor(),<EOL>q.event_shape_tensor())<EOL>]):<EOL><INDENT>num_reduce_dims = (<EOL>prefer_static.rank_from_shape(<EOL>a.event_shape_tensor, a.event_shape) -<EOL>prefer_static.rank_from_shape(<EOL>p.event_shape_tensor, a.event_shape))<EOL>reduce_dims = prefer_static.range(-num_reduce_dims - <NUM_LIT:1>, -<NUM_LIT:1>, <NUM_LIT:1>)<EOL>return tf.reduce_sum(<EOL>input_tensor=kullback_leibler.kl_divergence(p, q, name=name),<EOL>axis=reduce_dims)<EOL><DEDENT><DEDENT>

Batched KL divergence `KL(a || b)` for Independent distributions.

    We can leverage the fact that
    ```
    KL(Independent(a) || Independent(b)) = sum(KL(a || b))
    ```
    where the sum is over the `reinterpreted_batch_ndims`.

    Args:
      a: Instance of `Independent`.
      b: Instance of `Independent`.
      name: (optional) name to use for created ops. Default "kl_independent".

    Returns:
      Batchwise `KL(a || b)`.

    Raises:
      ValueError: If the event space for `a` and `b`, or their underlying
        distributions don't match.


def __init__(<EOL>self, distribution, reinterpreted_batch_ndims=None,<EOL>validate_args=False, name=None):

parameters = dict(locals())<EOL>name = name or "<STR_LIT>" + distribution.name<EOL>self._distribution = distribution<EOL>with tf.name_scope(name) as name:<EOL><INDENT>if reinterpreted_batch_ndims is None:<EOL><INDENT>reinterpreted_batch_ndims = self._get_default_reinterpreted_batch_ndims(<EOL>distribution)<EOL><DEDENT>reinterpreted_batch_ndims = tf.convert_to_tensor(<EOL>value=reinterpreted_batch_ndims,<EOL>dtype=tf.int32,<EOL>name="<STR_LIT>")<EOL>self._reinterpreted_batch_ndims = reinterpreted_batch_ndims<EOL>self._static_reinterpreted_batch_ndims = tf.get_static_value(<EOL>reinterpreted_batch_ndims)<EOL>if self._static_reinterpreted_batch_ndims is not None:<EOL><INDENT>self._reinterpreted_batch_ndims = self._static_reinterpreted_batch_ndims<EOL><DEDENT>super(Independent, self).__init__(<EOL>dtype=self._distribution.dtype,<EOL>reparameterization_type=self._distribution.reparameterization_type,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=self._distribution.allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=(<EOL>[reinterpreted_batch_ndims] +<EOL>distribution._graph_parents),  <EOL>name=name)<EOL>self._runtime_assertions = self._make_runtime_assertions(<EOL>distribution, reinterpreted_batch_ndims, validate_args)<EOL><DEDENT>

Construct a `Independent` distribution.

        Args:
          distribution: The base distribution instance to transform. Typically an
            instance of `Distribution`.
          reinterpreted_batch_ndims: Scalar, integer number of rightmost batch dims
            which will be regarded as event dims. When `None` all but the first
            batch axis (batch axis 0) will be transferred to event dimensions
            (analogous to `tf.layers.flatten`).
          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.
          name: The name for ops managed by the distribution.
            Default value: `Independent + distribution.name`.

        Raises:
          ValueError: if `reinterpreted_batch_ndims` exceeds
            `distribution.batch_ndims`


ndims = prefer_static.rank_from_shape(<EOL>distribution.batch_shape_tensor, distribution.batch_shape)<EOL>return prefer_static.maximum(<NUM_LIT:0>, ndims - <NUM_LIT:1>)<EOL>

Computes the default value for reinterpreted_batch_ndim __init__ arg.


def _build_iid_normal_model(self, num_timesteps, latent_size,<EOL>observation_size, transition_variance,<EOL>observation_variance):

transition_variance = self._build_placeholder(transition_variance)<EOL>observation_variance = self._build_placeholder(observation_variance)<EOL>random_orthogonal_matrix = lambda: np.linalg.qr(<EOL>np.random.randn(latent_size, latent_size))[<NUM_LIT:0>][:observation_size, :]<EOL>observation_matrix = self._build_placeholder(random_orthogonal_matrix())<EOL>model = tfd.LinearGaussianStateSpaceModel(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=self._build_placeholder(<EOL>np.zeros([latent_size, latent_size])),<EOL>transition_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.sqrt(transition_variance) *<EOL>tf.ones([latent_size], dtype=self.dtype)),<EOL>observation_matrix=observation_matrix,<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.sqrt(observation_variance) *<EOL>tf.ones([observation_size], dtype=self.dtype)),<EOL>initial_state_prior=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.sqrt(transition_variance) *<EOL>tf.ones([latent_size], dtype=self.dtype)),<EOL>validate_args=True)<EOL>return model<EOL>

Build a model whose outputs are IID normal by construction.


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

Convert a numpy array to a TF placeholder.

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

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


prior_batch_shape = (<EOL>[] if prior_batch_shape is None else prior_batch_shape)<EOL>transition_matrix_batch_shape = ([] if transition_matrix_batch_shape is None<EOL>else transition_matrix_batch_shape)<EOL>transition_noise_batch_shape = ([] if transition_noise_batch_shape is None<EOL>else transition_noise_batch_shape)<EOL>observation_matrix_batch_shape = ([]<EOL>if observation_matrix_batch_shape is None<EOL>else observation_matrix_batch_shape)<EOL>observation_noise_batch_shape = ([] if observation_noise_batch_shape is None<EOL>else observation_noise_batch_shape)<EOL>return tfd.LinearGaussianStateSpaceModel(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=tf.random.normal(transition_matrix_batch_shape +<EOL>[latent_size, latent_size]),<EOL>transition_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.nn.softplus(<EOL>tf.random.normal(transition_noise_batch_shape +<EOL>[latent_size]))),<EOL>observation_matrix=tf.random.normal(observation_matrix_batch_shape +<EOL>[observation_size, latent_size]),<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.nn.softplus(<EOL>tf.random.normal(observation_noise_batch_shape +<EOL>[observation_size]))),<EOL>initial_state_prior=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.nn.softplus(<EOL>tf.random.normal(prior_batch_shape + [latent_size]))),<EOL>validate_args=True)<EOL>

Builds a LGSSM with random normal ops of specified shape.


x = np.tril(linalg.circulant(np.arange(start, start + n)))<EOL>return np.dot(x, x.T)<EOL>

Deterministically create a positive definite matrix.


return linalg.cholesky(x).T<EOL>


df = np.asarray(df)[..., np.newaxis, np.newaxis]<EOL>x = np.sqrt(df) * np.asarray(x)<EOL>d = np.diagonal(x, axis1=-<NUM_LIT:2>, axis2=-<NUM_LIT:1>)<EOL>return x**<NUM_LIT:2> + np.matmul(<EOL>d[..., np.newaxis], d[..., np.newaxis, :])<EOL>

Compute Wishart variance for numpy scale matrix.


with tf.compat.v1.name_scope(<EOL>name, values=[amplitude, length_scale]) as name:<EOL><INDENT>dtype = util.maybe_get_common_dtype(<EOL>[amplitude, length_scale])<EOL>if amplitude is not None:<EOL><INDENT>amplitude = tf.convert_to_tensor(<EOL>value=amplitude, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._amplitude = _validate_arg_if_not_none(<EOL>amplitude, tf.compat.v1.assert_positive, validate_args)<EOL>if length_scale is not None:<EOL><INDENT>length_scale = tf.convert_to_tensor(<EOL>value=length_scale, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._length_scale = _validate_arg_if_not_none(<EOL>length_scale, tf.compat.v1.assert_positive, validate_args)<EOL><DEDENT>super(ExponentiatedQuadratic, self).__init__(<EOL>feature_ndims, dtype=dtype, name=name)<EOL>

Construct an ExponentiatedQuadratic kernel instance.

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


return self._amplitude<EOL>


return self._length_scale<EOL>


with tf.compat.v1.name_scope(<EOL>name, values=[base_kernel, fixed_inputs]) as name:<EOL><INDENT>if base_kernel.dtype is None:<EOL><INDENT>dtype = None if fixed_inputs is None else fixed_inputs.dtype<EOL><DEDENT>else:<EOL><INDENT>dtype = dtype_util.common_dtype([base_kernel, fixed_inputs], tf.float32)<EOL><DEDENT>self._base_kernel = base_kernel<EOL>self._fixed_inputs = (None if fixed_inputs is None else<EOL>tf.convert_to_tensor(value=fixed_inputs,<EOL>dtype=dtype))<EOL>if not self._is_empty_fixed_inputs():<EOL><INDENT>self._divisor_matrix = base_kernel.matrix(fixed_inputs, fixed_inputs)<EOL>if diag_shift is not None:<EOL><INDENT>self._divisor_matrix = _add_diagonal_shift(<EOL>self._divisor_matrix, diag_shift)<EOL><DEDENT><DEDENT>self._cholesky_bijector = invert.Invert(<EOL>cholesky_outer_product.CholeskyOuterProduct())<EOL><DEDENT>super(SchurComplement, self).__init__(<EOL>base_kernel.feature_ndims, dtype=dtype, name=name)<EOL>

Construct a SchurComplement kernel instance.

        Args:
          base_kernel: A `PositiveSemidefiniteKernel` instance, the kernel used to
            build the block matrices of which this kernel computes the  Schur
            complement.
          fixed_inputs: A Tensor, representing a collection of inputs. The Schur
            complement that this kernel computes comes from a block matrix, whose
            bottom-right corner is derived from `base_kernel.matrix(fixed_inputs,
            fixed_inputs)`, and whose top-right and bottom-left pieces are
            constructed by computing the base_kernel at pairs of input locations
            together with these `fixed_inputs`. `fixed_inputs` is allowed to be an
            empty collection (either `None` or having a zero shape entry), in which
            case the kernel falls back to the trivial application of `base_kernel`
            to inputs. See class-level docstring for more details on the exact
            computation this does; `fixed_inputs` correspond to the `Z` structure
            discussed there. `fixed_inputs` is assumed to have shape `[b1, ..., bB,
            N, f1, ..., fF]` where the `b`'s are batch shape entries, the `f`'s are
            feature_shape entries, and `N` is the number of fixed inputs. Use of
            this kernel entails a 1-time O(N^3) cost of computing the Cholesky
            decomposition of the k(Z, Z) matrix. The batch shape elements of
            `fixed_inputs` must be broadcast compatible with
            `base_kernel.batch_shape`.
          diag_shift: A floating point scalar to be added to the diagonal of the
            divisor_matrix before computing its Cholesky.
          validate_args: If `True`, parameters are checked for validity despite
            possibly degrading runtime performance.
            Default value: `False`
          name: Python `str` name prefixed to Ops created by this class.
            Default value: `"SchurComplement"`


if not (isinstance(ndims, int) and ndims >= <NUM_LIT:0>):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>.format(ndims))<EOL><DEDENT>if ndims == <NUM_LIT:0>:<EOL><INDENT>return x<EOL><DEDENT>x = tf.convert_to_tensor(value=x)<EOL>original_shape = x.shape<EOL>new_shape = distribution_util.pad(<EOL>tf.shape(input=x), axis=<NUM_LIT:0>, back=True, value=<NUM_LIT:1>, count=ndims)<EOL>x = tf.reshape(x, new_shape)<EOL>x.set_shape(original_shape.concatenate([<NUM_LIT:1>]*ndims))<EOL>return x<EOL>

Maybe add `ndims` ones to `x.shape` on the right.

    If `ndims` is zero, this is a no-op; otherwise, we will create and return a
    new `Tensor` whose shape is that of `x` with `ndims` ones concatenated on the
    right side. If the shape of `x` is known statically, the shape of the return
    value will be as well.

    Args:
      x: The `Tensor` we'll return a reshaping of.
      ndims: Python `integer` number of ones to pad onto `x.shape`.
    Returns:
      If `ndims` is zero, `x`; otherwise, a `Tensor` whose shape is that of `x`
      with `ndims` ones concatenated on the right side. If possible, returns a
      `Tensor` whose shape is known statically.
    Raises:
      ValueError: if `ndims` is not a Python `integer` greater than or equal to
      zero.


x = tf.convert_to_tensor(value=x)<EOL>if x.shape.ndims is not None:<EOL><INDENT>axes = tf.range(x.shape.ndims - ndims, x.shape.ndims)<EOL><DEDENT>else:<EOL><INDENT>axes = tf.range(tf.rank(x) - ndims, tf.rank(x))<EOL><DEDENT>return tf.reduce_sum(input_tensor=x, axis=axes)<EOL>

Return `Tensor` with right-most ndims summed.

    Args:
      x: the `Tensor` whose right-most `ndims` dimensions to sum
      ndims: number of right-most dimensions to sum.

    Returns:
      A `Tensor` resulting from calling `reduce_sum` on the `ndims` right-most
      dimensions. If the shape of `x` is statically known, the result will also
      have statically known shape. Otherwise, the resulting shape will only be
      known at runtime.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if not x.dtype.is_floating:<EOL><INDENT>raise TypeError('<STR_LIT>')<EOL><DEDENT>def grad(grad_ys):<EOL><INDENT>large_float_like_x = np.sqrt(np.finfo(x.dtype.as_numpy_dtype()).max)<EOL>safe_grads = tf.where(<EOL>tf.equal(x, <NUM_LIT:0>), tf.fill(tf.shape(input=x), large_float_like_x),<EOL><NUM_LIT:0.5> * tf.math.rsqrt(x))<EOL>return grad_ys * safe_grads<EOL><DEDENT>return tf.sqrt(x), grad<EOL><DEDENT>

A sqrt function whose gradient at zero is very large but finite.

    Args:
      x: a `Tensor` whose sqrt is to be computed.
      name: a Python `str` prefixed to all ops created by this function.
        Default `None` (i.e., "sqrt_with_finite_grads").

    Returns:
      sqrt: the square root of `x`, with an overridden gradient at zero
      grad: a gradient function, which is the same as sqrt's gradient everywhere
        except at zero, where it is given a large finite value, instead of `inf`.

    Raises:
      TypeError: if `tf.convert_to_tensor(x)` is not a `float` type.

    Often in kernel functions, we need to compute the L2 norm of the difference
    between two vectors, `x` and `y`: `sqrt(sum_i((x_i - y_i) ** 2))`. In the
    case where `x` and `y` are identical, e.g., on the diagonal of a kernel
    matrix, we get `NaN`s when we take gradients with respect to the inputs. To
    see, this consider the forward pass:

      ```
      [x_1 ... x_N]  -->  [x_1 ** 2 ... x_N ** 2]  -->
          (x_1 ** 2 + ... + x_N ** 2)  -->  sqrt((x_1 ** 2 + ... + x_N ** 2))
      ```

    When we backprop through this forward pass, the `sqrt` yields an `inf` because
    `grad_z(sqrt(z)) = 1 / (2 * sqrt(z))`. Continuing the backprop to the left, at
    the `x ** 2` term, we pick up a `2 * x`, and when `x` is zero, we get
    `0 * inf`, which is `NaN`.

    We'd like to avoid these `NaN`s, since they infect the rest of the connected
    computation graph. Practically, when two inputs to a kernel function are
    equal, we are in one of two scenarios:
      1. We are actually computing k(x, x), in which case norm(x - x) is
         identically zero, independent of x. In this case, we'd like the
         gradient to reflect this independence: it should be zero.
      2. We are computing k(x, y), and x just *happens* to have the same value
         as y. The gradient at such inputs is in fact ill-defined (there is a
         cusp in the sqrt((x - y) ** 2) surface along the line x = y). There are,
         however, an infinite number of sub-gradients, all of which are valid at
         all such inputs. By symmetry, there is exactly one which is "special":
         zero, and we elect to use that value here. In practice, having two
         identical inputs to a kernel matrix is probably a pathological
         situation to be avoided, but that is better resolved at a higher level
         than this.

    To avoid the infinite gradient at zero, we use tf.custom_gradient to redefine
    the gradient at zero. We assign it to be a very large value, specifically
    the sqrt of the max value of the floating point dtype of the input. We use
    the sqrt (as opposed to just using the max floating point value) to avoid
    potential overflow when combining this value with others downstream.


<EOL>if all(a is None for a in arg_list):<EOL><INDENT>return None<EOL><DEDENT>return dtype_util.common_dtype(arg_list, tf.float32)<EOL>

Return common dtype of arg_list, or None.

    Args:
      arg_list: an iterable of items which are either `None` or have a `dtype`
        property.

    Returns:
      dtype: The common dtype of items in `arg_list`, or `None` if the list is
        empty or all items are `None`.


flattened = []<EOL>for k in kernels:<EOL><INDENT>if isinstance(k, _SumKernel):<EOL><INDENT>flattened += k.kernels<EOL><DEDENT>else:<EOL><INDENT>flattened.append(k)<EOL><DEDENT><DEDENT>return flattened<EOL>

Flatten a list of kernels which may contain _SumKernel instances.

    Args:
      kernels: Python list of `PositiveSemidefiniteKernel` instances

    Returns:
      Python list containing the elements of kernels, with any _SumKernel
      instances replaced by their `kernels` property contents.


flattened = []<EOL>for k in kernels:<EOL><INDENT>if isinstance(k, _ProductKernel):<EOL><INDENT>flattened += k.kernels<EOL><DEDENT>else:<EOL><INDENT>flattened.append(k)<EOL><DEDENT><DEDENT>return flattened<EOL>

Flatten a list of kernels which may contain _ProductKernel instances.

    Args:
      kernels: Python list of `PositiveSemidefiniteKernel` instances

    Returns:
      Python list containing the elements of kernels, with any _ProductKernel
      instances replaced by their `kernels` property contents.


def __init__(self, feature_ndims, dtype=None, name=None):

if not (isinstance(feature_ndims, int) and feature_ndims > <NUM_LIT:0>):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>' +<EOL>'<STR_LIT>'.format(feature_ndims))<EOL><DEDENT>self._feature_ndims = feature_ndims<EOL>self._dtype = dtype<EOL>if not name or name[-<NUM_LIT:1>] != '<STR_LIT:/>':  <EOL><INDENT>name = tf.compat.v1.name_scope(name or type(self).__name__).name<EOL><DEDENT>self._name = name<EOL>

Construct a PositiveSemidefiniteKernel (subclass) instance.

        Args:
          feature_ndims: Python `integer` indicating the number of dims (the rank)
            of the feature space this kernel acts on.
          dtype: `DType` on which this kernel operates.
          name: Python `str` name prefixed to Ops created by this class. Default:
            subclass name.

        Raises:
          ValueError: if `feature_ndims` is not an integer greater than 0
        Inputs to PositiveSemidefiniteKernel methods partition into 3 pieces:

        ```none
        [b1, ..., bB, e, f1, ..., fF]
        '----------'  |  '---------'
             |        |       '-- Feature dimensions
             |        '-- Example dimension (`matrix`-only)
             '-- Batch dimensions
        ```

        The `feature_ndims` argument declares how many of the right-most shape
        dimensions belong to the feature dimensions. This enables us to predict
        which shape dimensions will be 'reduced' away during kernel computation.


return self._feature_ndims<EOL>

The number of feature dimensions.

        Kernel functions generally act on pairs of inputs from some space like

        ```none
        R^(d1 x ... x  dD)
        ```

        or, in words: rank-`D` real-valued tensors of shape `[d1, ..., dD]`. Inputs
        can be vectors in some `R^N`, but are not restricted to be. Indeed, one
        might consider kernels over matrices, tensors, or even more general spaces,
        like strings or graphs.

        Returns:
          The number of feature dimensions (feature rank) of this kernel.


return self._dtype<EOL>


return self._name<EOL>

Name prepended to all ops created by this class.


return self._batch_shape()<EOL>

The batch_shape property of a PositiveSemidefiniteKernel.

        This property describes the fully broadcast shape of all kernel parameters.
        For example, consider an ExponentiatedQuadratic kernel, which is
        parameterized by an amplitude and length_scale:

        ```none
        exp_quad(x, x') := amplitude * exp(||x - x'||**2 / length_scale**2)
        ```

        The batch_shape of such a kernel is derived from broadcasting the shapes of
        `amplitude` and `length_scale`. E.g., if their shapes were

        ```python
        amplitude.shape = [2, 1, 1]
        length_scale.shape = [1, 4, 3]
        ```

        then `exp_quad`'s batch_shape would be `[2, 4, 3]`.

        Note that this property defers to the private _batch_shape method, which
        concrete implementation sub-classes are obliged to provide.

        Returns:
            `TensorShape` instance describing the fully broadcast shape of all
            kernel parameters.


The batch_shape property of a PositiveSemidefiniteKernel as a `Tensor`.

        Returns:
            `Tensor` which evaluates to a vector of integers which are the
            fully-broadcast shapes of the kernel parameters.


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

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

Helper function to standardize op scope.


with self._name_scope(self._name, values=[x1, x2]):<EOL><INDENT>x1 = tf.convert_to_tensor(value=x1, name='<STR_LIT>')<EOL>x2 = tf.convert_to_tensor(value=x2, name='<STR_LIT>')<EOL>return self._apply(x1, x2)<EOL><DEDENT>

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

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

        Returns:
          `Tensor` containing the (batch of) results of applying the kernel function
          to inputs `x1` and `x2`. If the kernel parameters' batch shape is
          `[k1, ..., kK]` then the shape of the `Tensor` resulting from this method
          call is `broadcast([b1, ..., bB], [c1, ..., cC], [k1, ..., kK])`.

        Given an index set `S`, a kernel function is mathematically defined as a
        real- or complex-valued function on `S` satisfying the
        positive semi-definiteness constraint:

        ```none
        sum_i sum_j (c[i]*) c[j] k(x[i], x[j]) >= 0
        ```

        for any finite collections `{x[1], ..., x[N]}` in `S` and
        `{c[1], ..., c[N]}` in the reals (or the complex plane). '*' is the complex
        conjugate, in the complex case.

        This method most closely resembles the function described in the
        mathematical definition of a kernel. Given a PositiveSemidefiniteKernel `k`
        with scalar parameters and inputs `x` and `y` in `S`, `apply(x, y)` yields a
        single scalar value. Given the same kernel and, say, batched inputs of shape
        `[b1, ..., bB, f1, ..., fF]`, it will yield a batch of scalars of shape
        `[b1, ..., bB]`.

        #### Examples

        ```python
        import tensorflow_probability as tfp

        # Suppose `SomeKernel` acts on vectors (rank-1 tensors)
        scalar_kernel = tfp.positive_semidefinite_kernels.SomeKernel(param=.5)
        scalar_kernel.batch_shape
        # ==> []

        # `x` and `y` are batches of five 3-D vectors:
        x = np.ones([5, 3], np.float32)
        y = np.ones([5, 3], np.float32)
        scalar_kernel.apply(x, y).shape
        # ==> [5]
        ```

        The above output is the result of vectorized computation of the five values

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

        Now we can consider a kernel with batched parameters:

        ```python
        batch_kernel = tfp.positive_semidefinite_kernels.SomeKernel(param=[.2, .5])
        batch_kernel.batch_shape
        # ==> [2]
        batch_kernel.apply(x, y).shape
        # ==> Error! [2] and [5] can't broadcast.
        ```

        The parameter batch shape of `[2]` and the input batch shape of `[5]` can't
        be broadcast together. We can fix this by giving the parameter a shape of
        `[2, 1]` which will correctly broadcast with `[5]` to yield `[2, 5]`:

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


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

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

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

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

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


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

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

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

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

        Given inputs `x1` and `x2` of shapes

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

        and

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

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

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

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

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

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

        #### Examples

        First, consider a kernel with a single scalar parameter.

        ```python
        import tensorflow_probability as tfp

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

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

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

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

        Now consider a kernel with batched parameters with the same inputs

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        ```

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


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

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

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

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


return self._kernels<EOL>

The list of kernels this _SumKernel sums over.


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

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

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

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


return self._kernels<EOL>

The list of kernels this _ProductKernel multiplies over.


with tf.compat.v1.name_scope(<EOL>name, values=[amplitude, scale_mixture_rate, length_scale]) as name:<EOL><INDENT>dtype = util.maybe_get_common_dtype(<EOL>[amplitude, scale_mixture_rate, length_scale])<EOL>if amplitude is not None:<EOL><INDENT>amplitude = tf.convert_to_tensor(<EOL>value=amplitude, name="<STR_LIT>", dtype=dtype)<EOL><DEDENT>self._amplitude = _validate_arg_if_not_none(<EOL>amplitude, tf.compat.v1.assert_positive, validate_args)<EOL>if scale_mixture_rate is not None:<EOL><INDENT>scale_mixture_rate = tf.convert_to_tensor(<EOL>value=scale_mixture_rate, name="<STR_LIT>", dtype=dtype)<EOL><DEDENT>self._scale_mixture_rate = _validate_arg_if_not_none(<EOL>scale_mixture_rate, tf.compat.v1.assert_positive, validate_args)<EOL>if length_scale is not None:<EOL><INDENT>length_scale = tf.convert_to_tensor(<EOL>value=length_scale, name="<STR_LIT>", dtype=dtype)<EOL><DEDENT>self._length_scale = _validate_arg_if_not_none(<EOL>length_scale, tf.compat.v1.assert_positive, validate_args)<EOL><DEDENT>super(RationalQuadratic, self).__init__(<EOL>feature_ndims, dtype=dtype, name=name)<EOL>

Construct a RationalQuadratic kernel instance.

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


return self._amplitude<EOL>


return self._length_scale<EOL>


return self._scale_mixture_rate<EOL>


with tf.compat.v1.name_scope(<EOL>name, values=[amplitude, period, length_scale]) as name:<EOL><INDENT>dtype = util.maybe_get_common_dtype(<EOL>[amplitude, period, length_scale])<EOL>if amplitude is not None:<EOL><INDENT>amplitude = tf.convert_to_tensor(<EOL>value=amplitude, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._amplitude = _validate_arg_if_not_none(<EOL>amplitude, tf.compat.v1.assert_positive, validate_args)<EOL>if period is not None:<EOL><INDENT>period = tf.convert_to_tensor(value=period, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._period = _validate_arg_if_not_none(<EOL>period, tf.compat.v1.assert_positive, validate_args)<EOL>if length_scale is not None:<EOL><INDENT>length_scale = tf.convert_to_tensor(<EOL>value=length_scale, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._length_scale = _validate_arg_if_not_none(<EOL>length_scale, tf.compat.v1.assert_positive, validate_args)<EOL><DEDENT>super(ExpSinSquared, self).__init__(feature_ndims, dtype=dtype, name=name)<EOL>

Construct a ExpSinSquared kernel instance.

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


return self._period<EOL>


def _init_params(self, amplitude, length_scale, validate_args):

dtype = util.maybe_get_common_dtype(<EOL>[amplitude, length_scale])<EOL>if amplitude is not None:<EOL><INDENT>amplitude = tf.convert_to_tensor(<EOL>value=amplitude, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._amplitude = _validate_arg_if_not_none(<EOL>amplitude, tf.compat.v1.assert_positive, validate_args)<EOL>if length_scale is not None:<EOL><INDENT>length_scale = tf.convert_to_tensor(<EOL>value=length_scale, name='<STR_LIT>', dtype=dtype)<EOL><DEDENT>self._length_scale = _validate_arg_if_not_none(<EOL>length_scale, tf.compat.v1.assert_positive, validate_args)<EOL>return dtype<EOL>

Shared init logic for `amplitude` and `length_scale` params.

        Args:
          amplitude: `Tensor` (or convertible) or `None` to convert, validate.
          length_scale: `Tensor` (or convertible) or `None` to convert, validate.
          validate_args: If `True`, parameters are checked for validity despite
            possibly degrading runtime performance

        Returns:
          dtype: The common `DType` of the parameters.


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

Construct a MaternOneHalf kernel instance.

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


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

Construct a MaternThreeHalves kernel instance.

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


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

Construct a MaternFiveHalves kernel instance.

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