with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>_, observed_stddev, observed_initial = (<EOL>sts_util.empirical_statistics(observed_time_series)<EOL>if observed_time_series is not None else (<NUM_LIT:0.>, <NUM_LIT:1.>, <NUM_LIT:0.>))<EOL>if drift_scale_prior is None:<EOL><INDENT>drift_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev), scale=<NUM_LIT>)<EOL><DEDENT>if initial_effect_prior is None:<EOL><INDENT>initial_effect_prior = tfd.Normal(<EOL>loc=observed_initial,<EOL>scale=tf.abs(observed_initial) + observed_stddev)<EOL><DEDENT>dtype = tf.debugging.assert_same_float_dtype(<EOL>[drift_scale_prior, initial_effect_prior])<EOL>if isinstance(initial_effect_prior, tfd.Normal):<EOL><INDENT>initial_state_prior = tfd.MultivariateNormalDiag(<EOL>loc=tf.stack([initial_effect_prior.mean()] * num_seasons, axis=-<NUM_LIT:1>),<EOL>scale_diag=tf.stack([initial_effect_prior.stddev()] * num_seasons,<EOL>axis=-<NUM_LIT:1>))<EOL><DEDENT>else:<EOL><INDENT>initial_state_prior = initial_effect_prior<EOL><DEDENT>if constrain_mean_effect_to_zero:<EOL><INDENT>(effects_to_residuals, _) = build_effects_to_residuals_matrix(<EOL>num_seasons, dtype=dtype)<EOL>effects_to_residuals_linop = tf.linalg.LinearOperatorFullMatrix(<EOL>effects_to_residuals)  <EOL>initial_state_prior_loc = effects_to_residuals_linop.matvec(<EOL>initial_state_prior.mean())<EOL>initial_state_prior_scale_linop = effects_to_residuals_linop.matmul(<EOL>initial_state_prior.scale)  <EOL>initial_state_prior = tfd.MultivariateNormalFullCovariance(<EOL>loc=initial_state_prior_loc,<EOL>covariance_matrix=initial_state_prior_scale_linop.matmul(<EOL>initial_state_prior_scale_linop.to_dense(), adjoint_arg=True))<EOL><DEDENT>self._constrain_mean_effect_to_zero = constrain_mean_effect_to_zero<EOL>self._initial_state_prior = initial_state_prior<EOL>self._num_seasons = num_seasons<EOL>self._num_steps_per_season = num_steps_per_season<EOL>super(Seasonal, self).__init__(<EOL>parameters=[<EOL>Parameter('<STR_LIT>', drift_scale_prior, tfb.Softplus()),<EOL>],<EOL>latent_size=(num_seasons - <NUM_LIT:1><EOL>if self.constrain_mean_effect_to_zero else num_seasons),<EOL>name=name)<EOL><DEDENT>

Specify a seasonal effects model.

        Args:
          num_seasons: Scalar Python `int` number of seasons.
          num_steps_per_season: Python `int` number of steps in each
            season. This may be either a scalar (shape `[]`), in which case all
            seasons have the same length, or a NumPy array of shape `[num_seasons]`,
            in which seasons have different length, but remain constant around
            different cycles, or a NumPy array of shape `[num_cycles, num_seasons]`,
            in which num_steps_per_season for each season also varies in different
            cycle (e.g., a 4 years cycle with leap day).
            Default value: 1.
          drift_scale_prior: optional `tfd.Distribution` instance specifying a prior
            on the `drift_scale` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          initial_effect_prior: optional `tfd.Distribution` instance specifying a
            normal prior on the initial effect of each season. This may be either
            a scalar `tfd.Normal` prior, in which case it applies independently to
            every season, or it may be multivariate normal (e.g.,
            `tfd.MultivariateNormalDiag`) with event shape `[num_seasons]`, in
            which case it specifies a joint prior across all seasons. If `None`, a
            heuristic default prior is constructed based on the provided
            `observed_time_series`.
            Default value: `None`.
          constrain_mean_effect_to_zero: if `True`, use a model parameterization
            that constrains the mean effect across all seasons to be zero. This
            constraint is generally helpful in identifying the contributions of
            different model components and can lead to more interpretable
            posterior decompositions. It may be undesirable if you plan to directly
            examine the latent space of the underlying state space model.
            Default value: `True`.
          observed_time_series: optional `float` `Tensor` of shape
            `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
            supported when `T > 1`), specifying an observed time series.
            Any priors not explicitly set will be given default values according to
            the scale of the observed time series (or batch of time series). May
            optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
            a mask `Tensor` to specify timesteps with missing observations.
            Default value: `None`.
          name: the name of this model component.
            Default value: 'Seasonal'.


return self._constrain_mean_effect_to_zero<EOL>

Whether to constrain the mean effect to zero.


return self._num_seasons<EOL>


return self._num_steps_per_season<EOL>


return self._initial_state_prior<EOL>

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


def __init__(self, parameters, latent_size, name='<STR_LIT>'):

self._parameters = parameters<EOL>self._latent_size = latent_size<EOL>self._name = name<EOL>

Construct a specification for a structural time series model.

        Args:
          parameters: list of Parameter namedtuples, each specifying the name
            and prior distribution of a model parameter along with a
            bijective transformation from an unconstrained space to the support
            of that parameter. The order of this list determines the canonical
            parameter ordering used by fitting and inference algorithms.
          latent_size: Python `int` specifying the dimensionality of the latent
            state space for this model.
          name: Python `str` name for this model component.


return self._parameters<EOL>

List of Parameter(name, prior, bijector) namedtuples for this model.


return self._latent_size<EOL>

Python `int` dimensionality of the latent space in this model.


return self._name<EOL>


batch_shape = tf.TensorShape([])<EOL>for param in self.parameters:<EOL><INDENT>batch_shape = tf.broadcast_static_shape(<EOL>batch_shape, param.prior.batch_shape)<EOL><DEDENT>return batch_shape<EOL>

Static batch shape of models represented by this component.

        Returns:
          batch_shape: A `tf.TensorShape` giving the broadcast batch shape of
            all model parameters. This should match the batch shape of
            derived state space models, i.e.,
            `self.make_state_space_model(...).batch_shape`. It may be partially
            defined or unknown.


batch_shape = tf.constant([], dtype=tf.int32)<EOL>for param in self.parameters:<EOL><INDENT>batch_shape = tf.broadcast_dynamic_shape(<EOL>batch_shape, param.prior.batch_shape_tensor())<EOL><DEDENT>return batch_shape<EOL>

Runtime batch shape of models represented by this component.

        Returns:
          batch_shape: `int` `Tensor` giving the broadcast batch shape of
            all model parameters. This should match the batch shape of
            derived state space models, i.e.,
            `self.make_state_space_model(...).batch_shape_tensor()`.


if hasattr(param_vals, '<STR_LIT>'):<EOL><INDENT>param_map = param_vals<EOL><DEDENT>else:<EOL><INDENT>param_map = {p.name: v for (p, v) in zip(self.parameters, param_vals)}<EOL><DEDENT>return param_map<EOL>

If given an ordered list of parameter values, build a name:value map.

        This is a utility method that allows parameter values to be specified as
        either lists or dicts, by transforming lists to a canonical dict
        representation.

        Args:
          param_vals: Python list (or other `iterable`) of `Tensor` values
            corresponding to the parameters listed in `self.parameters`,
            OR a map (Python `dict`) of parameter names to values.

        Returns:
          param_map: Python `dict` mapping from the names given in `self.parameters`
            to the specified parameter values.


Instantiate this model as a Distribution over specified `num_timesteps`.

        Args:
          num_timesteps: Python `int` number of timesteps to model.
          param_vals: a list of `Tensor` parameter values in order corresponding to
            `self.parameters`, or a dict mapping from parameter names to values.
          initial_state_prior: an optional `Distribution` instance overriding the
            default prior on the model's initial state. This is used in forecasting
            ("today's prior is yesterday's posterior").
          initial_step: optional `int` specifying the initial timestep to model.
            This is relevant when the model contains time-varying components,
            e.g., holidays or seasonality.

        Returns:
          dist: a `LinearGaussianStateSpaceModel` Distribution object.


seed = distributions.SeedStream(<EOL>seed, salt='<STR_LIT>')<EOL>with tf.compat.v1.name_scope(<EOL>'<STR_LIT>',<EOL>values=[num_timesteps, params_sample_shape, trajectories_sample_shape]):<EOL><INDENT>param_samples = [<EOL>p.prior.sample(params_sample_shape, seed=seed(), name=p.name)<EOL>for p in self.parameters<EOL>]<EOL>model = self.make_state_space_model(<EOL>num_timesteps=num_timesteps,<EOL>initial_step=initial_step,<EOL>param_vals=param_samples)<EOL>return model.sample(trajectories_sample_shape, seed=seed()), param_samples<EOL><DEDENT>

Sample from the joint prior over model parameters and trajectories.

        Args:
          num_timesteps: Scalar `int` `Tensor` number of timesteps to model.
          initial_step: Optional scalar `int` `Tensor` specifying the starting
            timestep.
              Default value: 0.
          params_sample_shape: Number of possible worlds to sample iid from the
            parameter prior, or more generally, `Tensor` `int` shape to fill with
            iid samples.
              Default value: [] (i.e., draw a single sample and don't expand the
              shape).
          trajectories_sample_shape: For each sampled set of parameters, number
            of trajectories to sample, or more generally, `Tensor` `int` shape to
            fill with iid samples.
            Default value: [] (i.e., draw a single sample and don't expand the
              shape).
          seed: Python `int` random seed.

        Returns:
          trajectories: `float` `Tensor` of shape
            `trajectories_sample_shape + params_sample_shape + [num_timesteps, 1]`
            containing all sampled trajectories.
          param_samples: list of sampled parameter value `Tensor`s, in order
            corresponding to `self.parameters`, each of shape
            `params_sample_shape + prior.batch_shape + prior.event_shape`.


with tf.compat.v1.name_scope(<EOL>'<STR_LIT>', values=[observed_time_series]):<EOL><INDENT>[<EOL>observed_time_series,<EOL>mask<EOL>] = sts_util.canonicalize_observed_time_series_with_mask(<EOL>observed_time_series)<EOL>num_timesteps = distribution_util.prefer_static_value(<EOL>tf.shape(input=observed_time_series))[-<NUM_LIT:2>]<EOL>def log_joint_fn(*param_vals):<EOL><INDENT>"""<STR_LIT>"""<EOL>param_lp = sum([<EOL>param.prior.log_prob(param_val)<EOL>for (param, param_val) in zip(self.parameters, param_vals)<EOL>])<EOL>lgssm = self.make_state_space_model(<EOL>param_vals=param_vals, num_timesteps=num_timesteps)<EOL>observation_lp = lgssm.log_prob(observed_time_series, mask=mask)<EOL>sample_ndims = tf.maximum(<NUM_LIT:0>,<EOL>tf.rank(observation_lp) - tf.rank(param_lp))<EOL>observation_lp = tf.reduce_sum(<EOL>input_tensor=observation_lp, axis=tf.range(sample_ndims))<EOL>return param_lp + observation_lp<EOL><DEDENT><DEDENT>return log_joint_fn<EOL>

Build the joint density `log p(params) + log p(y|params)` as a callable.

        Args:
          observed_time_series: Observed `Tensor` trajectories of shape
            `sample_shape + batch_shape + [num_timesteps, 1]` (the trailing
            `1` dimension is optional if `num_timesteps > 1`), where
            `batch_shape` should match `self.batch_shape` (the broadcast batch
            shape of all priors on parameters for this structural time series
            model). May optionally be an instance of `tfp.sts.MaskedTimeSeries`,
            which includes a mask `Tensor` to specify timesteps with missing
            observations.

        Returns:
         log_joint_fn: A function taking a `Tensor` argument for each model
           parameter, in canonical order, and returning a `Tensor` log probability
           of shape `batch_shape`. Note that, *unlike* `tfp.Distributions`
           `log_prob` methods, the `log_joint` sums over the `sample_shape` from y,
           so that `sample_shape` does not appear in the output log_prob. This
           corresponds to viewing multiple samples in `y` as iid observations from a
           single model, which is typically the desired behavior for parameter
           inference.


start_dim = <NUM_LIT:0><EOL>marginals = []<EOL>for size in block_sizes:<EOL><INDENT>end_dim = start_dim + size<EOL>marginals.append(covariance[..., start_dim:end_dim, start_dim:end_dim])<EOL>start_dim = end_dim<EOL><DEDENT>return marginals<EOL>

Split a covariance matrix into block-diagonal marginals of given sizes.


def _decompose_from_posterior_marginals(<EOL>model, posterior_means, posterior_covs, parameter_samples):

try:<EOL><INDENT>model.components<EOL><DEDENT>except AttributeError:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(model))<EOL><DEDENT>with tf.compat.v1.name_scope('<STR_LIT>'):<EOL><INDENT>latent_sizes = [component.latent_size for component in model.components]<EOL>component_means = tf.split(posterior_means, latent_sizes, axis=-<NUM_LIT:1>)<EOL>component_covs = _split_covariance_into_marginals(<EOL>posterior_covs, latent_sizes)<EOL>num_timesteps = dist_util.prefer_static_value(<EOL>tf.shape(input=posterior_means))[-<NUM_LIT:2>]<EOL>component_ssms = model.make_component_state_space_models(<EOL>num_timesteps=num_timesteps,<EOL>param_vals=parameter_samples)<EOL>component_predictive_dists = collections.OrderedDict()<EOL>for (component, component_ssm,<EOL>component_mean, component_cov) in zip(model.components, component_ssms,<EOL>component_means, component_covs):<EOL><INDENT>component_obs_mean, component_obs_cov = (<EOL>component_ssm.latents_to_observations(<EOL>latent_means=component_mean,<EOL>latent_covs=component_cov))<EOL>component_predictive_dists[component] = sts_util.mix_over_posterior_draws(<EOL>means=component_obs_mean[..., <NUM_LIT:0>],<EOL>variances=component_obs_cov[..., <NUM_LIT:0>, <NUM_LIT:0>])<EOL><DEDENT><DEDENT>return component_predictive_dists<EOL>

Utility method to decompose a joint posterior into components.

    Args:
      model: `tfp.sts.Sum` instance defining an additive STS model.
      posterior_means: float `Tensor` of shape `concat(
        [[num_posterior_draws], batch_shape, num_timesteps, latent_size])`
        representing the posterior mean over latents in an
        `AdditiveStateSpaceModel`.
      posterior_covs: float `Tensor` of shape `concat(
        [[num_posterior_draws], batch_shape, num_timesteps,
        latent_size, latent_size])`
        representing the posterior marginal covariances over latents in an
        `AdditiveStateSpaceModel`.
      parameter_samples: Python `list` of `Tensors` representing posterior
        samples of model parameters, with shapes `[concat([
        [num_posterior_draws], param.prior.batch_shape,
        param.prior.event_shape]) for param in model.parameters]`. This may
        optionally also be a map (Python `dict`) of parameter names to
        `Tensor` values.

    Returns:
      component_dists: A `collections.OrderedDict` instance mapping
        component StructuralTimeSeries instances (elements of `model.components`)
        to `tfd.Distribution` instances representing the posterior marginal
        distributions on the process modeled by each component. Each distribution
        has batch shape matching that of `posterior_means`/`posterior_covs`, and
        event shape of `[num_timesteps]`.


def decompose_by_component(model, observed_time_series, parameter_samples):

with tf.compat.v1.name_scope('<STR_LIT>',<EOL>values=[observed_time_series]):<EOL><INDENT>[<EOL>observed_time_series,<EOL>is_missing<EOL>] = sts_util.canonicalize_observed_time_series_with_mask(<EOL>observed_time_series)<EOL>num_timesteps = dist_util.prefer_static_value(<EOL>tf.shape(input=observed_time_series))[-<NUM_LIT:2>]<EOL>ssm = model.make_state_space_model(num_timesteps=num_timesteps,<EOL>param_vals=parameter_samples)<EOL>posterior_means, posterior_covs = ssm.posterior_marginals(<EOL>observed_time_series, mask=is_missing)<EOL>return _decompose_from_posterior_marginals(<EOL>model, posterior_means, posterior_covs, parameter_samples)<EOL><DEDENT>

Decompose an observed time series into contributions from each component.

    This method decomposes a time series according to the posterior represention
    of a structural time series model. In particular, it:
      - Computes the posterior marginal mean and covariances over the additive
        model's latent space.
      - Decomposes the latent posterior into the marginal blocks for each
        model component.
      - Maps the per-component latent posteriors back through each component's
        observation model, to generate the time series modeled by that component.

    Args:
      model: An instance of `tfp.sts.Sum` representing a structural time series
        model.
      observed_time_series: `float` `Tensor` of shape
        `batch_shape + [num_timesteps, 1]` (omitting the trailing unit dimension
        is also supported when `num_timesteps > 1`), specifying an observed time
        series. May optionally be an instance of `tfp.sts.MaskedTimeSeries`, which
        includes a mask `Tensor` to specify timesteps with missing observations.
      parameter_samples: Python `list` of `Tensors` representing posterior
        samples of model parameters, with shapes `[concat([
        [num_posterior_draws], param.prior.batch_shape,
        param.prior.event_shape]) for param in model.parameters]`. This may
        optionally also be a map (Python `dict`) of parameter names to
        `Tensor` values.
    Returns:
      component_dists: A `collections.OrderedDict` instance mapping
        component StructuralTimeSeries instances (elements of `model.components`)
        to `tfd.Distribution` instances representing the posterior marginal
        distributions on the process modeled by each component. Each distribution
        has batch shape matching that of `posterior_means`/`posterior_covs`, and
        event shape of `[num_timesteps]`.

    #### Examples

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

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

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

    To extract the contributions of individual components, pass the time series
    and sampled parameters into `decompose_by_component`:

    ```python
      component_dists = decompose_by_component(
        model,
        observed_time_series=observed_time_series,
        parameter_samples=samples)

      # Component mean and stddev have shape `[len(observed_time_series)]`.
      day_of_week_effect_mean = component_dists[day_of_week].mean()
      day_of_week_effect_stddev = component_dists[day_of_week].stddev()
    ```

    Using the component distributions, we can visualize the uncertainty for
    each component:

    ```
    from matplotlib import pylab as plt
    num_components = len(component_dists)
    xs = np.arange(len(observed_time_series))
    fig = plt.figure(figsize=(12, 3 * num_components))
    for i, (component, component_dist) in enumerate(component_dists.items()):

      # If in graph mode, replace `.numpy()` with `.eval()` or `sess.run()`.
      component_mean = component_dist.mean().numpy()
      component_stddev = component_dist.stddev().numpy()

      ax = fig.add_subplot(num_components, 1, 1 + i)
      ax.plot(xs, component_mean, lw=2)
      ax.fill_between(xs,
                      component_mean - 2 * component_stddev,
                      component_mean + 2 * component_stddev,
                      alpha=0.5)
      ax.set_title(component.name)
    ```


def decompose_forecast_by_component(model, forecast_dist, parameter_samples):

with tf.compat.v1.name_scope('<STR_LIT>'):<EOL><INDENT>try:<EOL><INDENT>forecast_lgssm = forecast_dist.components_distribution<EOL>forecast_latent_mean, _ = forecast_lgssm._joint_mean()  <EOL>forecast_latent_covs, _ = forecast_lgssm._joint_covariances()  <EOL><DEDENT>except AttributeError as e:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(e))<EOL><DEDENT>forecast_latent_mean = dist_util.move_dimension(<EOL>forecast_latent_mean, source_idx=-<NUM_LIT:3>, dest_idx=<NUM_LIT:0>)<EOL>forecast_latent_covs = dist_util.move_dimension(<EOL>forecast_latent_covs, source_idx=-<NUM_LIT:4>, dest_idx=<NUM_LIT:0>)<EOL>return _decompose_from_posterior_marginals(<EOL>model, forecast_latent_mean, forecast_latent_covs, parameter_samples)<EOL><DEDENT>

Decompose a forecast distribution into contributions from each component.

    Args:
      model: An instance of `tfp.sts.Sum` representing a structural time series
        model.
      forecast_dist: A `Distribution` instance returned by `tfp.sts.forecast()`.
        (specifically, must be a `tfd.MixtureSameFamily` over a
        `tfd.LinearGaussianStateSpaceModel` parameterized by posterior samples).
      parameter_samples: Python `list` of `Tensors` representing posterior samples
        of model parameters, with shapes `[concat([[num_posterior_draws],
        param.prior.batch_shape, param.prior.event_shape]) for param in
        model.parameters]`. This may optionally also be a map (Python `dict`) of
        parameter names to `Tensor` values.
    Returns:
      component_forecasts: A `collections.OrderedDict` instance mapping
        component StructuralTimeSeries instances (elements of `model.components`)
        to `tfd.Distribution` instances representing the marginal forecast for
        each component. Each distribution has batch and event shape matching
        `forecast_dist` (specifically, the event shape is
        `[num_steps_forecast]`).

    #### Examples

    Suppose we've built a model, fit it to data, and constructed a forecast
    distribution:

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

      num_steps_forecast = 50
      samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
      forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                   parameter_samples=samples,
                                   num_steps_forecast=num_steps_forecast)
    ```

    To extract the forecast for individual components, pass the forecast
    distribution into `decompose_forecast_by_components`:

    ```python
      component_forecasts = decompose_forecast_by_component(
        model, forecast_dist, samples)

      # Component mean and stddev have shape `[num_steps_forecast]`.
      day_of_week_effect_mean = forecast_components[day_of_week].mean()
      day_of_week_effect_stddev = forecast_components[day_of_week].stddev()
    ```

    Using the component forecasts, we can visualize the uncertainty for each
    component:

    ```
    from matplotlib import pylab as plt
    num_components = len(component_forecasts)
    xs = np.arange(num_steps_forecast)
    fig = plt.figure(figsize=(12, 3 * num_components))
    for i, (component, component_dist) in enumerate(component_forecasts.items()):

      # If in graph mode, replace `.numpy()` with `.eval()` or `sess.run()`.
      component_mean = component_dist.mean().numpy()
      component_stddev = component_dist.stddev().numpy()

      ax = fig.add_subplot(num_components, 1, 1 + i)
      ax.plot(xs, component_mean, lw=2)
      ax.fill_between(xs,
                      component_mean - 2 * component_stddev,
                      component_mean + 2 * component_stddev,
                      alpha=0.5)
      ax.set_title(component.name)
    ```


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>values=[drift_scale]) as name:<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[design_matrix, drift_scale, initial_state_prior])<EOL>design_matrix = tf.convert_to_tensor(<EOL>value=design_matrix, name='<STR_LIT>', dtype=dtype)<EOL>drift_scale = tf.convert_to_tensor(<EOL>value=drift_scale, name='<STR_LIT>', dtype=dtype)<EOL>observation_noise_scale = tf.convert_to_tensor(<EOL>value=observation_noise_scale,<EOL>name='<STR_LIT>',<EOL>dtype=dtype)<EOL>num_features = prefer_static.shape(design_matrix)[-<NUM_LIT:1>]<EOL>def observation_matrix_fn(t):<EOL><INDENT>observation_matrix = tf.linalg.LinearOperatorFullMatrix(<EOL>design_matrix[..., t, tf.newaxis, :], name='<STR_LIT>')<EOL>return observation_matrix<EOL><DEDENT>self._drift_scale = drift_scale<EOL>self._observation_noise_scale = observation_noise_scale<EOL>super(DynamicLinearRegressionStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=tf.linalg.LinearOperatorIdentity(<EOL>num_rows=num_features,<EOL>dtype=dtype,<EOL>name='<STR_LIT>'),<EOL>transition_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=(drift_scale[..., tf.newaxis] *<EOL>tf.ones([num_features], dtype=dtype)),<EOL>name='<STR_LIT>'),<EOL>observation_matrix=observation_matrix_fn,<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=observation_noise_scale[..., tf.newaxis],<EOL>name='<STR_LIT>'),<EOL>initial_state_prior=initial_state_prior,<EOL>initial_step=initial_step,<EOL>allow_nan_stats=allow_nan_stats,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

State space model for a dynamic linear regression.

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


return self._drift_scale<EOL>

Standard deviation of the drift in weights at each timestep.


return self._observation_noise_scale<EOL>

Standard deviation of the observation noise.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[design_matrix, drift_scale_prior, initial_weights_prior])<EOL>num_features = prefer_static.shape(design_matrix)[-<NUM_LIT:1>]<EOL>if initial_weights_prior is None:<EOL><INDENT>initial_weights_prior = tfd.MultivariateNormalDiag(<EOL>scale_diag=<NUM_LIT> * tf.ones([num_features], dtype=dtype))<EOL><DEDENT>if drift_scale_prior is None:<EOL><INDENT>if observed_time_series is None:<EOL><INDENT>observed_stddev = tf.constant(<NUM_LIT:1.0>, dtype=dtype)<EOL><DEDENT>else:<EOL><INDENT>_, observed_stddev, _ = sts_util.empirical_statistics(<EOL>observed_time_series)<EOL><DEDENT>drift_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev),<EOL>scale=<NUM_LIT>,<EOL>name='<STR_LIT>')<EOL><DEDENT>self._initial_state_prior = initial_weights_prior<EOL>self._design_matrix = design_matrix<EOL>super(DynamicLinearRegression, self).__init__(<EOL>parameters=[<EOL>Parameter('<STR_LIT>', drift_scale_prior, tfb.Softplus())<EOL>],<EOL>latent_size=num_features,<EOL>name=name)<EOL><DEDENT>

Specify a dynamic linear regression.

        Args:
          design_matrix: float `Tensor` of shape `concat([batch_shape,
            [num_timesteps, num_features]])`.
          drift_scale_prior: instance of `tfd.Distribution` specifying a prior on
            the `drift_scale` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          initial_weights_prior: instance of `tfd.MultivariateNormal` representing
            the prior distribution on the latent states (the regression weights).
            Must have event shape `[num_features]`. If `None`, a weakly-informative
            Normal(0., 10.) prior is used.
            Default value: `None`.
          observed_time_series: `float` `Tensor` of shape `batch_shape + [T, 1]`
            (omitting the trailing unit dimension is also supported when `T > 1`),
            specifying an observed time series. Any priors not explicitly set will
            be given default values according to the scale of the observed time
            series (or batch of time series). May optionally be an instance of
            `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify
            timesteps with missing observations.
            Default value: `None`.
          name: Python `str` for the name of this component.
            Default value: 'DynamicLinearRegression'.


return self._initial_state_prior<EOL>

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


return self._design_matrix<EOL>


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

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

Convert a numpy array to a TF placeholder.

        Args:
          ndarray: any object convertible to a numpy array via `np.asarray()`.
          dtype: numpy `dtype` of the returned placeholder. If `None`, uses
            the default `self.dtype`.

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


def expectation(f, samples, log_prob=None, use_reparametrization=True,<EOL>axis=<NUM_LIT:0>, keep_dims=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [samples]):<EOL><INDENT>if not callable(f):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>if use_reparametrization:<EOL><INDENT>return tf.reduce_mean(<EOL>input_tensor=f(samples), axis=axis, keepdims=keep_dims)<EOL><DEDENT>else:<EOL><INDENT>if not callable(log_prob):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>stop = tf.stop_gradient  <EOL>x = stop(samples)<EOL>logpx = log_prob(x)<EOL>fx = f(x)  <EOL>dice = fx * tf.exp(logpx - stop(logpx))<EOL>return tf.reduce_mean(input_tensor=dice, axis=axis, keepdims=keep_dims)<EOL><DEDENT><DEDENT>

Computes the Monte-Carlo approximation of `E_p[f(X)]`.

    This function computes the Monte-Carlo approximation of an expectation, i.e.,

    ```none
    E_p[f(X)] approx= m**-1 sum_i^m f(x_j),  x_j ~iid p(X)
    ```

    where:

    - `x_j = samples[j, ...]`,
    - `log(p(samples)) = log_prob(samples)` and
    - `m = prod(shape(samples)[axis])`.

    Tricks: Reparameterization and Score-Gradient

    When p is "reparameterized", i.e., a diffeomorphic transformation of a
    parameterless distribution (e.g.,
    `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and
    expectation, i.e.,
    `grad[ Avg{ s_i : i=1...n } ] = Avg{ grad[s_i] : i=1...n }` where
    `S_n = Avg{s_i}` and `s_i = f(x_i), x_i ~ p`.

    However, if p is not reparameterized, TensorFlow's gradient will be incorrect
    since the chain-rule stops at samples of non-reparameterized distributions.
    (The non-differentiated result, `approx_expectation`, is the same regardless
    of `use_reparametrization`.) In this circumstance using the Score-Gradient
    trick results in an unbiased gradient, i.e.,

    ```none
    grad[ E_p[f(X)] ]
    = grad[ int dx p(x) f(x) ]
    = int dx grad[ p(x) f(x) ]
    = int dx [ p'(x) f(x) + p(x) f'(x) ]
    = int dx p(x) [p'(x) / p(x) f(x) + f'(x) ]
    = int dx p(x) grad[ f(x) p(x) / stop_grad[p(x)] ]
    = E_p[ grad[ f(x) p(x) / stop_grad[p(x)] ] ]
    ```

    Unless p is not reparametrized, it is usually preferable to
    `use_reparametrization = True`.

    Warning: users are responsible for verifying `p` is a "reparameterized"
    distribution.

    Example Use:

    ```python
    # Monte-Carlo approximation of a reparameterized distribution, e.g., Normal.

    num_draws = int(1e5)
    p = tfp.distributions.Normal(loc=0., scale=1.)
    q = tfp.distributions.Normal(loc=1., scale=2.)
    exact_kl_normal_normal = tfp.distributions.kl_divergence(p, q)
    # ==> 0.44314718
    approx_kl_normal_normal = tfp.monte_carlo.expectation(
        f=lambda x: p.log_prob(x) - q.log_prob(x),
        samples=p.sample(num_draws, seed=42),
        log_prob=p.log_prob,
        use_reparametrization=(p.reparameterization_type
                               == tfp.distributions.FULLY_REPARAMETERIZED))
    # ==> 0.44632751
    # Relative Error: <1%

    # Monte-Carlo approximation of non-reparameterized distribution,
    # e.g., Bernoulli.

    num_draws = int(1e5)
    p = tfp.distributions.Bernoulli(probs=0.4)
    q = tfp.distributions.Bernoulli(probs=0.8)
    exact_kl_bernoulli_bernoulli = tfp.distributions.kl_divergence(p, q)
    # ==> 0.38190854
    approx_kl_bernoulli_bernoulli = tfp.monte_carlo.expectation(
        f=lambda x: p.log_prob(x) - q.log_prob(x),
        samples=p.sample(num_draws, seed=42),
        log_prob=p.log_prob,
        use_reparametrization=(p.reparameterization_type
                               == tfp.distributions.FULLY_REPARAMETERIZED))
    # ==> 0.38336259
    # Relative Error: <1%

    # For comparing the gradients, see `expectation_test.py`.
    ```

    Note: The above example is for illustration only. To compute approximate
    KL-divergence, the following is preferred:

    ```python
    approx_kl_p_q = bf.monte_carlo_csiszar_f_divergence(
        f=bf.kl_reverse,
        p_log_prob=q.log_prob,
        q=p,
        num_draws=num_draws)
    ```

    Args:
      f: Python callable which can return `f(samples)`.
      samples: `Tensor` of samples used to form the Monte-Carlo approximation of
        `E_p[f(X)]`.  A batch of samples should be indexed by `axis` dimensions.
      log_prob: Python callable which can return `log_prob(samples)`. Must
        correspond to the natural-logarithm of the pdf/pmf of each sample. Only
        required/used if `use_reparametrization=False`.
        Default value: `None`.
      use_reparametrization: Python `bool` indicating that the approximation
        should use the fact that the gradient of samples is unbiased. Whether
        `True` or `False`, this arg only affects the gradient of the resulting
        `approx_expectation`.
        Default value: `True`.
      axis: The dimensions to average. If `None`, averages all
        dimensions.
        Default value: `0` (the left-most dimension).
      keep_dims: If True, retains averaged dimensions using size `1`.
        Default value: `False`.
      name: A `name_scope` for operations created by this function.
        Default value: `None` (which implies "expectation").

    Returns:
      approx_expectation: `Tensor` corresponding to the Monte-Carlo approximation
        of `E_p[f(X)]`.

    Raises:
      ValueError: if `f` is not a Python `callable`.
      ValueError: if `use_reparametrization=False` and `log_prob` is not a Python
        `callable`.


return tf.reduce_mean(input_tensor=values, axis=[<NUM_LIT:0>])<EOL>

Mean over sample indices.  In this module this is always [0].


return tf.reduce_max(input_tensor=values, axis=[<NUM_LIT:0>])<EOL>

Max over sample indices.  In this module this is always [0].


with tf.compat.v1.name_scope('<STR_LIT>', values=[z, n]):<EOL><INDENT>if (n is None) == (z is None):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>' % (n, z))<EOL><DEDENT>if n is not None:<EOL><INDENT>return dist.sample(n, seed=seed)<EOL><DEDENT>else:<EOL><INDENT>return tf.convert_to_tensor(value=z, name='<STR_LIT:z>')<EOL><DEDENT><DEDENT>


states_was_list = _is_list_like(states)<EOL>if not states_was_list:<EOL><INDENT>states = [states]<EOL><DEDENT>filter_beyond_lag = _broadcast_maybelist_arg(states, filter_beyond_lag,<EOL>'<STR_LIT>')<EOL>filter_threshold = _broadcast_maybelist_arg(states, filter_threshold,<EOL>'<STR_LIT>')<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>'):<EOL><INDENT>ess_list = [<EOL>_effective_sample_size_single_state(s, ml, mlt)<EOL>for (s, ml, mlt) in zip(states, filter_beyond_lag, filter_threshold)<EOL>]<EOL><DEDENT>if states_was_list:<EOL><INDENT>return ess_list<EOL><DEDENT>return ess_list[<NUM_LIT:0>]<EOL>

Estimate a lower bound on effective sample size for each independent chain.

    Roughly speaking, "effective sample size" (ESS) is the size of an iid sample
    with the same variance as `state`.

    More precisely, given a stationary sequence of possibly correlated random
    variables `X_1, X_2,...,X_N`, each identically distributed ESS is the number
    such that

    ```Variance{ N**-1 * Sum{X_i} } = ESS**-1 * Variance{ X_1 }.```

    If the sequence is uncorrelated, `ESS = N`.  In general, one should expect
    `ESS <= N`, with more highly correlated sequences having smaller `ESS`.

    Args:
      states:  `Tensor` or list of `Tensor` objects.  Dimension zero should index
        identically distributed states.
      filter_threshold:  `Tensor` or list of `Tensor` objects.
        Must broadcast with `state`.  The auto-correlation sequence is truncated
        after the first appearance of a term less than `filter_threshold`.
        Setting to `None` means we use no threshold filter.  Since `|R_k| <= 1`,
        setting to any number less than `-1` has the same effect.
      filter_beyond_lag:  `Tensor` or list of `Tensor` objects.  Must be
        `int`-like and scalar valued.  The auto-correlation sequence is truncated
        to this length.  Setting to `None` means we do not filter based on number
        of lags.
      name:  `String` name to prepend to created ops.

    Returns:
      ess:  `Tensor` or list of `Tensor` objects.  The effective sample size of
        each component of `states`.  Shape will be `states.shape[1:]`.

    Raises:
      ValueError:  If `states` and `filter_threshold` or `states` and
        `filter_beyond_lag` are both lists with different lengths.

    #### Examples

    We use ESS to estimate standard error.

    ```
    import tensorflow as tf
    import tensorflow_probability as tfp
    tfd = tfp.distributions

    target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.])

    # Get 1000 states from one chain.
    states = tfp.mcmc.sample_chain(
        num_burnin_steps=200,
        num_results=1000,
        current_state=tf.constant([0., 0.]),
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target.log_prob,
          step_size=0.05,
          num_leapfrog_steps=20))
    states.shape
    ==> (1000, 2)

    ess = effective_sample_size(states)
    ==> Shape (2,) Tensor

    mean, variance = tf.nn.moments(states, axis=0)
    standard_error = tf.sqrt(variance / ess)
    ```

    Some math shows that, with `R_k` the auto-correlation sequence,
    `R_k := Covariance{X_1, X_{1+k}} / Variance{X_1}`, we have

    ```ESS(N) =  N / [ 1 + 2 * ( (N - 1) / N * R_1 + ... + 1 / N * R_{N-1}  ) ]```

    This function estimates the above by first estimating the auto-correlation.
    Since `R_k` must be estimated using only `N - k` samples, it becomes
    progressively noisier for larger `k`.  For this reason, the summation over
    `R_k` should be truncated at some number `filter_beyond_lag < N`.  Since many
    MCMC methods generate chains where `R_k > 0`, a reasonable criteria is to
    truncate at the first index where the estimated auto-correlation becomes
    negative.

    The arguments `filter_beyond_lag`, `filter_threshold` are filters intended to
    remove noisy tail terms from `R_k`.  They combine in an "OR" manner meaning
    terms are removed if they were to be filtered under the `filter_beyond_lag` OR
    `filter_threshold` criteria.


with tf.compat.v1.name_scope(<EOL>'<STR_LIT>',<EOL>values=[states, filter_beyond_lag, filter_threshold]):<EOL><INDENT>states = tf.convert_to_tensor(value=states, name='<STR_LIT>')<EOL>dt = states.dtype<EOL>auto_corr = stats.auto_correlation(<EOL>states, axis=<NUM_LIT:0>, max_lags=filter_beyond_lag)<EOL>if filter_threshold is not None:<EOL><INDENT>filter_threshold = tf.convert_to_tensor(<EOL>value=filter_threshold, dtype=dt, name='<STR_LIT>')<EOL>mask = auto_corr < filter_threshold<EOL>mask = tf.cast(mask, dtype=dt)<EOL>mask = tf.cumsum(mask, axis=<NUM_LIT:0>)<EOL>mask = tf.maximum(<NUM_LIT:1.> - mask, <NUM_LIT:0.>)<EOL>auto_corr *= mask<EOL><DEDENT>n = _axis_size(states, axis=<NUM_LIT:0>)<EOL>k = tf.range(<NUM_LIT:0.>, _axis_size(auto_corr, axis=<NUM_LIT:0>))<EOL>nk_factor = (n - k) / n<EOL>if auto_corr.shape.ndims is not None:<EOL><INDENT>new_shape = [-<NUM_LIT:1>] + [<NUM_LIT:1>] * (auto_corr.shape.ndims - <NUM_LIT:1>)<EOL><DEDENT>else:<EOL><INDENT>new_shape = tf.concat(<EOL>([-<NUM_LIT:1>],<EOL>tf.ones([tf.rank(auto_corr) - <NUM_LIT:1>], dtype=tf.int32)),<EOL>axis=<NUM_LIT:0>)<EOL><DEDENT>nk_factor = tf.reshape(nk_factor, new_shape)<EOL>return n / (-<NUM_LIT:1> +<EOL><NUM_LIT:2> * tf.reduce_sum(input_tensor=nk_factor * auto_corr, axis=<NUM_LIT:0>))<EOL><DEDENT>

ESS computation for one single Tensor argument.


chains_states_was_list = _is_list_like(chains_states)<EOL>if not chains_states_was_list:<EOL><INDENT>chains_states = [chains_states]<EOL><DEDENT>icn_const_ = tf.get_static_value(<EOL>tf.convert_to_tensor(value=independent_chain_ndims))<EOL>if icn_const_ is not None:<EOL><INDENT>independent_chain_ndims = icn_const_<EOL>if icn_const_ < <NUM_LIT:1>:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(<EOL>independent_chain_ndims))<EOL><DEDENT><DEDENT>with tf.compat.v1.name_scope(name, '<STR_LIT>'):<EOL><INDENT>rhat_list = [<EOL>_potential_scale_reduction_single_state(s, independent_chain_ndims)<EOL>for s in chains_states<EOL>]<EOL><DEDENT>if chains_states_was_list:<EOL><INDENT>return rhat_list<EOL><DEDENT>return rhat_list[<NUM_LIT:0>]<EOL>

Gelman and Rubin (1992)'s potential scale reduction for chain convergence.

    Given `N > 1` states from each of `C > 1` independent chains, the potential
    scale reduction factor, commonly referred to as R-hat, measures convergence of
    the chains (to the same target) by testing for equality of means.
    Specifically, R-hat measures the degree to which variance (of the means)
    between chains exceeds what one would expect if the chains were identically
    distributed. See [Gelman and Rubin (1992)][1]; [Brooks and Gelman (1998)][2].

    Some guidelines:

    * The initial state of the chains should be drawn from a distribution
      overdispersed with respect to the target.
    * If all chains converge to the target, then as `N --> infinity`, R-hat --> 1.
      Before that, R-hat > 1 (except in pathological cases, e.g. if the chain
      paths were identical).
    * The above holds for any number of chains `C > 1`.  Increasing `C` does
      improves effectiveness of the diagnostic.
    * Sometimes, R-hat < 1.2 is used to indicate approximate convergence, but of
      course this is problem dependent. See [Brooks and Gelman (1998)][2].
    * R-hat only measures non-convergence of the mean. If higher moments, or
      other statistics are desired, a different diagnostic should be used. See
      [Brooks and Gelman (1998)][2].

    Args:
      chains_states:  `Tensor` or Python `list` of `Tensor`s representing the
        state(s) of a Markov Chain at each result step.  The `ith` state is
        assumed to have shape `[Ni, Ci1, Ci2,...,CiD] + A`.
        Dimension `0` indexes the `Ni > 1` result steps of the Markov Chain.
        Dimensions `1` through `D` index the `Ci1 x ... x CiD` independent
        chains to be tested for convergence to the same target.
        The remaining dimensions, `A`, can have any shape (even empty).
      independent_chain_ndims: Integer type `Tensor` with value `>= 1` giving the
        number of giving the number of dimensions, from `dim = 1` to `dim = D`,
        holding independent chain results to be tested for convergence.
      name: `String` name to prepend to created tf.  Default:
        `potential_scale_reduction`.

    Returns:
      `Tensor` or Python `list` of `Tensor`s representing the R-hat statistic for
      the state(s).  Same `dtype` as `state`, and shape equal to
      `state.shape[1 + independent_chain_ndims:]`.

    Raises:
      ValueError:  If `independent_chain_ndims < 1`.

    #### Examples

    Diagnosing convergence by monitoring 10 chains that each attempt to
    sample from a 2-variate normal.

    ```python
    import tensorflow as tf
    import tensorflow_probability as tfp
    tfd = tfp.distributions

    target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.])

    # Get 10 (2x) overdispersed initial states.
    initial_state = target.sample(10) * 2.
    ==> (10, 2)

    # Get 1000 samples from the 10 independent chains.
    chains_states, _ = tfp.mcmc.sample_chain(
        num_burnin_steps=200,
        num_results=1000,
        current_state=initial_state,
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=target.log_prob,
            step_size=0.05,
            num_leapfrog_steps=20))
    chains_states.shape
    ==> (1000, 10, 2)

    rhat = tfp.mcmc.diagnostic.potential_scale_reduction(
        chains_states, independent_chain_ndims=1)

    # The second dimension needed a longer burn-in.
    rhat.eval()
    ==> [1.05, 1.3]
    ```

    To see why R-hat is reasonable, let `X` be a random variable drawn uniformly
    from the combined states (combined over all chains).  Then, in the limit
    `N, C --> infinity`, with `E`, `Var` denoting expectation and variance,

    ```R-hat = ( E[Var[X | chain]] + Var[E[X | chain]] ) / E[Var[X | chain]].```

    Using the law of total variance, the numerator is the variance of the combined
    states, and the denominator is the total variance minus the variance of the
    the individual chain means.  If the chains are all drawing from the same
    distribution, they will have the same mean, and thus the ratio should be one.

    #### References

    [1]: Stephen P. Brooks and Andrew Gelman. General Methods for Monitoring
         Convergence of Iterative Simulations. _Journal of Computational and
         Graphical Statistics_, 7(4), 1998.

    [2]: Andrew Gelman and Donald B. Rubin. Inference from Iterative Simulation
         Using Multiple Sequences. _Statistical Science_, 7(4):457-472, 1992.


with tf.compat.v1.name_scope(<EOL>'<STR_LIT>',<EOL>values=[state, independent_chain_ndims]):<EOL><INDENT>state = tf.convert_to_tensor(value=state, name='<STR_LIT:state>')<EOL>sample_ndims = <NUM_LIT:1><EOL>sample_axis = tf.range(<NUM_LIT:0>, sample_ndims)<EOL>chain_axis = tf.range(sample_ndims,<EOL>sample_ndims + independent_chain_ndims)<EOL>sample_and_chain_axis = tf.range(<EOL><NUM_LIT:0>, sample_ndims + independent_chain_ndims)<EOL>n = _axis_size(state, sample_axis)<EOL>m = _axis_size(state, chain_axis)<EOL>b_div_n = _reduce_variance(<EOL>tf.reduce_mean(input_tensor=state, axis=sample_axis, keepdims=True),<EOL>sample_and_chain_axis,<EOL>biased=False)<EOL>w = tf.reduce_mean(<EOL>input_tensor=_reduce_variance(<EOL>state, sample_axis, keepdims=True, biased=True),<EOL>axis=sample_and_chain_axis)<EOL>sigma_2_plus = w + b_div_n<EOL>return ((m + <NUM_LIT:1.>) / m) * sigma_2_plus / w - (n - <NUM_LIT:1.>) / (m * n)<EOL><DEDENT>

potential_scale_reduction for one single state `Tensor`.


if axis is None:<EOL><INDENT>return tf.cast(tf.size(input=x), x.dtype)<EOL><DEDENT>return tf.cast(<EOL>tf.reduce_prod(input_tensor=tf.gather(tf.shape(input=x), axis)), x.dtype)<EOL>

Get number of elements of `x` in `axis`, as type `x.dtype`.


return isinstance(x, (tuple, list))<EOL>

Helper which returns `True` if input is `list`-like.


def _broadcast_maybelist_arg(states, secondary_arg, name):

if _is_list_like(secondary_arg):<EOL><INDENT>if len(secondary_arg) != len(states):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(name, len(states)))<EOL><DEDENT><DEDENT>else:<EOL><INDENT>secondary_arg = [secondary_arg] * len(states)<EOL><DEDENT>return secondary_arg<EOL>

Broadcast a listable secondary_arg to that of states.


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

Helper which uses graph seed if using eager.


def _log_gamma_log_prob(self, x, event_dims=()):

return tf.reduce_sum(<EOL>input_tensor=self._shape_param * x - self._rate_param * tf.exp(x),<EOL>axis=event_dims)<EOL>

Computes log-pdf of a log-gamma random variable.

        Args:
          x: Value of the random variable.
          event_dims: Dimensions not to treat as independent.

        Returns:
          log_prob: The log-pdf up to a normalizing constant.


x = tf.constant(np.random.rand(<NUM_LIT:50>, <NUM_LIT:10>, <NUM_LIT:2>), np.float32)<EOL>self._integrator_conserves_energy(x, independent_chain_ndims)<EOL>

Tests the long-term energy conservation of the leapfrog integrator.

        The leapfrog integrator is symplectic, so for sufficiently small step
        sizes it should be possible to run it more or less indefinitely without
        the energy of the system blowing up or collapsing.

        Args:
          independent_chain_ndims: Python `int` scalar representing the number of
            dims associated with independent chains.


initial_draws = np.log(np.random.gamma(self._shape_param,<EOL>size=[<NUM_LIT>, <NUM_LIT:2>, <NUM_LIT:2>]))<EOL>initial_draws -= np.log(self._rate_param)<EOL>x = tf.constant(initial_draws, np.float32)<EOL>self._kernel_leaves_target_invariant(x, independent_chain_ndims)<EOL>

Tests that the kernel leaves the target distribution invariant.

        Draws some independent samples from the target distribution,
        applies an iteration of the MCMC kernel, then runs a
        Kolmogorov-Smirnov test to determine if the distribution of the
        MCMC-updated samples has changed.

        We also confirm that running the kernel with a different log-pdf
        does change the target distribution. (And that we can detect that.)

        Args:
          independent_chain_ndims: Python `int` scalar representing the number of
            dims associated with independent chains.


if dim < <NUM_LIT:1> or dim > _MAX_DIMENSION:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(_MAX_DIMENSION,<EOL>dim))<EOL><DEDENT>if (num_results is None) == (sequence_indices is None):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if not dtype.is_floating:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[num_results, sequence_indices]):<EOL><INDENT>if num_results is not None:<EOL><INDENT>num_results = tf.convert_to_tensor(value=num_results)<EOL><DEDENT>if sequence_indices is not None:<EOL><INDENT>sequence_indices = tf.convert_to_tensor(value=sequence_indices)<EOL><DEDENT>indices = _get_indices(num_results, sequence_indices, dtype)<EOL>radixes = tf.constant(_PRIMES[<NUM_LIT:0>:dim], dtype=dtype, shape=[dim, <NUM_LIT:1>])<EOL>max_sizes_by_axes = _base_expansion_size(<EOL>tf.reduce_max(input_tensor=indices), radixes)<EOL>max_size = tf.reduce_max(input_tensor=max_sizes_by_axes)<EOL>exponents_by_axes = tf.tile([tf.range(max_size)], [dim, <NUM_LIT:1>])<EOL>weight_mask = exponents_by_axes >= max_sizes_by_axes<EOL>capped_exponents = tf.where(<EOL>weight_mask,<EOL>tf.zeros_like(exponents_by_axes),<EOL>exponents_by_axes)<EOL>weights = radixes ** capped_exponents<EOL>coeffs = tf.math.floordiv(indices, weights)<EOL>coeffs *= <NUM_LIT:1.> - tf.cast(weight_mask, dtype)<EOL>coeffs %= radixes<EOL>if not randomized:<EOL><INDENT>coeffs /= radixes<EOL>return tf.reduce_sum(input_tensor=coeffs / weights, axis=-<NUM_LIT:1>)<EOL><DEDENT>stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>coeffs = _randomize(coeffs, radixes, seed=stream())<EOL>coeffs *= <NUM_LIT:1.> - tf.cast(weight_mask, dtype)<EOL>coeffs /= radixes<EOL>base_values = tf.reduce_sum(input_tensor=coeffs / weights, axis=-<NUM_LIT:1>)<EOL>zero_correction = tf.random.uniform([dim, <NUM_LIT:1>], seed=stream(), dtype=dtype)<EOL>zero_correction /= radixes ** max_sizes_by_axes<EOL>return base_values + tf.reshape(zero_correction, [-<NUM_LIT:1>])<EOL><DEDENT>

r"""Returns a sample from the `dim` dimensional Halton sequence.

    Warning: The sequence elements take values only between 0 and 1. Care must be
    taken to appropriately transform the domain of a function if it differs from
    the unit cube before evaluating integrals using Halton samples. It is also
    important to remember that quasi-random numbers without randomization are not
    a replacement for pseudo-random numbers in every context. Quasi random numbers
    are completely deterministic and typically have significant negative
    autocorrelation unless randomization is used.

    Computes the members of the low discrepancy Halton sequence in dimension
    `dim`. The `dim`-dimensional sequence takes values in the unit hypercube in
    `dim` dimensions. Currently, only dimensions up to 1000 are supported. The
    prime base for the k-th axes is the k-th prime starting from 2. For example,
    if `dim` = 3, then the bases will be [2, 3, 5] respectively and the first
    element of the non-randomized sequence will be: [0.5, 0.333, 0.2]. For a more
    complete description of the Halton sequences see
    [here](https://en.wikipedia.org/wiki/Halton_sequence). For low discrepancy
    sequences and their applications see
    [here](https://en.wikipedia.org/wiki/Low-discrepancy_sequence).

    If `randomized` is true, this function produces a scrambled version of the
    Halton sequence introduced by [Owen (2017)][1]. For the advantages of
    randomization of low discrepancy sequences see [here](
    https://en.wikipedia.org/wiki/Quasi-Monte_Carlo_method#Randomization_of_quasi-Monte_Carlo).

    The number of samples produced is controlled by the `num_results` and
    `sequence_indices` parameters. The user must supply either `num_results` or
    `sequence_indices` but not both.
    The former is the number of samples to produce starting from the first
    element. If `sequence_indices` is given instead, the specified elements of
    the sequence are generated. For example, sequence_indices=tf.range(10) is
    equivalent to specifying n=10.

    #### Examples

    ```python
    import tensorflow as tf
    import tensorflow_probability as tfp

    # Produce the first 1000 members of the Halton sequence in 3 dimensions.
    num_results = 1000
    dim = 3
    sample = tfp.mcmc.sample_halton_sequence(
      dim,
      num_results=num_results,
      seed=127)

    # Evaluate the integral of x_1 * x_2^2 * x_3^3  over the three dimensional
    # hypercube.
    powers = tf.range(1.0, limit=dim + 1)
    integral = tf.reduce_mean(tf.reduce_prod(sample ** powers, axis=-1))
    true_value = 1.0 / tf.reduce_prod(powers + 1.0)
    with tf.Session() as session:
      values = session.run((integral, true_value))

    # Produces a relative absolute error of 1.7%.
    print ("Estimated: %f, True Value: %f" % values)

    # Now skip the first 1000 samples and recompute the integral with the next
    # thousand samples. The sequence_indices argument can be used to do this.

    sequence_indices = tf.range(start=1000, limit=1000 + num_results,
                                dtype=tf.int32)
    sample_leaped = tfp.mcmc.sample_halton_sequence(
        dim,
        sequence_indices=sequence_indices,
        seed=111217)

    integral_leaped = tf.reduce_mean(tf.reduce_prod(sample_leaped ** powers,
                                                    axis=-1))
    with tf.Session() as session:
      values = session.run((integral_leaped, true_value))
    # Now produces a relative absolute error of 0.05%.
    print ("Leaped Estimated: %f, True Value: %f" % values)
    ```

    Args:
      dim: Positive Python `int` representing each sample's `event_size.` Must
        not be greater than 1000.
      num_results: (Optional) Positive scalar `Tensor` of dtype int32. The number
        of samples to generate. Either this parameter or sequence_indices must
        be specified but not both. If this parameter is None, then the behaviour
        is determined by the `sequence_indices`.
        Default value: `None`.
      sequence_indices: (Optional) `Tensor` of dtype int32 and rank 1. The
        elements of the sequence to compute specified by their position in the
        sequence. The entries index into the Halton sequence starting with 0 and
        hence, must be whole numbers. For example, sequence_indices=[0, 5, 6] will
        produce the first, sixth and seventh elements of the sequence. If this
        parameter is None, then the `num_results` parameter must be specified
        which gives the number of desired samples starting from the first sample.
        Default value: `None`.
      dtype: (Optional) The dtype of the sample. One of: `float16`, `float32` or
        `float64`.
        Default value: `tf.float32`.
      randomized: (Optional) bool indicating whether to produce a randomized
        Halton sequence. If True, applies the randomization described in
        [Owen (2017)][1].
        Default value: `True`.
      seed: (Optional) Python integer to seed the random number generator. Only
        used if `randomized` is True. If not supplied and `randomized` is True,
        no seed is set.
        Default value: `None`.
      name:  (Optional) Python `str` describing ops managed by this function. If
        not supplied the name of this function is used.
        Default value: "sample_halton_sequence".

    Returns:
      halton_elements: Elements of the Halton sequence. `Tensor` of supplied dtype
        and `shape` `[num_results, dim]` if `num_results` was specified or shape
        `[s, dim]` where s is the size of `sequence_indices` if `sequence_indices`
        were specified.

    Raises:
      ValueError: if both `sequence_indices` and `num_results` were specified or
        if dimension `dim` is less than 1 or greater than 1000.

    #### References

    [1]: Art B. Owen. A randomized Halton algorithm in R. _arXiv preprint
         arXiv:1706.02808_, 2017. https://arxiv.org/abs/1706.02808


def _randomize(coeffs, radixes, seed=None):

given_dtype = coeffs.dtype<EOL>coeffs = tf.cast(coeffs, dtype=tf.int32)<EOL>num_coeffs = tf.shape(input=coeffs)[-<NUM_LIT:1>]<EOL>radixes = tf.reshape(tf.cast(radixes, dtype=tf.int32), shape=[-<NUM_LIT:1>])<EOL>stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>perms = _get_permutations(num_coeffs, radixes, seed=stream())<EOL>perms = tf.reshape(perms, shape=[-<NUM_LIT:1>])<EOL>radix_sum = tf.reduce_sum(input_tensor=radixes)<EOL>radix_offsets = tf.reshape(tf.cumsum(radixes, exclusive=True),<EOL>shape=[-<NUM_LIT:1>, <NUM_LIT:1>])<EOL>offsets = radix_offsets + tf.range(num_coeffs) * radix_sum<EOL>permuted_coeffs = tf.gather(perms, coeffs + offsets)<EOL>return tf.cast(permuted_coeffs, dtype=given_dtype)<EOL>

Applies the Owen (2017) randomization to the coefficients.


def _get_permutations(num_results, dims, seed=None):

sample_range = tf.range(num_results)<EOL>stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>def generate_one(d):<EOL><INDENT>seed = stream()<EOL>fn = lambda _: tf.random.shuffle(tf.range(d), seed=seed)<EOL>return tf.map_fn(<EOL>fn,<EOL>sample_range,<EOL>parallel_iterations=<NUM_LIT:1> if seed is not None else <NUM_LIT:10>)<EOL><DEDENT>return tf.concat([generate_one(d) for d in tf.unstack(dims)],<EOL>axis=-<NUM_LIT:1>)<EOL>

Uniform iid sample from the space of permutations.

    Draws a sample of size `num_results` from the group of permutations of degrees
    specified by the `dims` tensor. These are packed together into one tensor
    such that each row is one sample from each of the dimensions in `dims`. For
    example, if dims = [2,3] and num_results = 2, the result is a tensor of shape
    [2, 2 + 3] and the first row of the result might look like:
    [1, 0, 2, 0, 1]. The first two elements are a permutation over 2 elements
    while the next three are a permutation over 3 elements.

    Args:
      num_results: A positive scalar `Tensor` of integral type. The number of
        draws from the discrete uniform distribution over the permutation groups.
      dims: A 1D `Tensor` of the same dtype as `num_results`. The degree of the
        permutation groups from which to sample.
      seed: (Optional) Python integer to seed the random number generator.

    Returns:
      permutations: A `Tensor` of shape `[num_results, sum(dims)]` and the same
      dtype as `dims`.


def _get_indices(num_results, sequence_indices, dtype, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[num_results, sequence_indices]):<EOL><INDENT>if sequence_indices is None:<EOL><INDENT>num_results = tf.cast(num_results, dtype=dtype)<EOL>sequence_indices = tf.range(num_results, dtype=dtype)<EOL><DEDENT>else:<EOL><INDENT>sequence_indices = tf.cast(sequence_indices, dtype)<EOL><DEDENT>indices = sequence_indices + <NUM_LIT:1><EOL>return tf.reshape(indices, [-<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT:1>])<EOL><DEDENT>

Generates starting points for the Halton sequence procedure.

    The k'th element of the sequence is generated starting from a positive integer
    which must be distinct for each `k`. It is conventional to choose the starting
    point as `k` itself (or `k+1` if k is zero based). This function generates
    the starting integers for the required elements and reshapes the result for
    later use.

    Args:
      num_results: Positive scalar `Tensor` of dtype int32. The number of samples
        to generate. If this parameter is supplied, then `sequence_indices`
        should be None.
      sequence_indices: `Tensor` of dtype int32 and rank 1. The entries
        index into the Halton sequence starting with 0 and hence, must be whole
        numbers. For example, sequence_indices=[0, 5, 6] will produce the first,
        sixth and seventh elements of the sequence. If this parameter is not None
        then `n` must be None.
      dtype: The dtype of the sample. One of `float32` or `float64`.
        Default is `float32`.
      name: Python `str` name which describes ops created by this function.

    Returns:
      indices: `Tensor` of dtype `dtype` and shape = `[n, 1, 1]`.


return tf.floor(tf.math.log(num) / tf.math.log(bases)) + <NUM_LIT:1><EOL>

Computes the number of terms in the place value expansion.

    Let num = a0 + a1 b + a2 b^2 + ... ak b^k be the place value expansion of
    `num` in base b (ak <> 0). This function computes and returns `k+1` for each
    base `b` specified in `bases`.

    This can be inferred from the base `b` logarithm of `num` as follows:
      $$k = Floor(log_b (num)) + 1  = Floor( log(num) / log(b)) + 1$$

    Args:
      num: Scalar `Tensor` of dtype either `float32` or `float64`. The number to
        compute the base expansion size of.
      bases: `Tensor` of the same dtype as num. The bases to compute the size
        against.

    Returns:
      Tensor of same dtype and shape as `bases` containing the size of num when
      written in that base.


if not mcmc_util.is_list_like(bijector):<EOL><INDENT>bijector = [bijector]<EOL><DEDENT>def fn(transformed_state_parts, event_ndims):<EOL><INDENT>return sum([<EOL>b.forward_log_det_jacobian(sp, event_ndims=e)<EOL>for b, e, sp in zip(bijector, event_ndims, transformed_state_parts)<EOL>])<EOL><DEDENT>return fn<EOL>

Makes a function which applies a list of Bijectors' `log_det_jacobian`s.


if not mcmc_util.is_list_like(bijector):<EOL><INDENT>bijector = [bijector]<EOL><DEDENT>def fn(transformed_state_parts):<EOL><INDENT>return [b.forward(sp) for b, sp in zip(bijector, transformed_state_parts)]<EOL><DEDENT>return fn<EOL>

Makes a function which applies a list of Bijectors' `forward`s.


if not mcmc_util.is_list_like(bijector):<EOL><INDENT>bijector = [bijector]<EOL><DEDENT>def fn(state_parts):<EOL><INDENT>return [b.inverse(sp)<EOL>for b, sp in zip(bijector, state_parts)]<EOL><DEDENT>return fn<EOL>

Makes a function which applies a list of Bijectors' `inverse`s.


def __init__(self, inner_kernel, bijector, name=None):

self._parameters = dict(<EOL>inner_kernel=inner_kernel,<EOL>bijector=bijector,<EOL>name=name or '<STR_LIT>')<EOL>inner_kernel_kwargs = inner_kernel.parameters.copy()<EOL>target_log_prob_fn = inner_kernel_kwargs['<STR_LIT>']<EOL>self._forward_transform = forward_transform_fn(bijector)<EOL>self._inverse_transform = inverse_transform_fn(bijector)<EOL>self._forward_log_det_jacobian = forward_log_det_jacobian_fn(bijector)<EOL>def new_target_log_prob(*transformed_state_parts):<EOL><INDENT>"""<STR_LIT>"""<EOL>transformed_state_parts = [<EOL>tf.identity(sp) for sp in transformed_state_parts<EOL>]<EOL>tlp = target_log_prob_fn(<EOL>*self._forward_transform(transformed_state_parts))<EOL>event_ndims = [<EOL>tf.rank(sp) - tf.rank(tlp) for sp in transformed_state_parts<EOL>]<EOL>return tlp + self._forward_log_det_jacobian(<EOL>transformed_state_parts=transformed_state_parts,<EOL>event_ndims=event_ndims)<EOL><DEDENT>inner_kernel_kwargs.update(target_log_prob_fn=new_target_log_prob)<EOL>with deprecation.silence():<EOL><INDENT>self._inner_kernel = type(inner_kernel)(**inner_kernel_kwargs)<EOL><DEDENT>

Instantiates this object.

        Args:
          inner_kernel: `TransitionKernel`-like object which has a
            `target_log_prob_fn` argument.
          bijector: `tfp.distributions.Bijector` or list of
            `tfp.distributions.Bijector`s. These bijectors use `forward` to map the
            `inner_kernel` state space to the state expected by
            `inner_kernel.target_log_prob_fn`.
          name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., "transformed_kernel").

        Returns:
          transformed_kernel: Instance of `TransitionKernel` which copies the input
            transition kernel then modifies its `target_log_prob_fn` by applying the
            provided bijector(s).


return self._parameters<EOL>

Return `dict` of ``__init__`` arguments and their values.


def one_step(self, current_state, previous_kernel_results):

with tf.compat.v1.name_scope(<EOL>name=mcmc_util.make_name(self.name, '<STR_LIT>', '<STR_LIT>'),<EOL>values=[previous_kernel_results]):<EOL><INDENT>transformed_next_state, kernel_results = self._inner_kernel.one_step(<EOL>previous_kernel_results.transformed_state,<EOL>previous_kernel_results.inner_results)<EOL>transformed_next_state_parts = (<EOL>transformed_next_state<EOL>if mcmc_util.is_list_like(transformed_next_state) else<EOL>[transformed_next_state])<EOL>next_state_parts = self._forward_transform(transformed_next_state_parts)<EOL>next_state = (<EOL>next_state_parts if mcmc_util.is_list_like(transformed_next_state)<EOL>else next_state_parts[<NUM_LIT:0>])<EOL>kernel_results = TransformedTransitionKernelResults(<EOL>transformed_state=transformed_next_state,<EOL>inner_results=kernel_results)<EOL>return next_state, kernel_results<EOL><DEDENT>

Runs one iteration of the Transformed Kernel.

        Args:
          current_state: `Tensor` or Python `list` of `Tensor`s
            representing the current state(s) of the Markov chain(s),
            _after_ application of `bijector.forward`. The first `r`
            dimensions index independent chains,
            `r = tf.rank(target_log_prob_fn(*current_state))`. The
            `inner_kernel.one_step` does not actually use `current_state`,
            rather it takes as input
            `previous_kernel_results.transformed_state` (because
            `TransformedTransitionKernel` creates a copy of the input
            inner_kernel with a modified `target_log_prob_fn` which
            internally applies the `bijector.forward`).
          previous_kernel_results: `collections.namedtuple` containing `Tensor`s
            representing values from previous calls to this function (or from the
            `bootstrap_results` function.)

        Returns:
          next_state: Tensor or Python list of `Tensor`s representing the state(s)
            of the Markov chain(s) after taking exactly one step. Has same type and
            shape as `current_state`.
          kernel_results: `collections.namedtuple` of internal calculations used to
            advance the chain.


def bootstrap_results(self, init_state=None, transformed_init_state=None):

if (init_state is None) == (transformed_init_state is None):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>with tf.compat.v1.name_scope(<EOL>name=mcmc_util.make_name(self.name, '<STR_LIT>',<EOL>'<STR_LIT>'),<EOL>values=[init_state, transformed_init_state]):<EOL><INDENT>if transformed_init_state is None:<EOL><INDENT>init_state_parts = (init_state if mcmc_util.is_list_like(init_state)<EOL>else [init_state])<EOL>transformed_init_state_parts = self._inverse_transform(init_state_parts)<EOL>transformed_init_state = (<EOL>transformed_init_state_parts if mcmc_util.is_list_like(init_state)<EOL>else transformed_init_state_parts[<NUM_LIT:0>])<EOL><DEDENT>else:<EOL><INDENT>if mcmc_util.is_list_like(transformed_init_state):<EOL><INDENT>transformed_init_state = [<EOL>tf.convert_to_tensor(value=s, name='<STR_LIT>')<EOL>for s in transformed_init_state<EOL>]<EOL><DEDENT>else:<EOL><INDENT>transformed_init_state = tf.convert_to_tensor(<EOL>value=transformed_init_state, name='<STR_LIT>')<EOL><DEDENT><DEDENT>kernel_results = TransformedTransitionKernelResults(<EOL>transformed_state=transformed_init_state,<EOL>inner_results=self._inner_kernel.bootstrap_results(<EOL>transformed_init_state))<EOL>return kernel_results<EOL><DEDENT>

Returns an object with the same type as returned by `one_step`.

        Unlike other `TransitionKernel`s,
        `TransformedTransitionKernel.bootstrap_results` has the option of
        initializing the `TransformedTransitionKernelResults` from either an initial
        state, eg, requiring computing `bijector.inverse(init_state)`, or
        directly from `transformed_init_state`, i.e., a `Tensor` or list
        of `Tensor`s which is interpretted as the `bijector.inverse`
        transformed state.

        Args:
          init_state: `Tensor` or Python `list` of `Tensor`s representing the a
            state(s) of the Markov chain(s). Must specify `init_state` or
            `transformed_init_state` but not both.
          transformed_init_state: `Tensor` or Python `list` of `Tensor`s
            representing the a state(s) of the Markov chain(s). Must specify
            `init_state` or `transformed_init_state` but not both.

        Returns:
          kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
            `Tensor`s representing internal calculations made within this function.

        Raises:
          ValueError: if `inner_kernel` results doesn't contain the member
            "target_log_prob".

        #### Examples

        To use `transformed_init_state` in context of
        `tfp.mcmc.sample_chain`, you need to explicitly pass the
        `previous_kernel_results`, e.g.,

        ```python
        transformed_kernel = tfp.mcmc.TransformedTransitionKernel(...)
        init_state = ...        # Doesnt matter.
        transformed_init_state = ... # Does matter.
        results, _ = tfp.mcmc.sample_chain(
            num_results=...,
            current_state=init_state,
            previous_kernel_results=transformed_kernel.bootstrap_results(
                transformed_init_state=transformed_init_state),
            kernel=transformed_kernel)
        ```


return isinstance(x, (tuple, list))<EOL>

Helper which returns `True` if input is `list`-like.


try:<EOL><INDENT>for fn in x._fields:<EOL><INDENT>_ = getattr(x, fn)<EOL><DEDENT>return True<EOL><DEDENT>except AttributeError:<EOL><INDENT>return False<EOL><DEDENT>

Helper which returns `True` if input is `collections.namedtuple`-like.


def make_name(super_name, default_super_name, sub_name):

name = super_name if super_name is not None else default_super_name<EOL>if sub_name is not None:<EOL><INDENT>name += '<STR_LIT:_>' + sub_name<EOL><DEDENT>return name<EOL>

Helper which makes a `str` name; useful for tf.compat.v1.name_scope.


def _expand_is_accepted_like(x):<EOL><INDENT>"""<STR_LIT>"""<EOL>with tf.compat.v1.name_scope('<STR_LIT>'):<EOL><INDENT>expand_shape = tf.concat([<EOL>tf.shape(input=is_accepted),<EOL>tf.ones([tf.rank(x) - tf.rank(is_accepted)], dtype=tf.int32),<EOL>],<EOL>axis=<NUM_LIT:0>)<EOL>multiples = tf.concat([<EOL>tf.ones([tf.rank(is_accepted)], dtype=tf.int32),<EOL>tf.shape(input=x)[tf.rank(is_accepted):],<EOL>],<EOL>axis=<NUM_LIT:0>)<EOL>m = tf.tile(tf.reshape(is_accepted, expand_shape),<EOL>multiples)<EOL>m.set_shape(m.shape.merge_with(x.shape))<EOL>return m<EOL><DEDENT><DEDENT>def _where(accepted, rejected):<EOL><INDENT>if accepted is rejected:<EOL><INDENT>return accepted<EOL><DEDENT>accepted = tf.convert_to_tensor(value=accepted, name='<STR_LIT>')<EOL>rejected = tf.convert_to_tensor(value=rejected, name='<STR_LIT>')<EOL>r = tf.where(_expand_is_accepted_like(accepted), accepted, rejected)<EOL>r.set_shape(r.shape.merge_with(accepted.shape.merge_with(rejected.shape)))<EOL>return r<EOL><DEDENT>with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[is_accepted, accepted, rejected]):<EOL><INDENT>if not is_list_like(accepted):<EOL><INDENT>return _where(accepted, rejected)<EOL><DEDENT>return [(choose(is_accepted, a, r, name=name) if is_namedtuple_like(a)<EOL>else _where(a, r))<EOL>for a, r in zip(accepted, rejected)]<EOL><DEDENT>

Helper to `choose` which expand_dims `is_accepted` and applies tf.where.


def choose(is_accepted, accepted, rejected, name=None):

if not is_namedtuple_like(accepted):<EOL><INDENT>return _choose_base_case(is_accepted, accepted, rejected, name=name)<EOL><DEDENT>if not isinstance(accepted, type(rejected)):<EOL><INDENT>raise TypeError('<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>type(accepted).__name__,<EOL>type(rejected).__name__))<EOL><DEDENT>return type(accepted)(**dict(<EOL>[(fn,<EOL>choose(is_accepted,<EOL>getattr(accepted, fn),<EOL>getattr(rejected, fn),<EOL>name=name))<EOL>for fn in accepted._fields]))<EOL>

Helper which expand_dims `is_accepted` then applies tf.where.


def safe_sum(x, alt_value=-np.inf, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [x, alt_value]):<EOL><INDENT>if not is_list_like(x):<EOL><INDENT>raise TypeError('<STR_LIT>')<EOL><DEDENT>if not x:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>in_shape = x[<NUM_LIT:0>].shape<EOL>x = tf.stack(x, axis=-<NUM_LIT:1>)<EOL>x = tf.reduce_sum(input_tensor=x, axis=-<NUM_LIT:1>)<EOL>alt_value = np.array(alt_value, x.dtype.as_numpy_dtype)<EOL>alt_fill = tf.fill(tf.shape(input=x), value=alt_value)<EOL>x = tf.where(tf.math.is_finite(x), x, alt_fill)<EOL>x.set_shape(x.shape.merge_with(in_shape))<EOL>return x<EOL><DEDENT>

Elementwise adds list members, replacing non-finite results with alt_value.

    Typically the `alt_value` is chosen so the `MetropolisHastings`
    `TransitionKernel` always rejects the proposal.

    Args:
      x: Python `list` of `Tensors` to elementwise add.
      alt_value: Python scalar used to replace any elementwise sums which would
        otherwise be non-finite.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "safe_sum").

    Returns:
      safe_sum: `Tensor` representing the elementwise sum of list of `Tensor`s
        `x` or `alt_value` where sums are non-finite.

    Raises:
      TypeError: if `x` is not list-like.
      ValueError: if `x` is empty.


def _doc(func):<EOL><INDENT>func.__doc__ = value<EOL>return func<EOL><DEDENT>return _doc<EOL>

Decorator to programmatically set a function docstring.


def _value_and_gradients(fn, fn_arg_list, result=None, grads=None, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[fn_arg_list, result, grads]):<EOL><INDENT>def _convert_to_tensor(x, name):<EOL><INDENT>ctt = lambda x_: x_ if x_ is None else tf.convert_to_tensor(<EOL>value=x_, name=name)<EOL>return [ctt(x_) for x_ in x] if is_list_like(x) else ctt(x)<EOL><DEDENT>fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)<EOL>else [fn_arg_list])<EOL>fn_arg_list = _convert_to_tensor(fn_arg_list, '<STR_LIT>')<EOL>if result is None:<EOL><INDENT>result = fn(*fn_arg_list)<EOL>if grads is None and tf.executing_eagerly():<EOL><INDENT>fn_arg_list = [<NUM_LIT:0> + x for x in fn_arg_list]<EOL><DEDENT><DEDENT>result = _convert_to_tensor(result, '<STR_LIT>')<EOL>if grads is not None:<EOL><INDENT>grads = _convert_to_tensor(grads, '<STR_LIT>')<EOL>return result, grads<EOL><DEDENT>if is_list_like(result) and len(result) == len(fn_arg_list):<EOL><INDENT>def fn_slice(i):<EOL><INDENT>"""<STR_LIT>"""<EOL>return lambda x: fn(*(fn_arg_list[:i] + [x] + fn_arg_list[i+<NUM_LIT:1>:]))<EOL><DEDENT>grads = [<EOL>tfp_math_value_and_gradients(fn_slice(i), fn_arg_list[i])[<NUM_LIT:1>]<EOL>for i in range(len(result))<EOL>]<EOL><DEDENT>else:<EOL><INDENT>_, grads = tfp_math_value_and_gradients(fn, fn_arg_list)<EOL><DEDENT>return result, grads<EOL><DEDENT>


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[fn_arg_list, result, grads]):<EOL><INDENT>fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)<EOL>else [fn_arg_list])<EOL>result, grads = _value_and_gradients(fn, fn_arg_list, result, grads)<EOL>if not all(r.dtype.is_floating<EOL>for r in (result if is_list_like(result) else [result])):  <EOL><INDENT>raise TypeError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if len(fn_arg_list) != len(grads):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if check_non_none_grads and any(g is None for g in grads):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(fn_arg_list, grads))<EOL><DEDENT>return result, grads<EOL><DEDENT>

Calls `fn` and computes the gradient of the result wrt `args_list`.


def smart_for_loop(loop_num_iter, body_fn, initial_loop_vars,<EOL>parallel_iterations=<NUM_LIT:10>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[loop_num_iter, initial_loop_vars]):<EOL><INDENT>loop_num_iter_ = tf.get_static_value(loop_num_iter)<EOL>if (loop_num_iter_ is None or tf.executing_eagerly() or<EOL>control_flow_util.GraphOrParentsInXlaContext(<EOL>tf.compat.v1.get_default_graph())):<EOL><INDENT>loop_num_iter = tf.cast(loop_num_iter, dtype=tf.int32)<EOL>return tf.while_loop(<EOL>cond=lambda i, *args: i < loop_num_iter,<EOL>body=lambda i, *args: [i + <NUM_LIT:1>] + list(body_fn(*args)),<EOL>loop_vars=[np.int32(<NUM_LIT:0>)] + initial_loop_vars,<EOL>parallel_iterations=parallel_iterations<EOL>)[<NUM_LIT:1>:]<EOL><DEDENT>result = initial_loop_vars<EOL>for _ in range(loop_num_iter_):<EOL><INDENT>result = body_fn(*result)<EOL><DEDENT>return result<EOL><DEDENT>

Construct a for loop, preferring a python loop if `n` is staticaly known.

    Given `loop_num_iter` and `body_fn`, return an op corresponding to executing
    `body_fn` `loop_num_iter` times, feeding previous outputs of `body_fn` into
    the next iteration.

    If `loop_num_iter` is statically known, the op is constructed via python for
    loop, and otherwise a `tf.while_loop` is used.

    Args:
      loop_num_iter: `Integer` `Tensor` representing the number of loop
        iterations.
      body_fn: Callable to be executed `loop_num_iter` times.
      initial_loop_vars: Listlike object of `Tensors` to be passed in to
        `body_fn`'s first execution.
      parallel_iterations: The number of iterations allowed to run in parallel.
        It must be a positive integer. See `tf.while_loop` for more details.
        Default value: `10`.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "smart_for_loop").
    Returns:
      result: `Tensor` representing applying `body_fn` iteratively `n` times.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', [initial_state, elems]), tf.compat.v1.variable_scope(<EOL>tf.compat.v1.get_variable_scope()) as vs:<EOL><INDENT>if vs.caching_device is None and not tf.executing_eagerly():<EOL><INDENT>vs.set_caching_device(lambda op: op.device)<EOL><DEDENT>initial_state = tf.nest.map_structure(<EOL>lambda x: tf.convert_to_tensor(value=x, name='<STR_LIT>'),<EOL>initial_state)<EOL>elems = tf.convert_to_tensor(value=elems, name='<STR_LIT>')<EOL>static_length = elems.shape[<NUM_LIT:0>]<EOL>if tf.compat.dimension_value(static_length) is None:<EOL><INDENT>length = tf.shape(input=elems)[<NUM_LIT:0>]<EOL><DEDENT>else:<EOL><INDENT>length = tf.convert_to_tensor(<EOL>value=static_length, dtype=tf.int32, name='<STR_LIT>')<EOL><DEDENT>elems_array = tf.TensorArray(<EOL>elems.dtype, size=length, element_shape=elems.shape[<NUM_LIT:1>:])<EOL>elems_array = elems_array.unstack(elems)<EOL>trace_arrays = tf.nest.map_structure(<EOL>lambda x: tf.TensorArray(x.dtype, size=length, element_shape=x.shape),<EOL>trace_fn(initial_state))<EOL>def _body(i, state, trace_arrays):<EOL><INDENT>state = loop_fn(state, elems_array.read(i))<EOL>trace_arrays = tf.nest.pack_sequence_as(trace_arrays, [<EOL>a.write(i, v) for a, v in zip(<EOL>tf.nest.flatten(trace_arrays), tf.nest.flatten(trace_fn(state)))<EOL>])<EOL>return i + <NUM_LIT:1>, state, trace_arrays<EOL><DEDENT>_, final_state, trace_arrays = tf.while_loop(<EOL>cond=lambda i, *args: i < length,<EOL>body=_body,<EOL>loop_vars=(<NUM_LIT:0>, initial_state, trace_arrays),<EOL>parallel_iterations=parallel_iterations)<EOL>stacked_trace = tf.nest.map_structure(lambda x: x.stack(), trace_arrays)<EOL>def _merge_static_length(x):<EOL><INDENT>x.set_shape(tf.TensorShape(static_length).concatenate(x.shape[<NUM_LIT:1>:]))<EOL>return x<EOL><DEDENT>stacked_trace = tf.nest.map_structure(_merge_static_length, stacked_trace)<EOL>return final_state, stacked_trace<EOL><DEDENT>

A simplified version of `tf.scan` that has configurable tracing.

    This function repeatedly calls `loop_fn(state, elem)`, where `state` is the
    `initial_state` during the first iteration, and the return value of `loop_fn`
    for every iteration thereafter. `elem` is a slice of `elements` along the
    first dimension, accessed in order. Additionally, it calls `trace_fn` on the
    return value of `loop_fn`. The `Tensor`s in return values of `trace_fn` are
    stacked and returned from this function, such that the first dimension of
    those `Tensor`s matches the size of `elems`.

    Args:
      loop_fn: A callable that takes in a `Tensor` or a nested collection of
        `Tensor`s with the same structure as `initial_state`, a slice of `elems`
        and returns the same structure as `initial_state`.
      initial_state: A `Tensor` or a nested collection of `Tensor`s passed to
        `loop_fn` in the first iteration.
      elems: A `Tensor` that is split along the first dimension and each element
        of which is passed to `loop_fn`.
      trace_fn: A callable that takes in the return value of `loop_fn` and returns
        a `Tensor` or a nested collection of `Tensor`s.
      parallel_iterations: Passed to the internal `tf.while_loop`.
      name: Name scope used in this function. Default: 'trace_scan'.

    Returns:
      final_state: The final return value of `loop_fn`.
      trace: The same structure as the return value of `trace_fn`, but with each
        `Tensor` being a stack of the corresponding `Tensors` in the return value
        of `trace_fn` for each slice of `elems`.


@functools.wraps(setter)<EOL>def _new_setter(kernel_results, *args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>results_stack = []<EOL>while hasattr(kernel_results, '<STR_LIT>'):<EOL><INDENT>results_stack.append(kernel_results)<EOL>kernel_results = kernel_results.inner_results<EOL><DEDENT>new_kernel_results = setter(kernel_results, *args, **kwargs)<EOL>for outer_results in reversed(results_stack):<EOL><INDENT>new_kernel_results = outer_results._replace(<EOL>inner_results=new_kernel_results)<EOL><DEDENT>return new_kernel_results<EOL><DEDENT>return _new_setter<EOL>

Wraps a setter so it applies to the inner-most results in `kernel_results`.

    The wrapped setter unwraps `kernel_results` and applies `setter` to the first
    results without an `inner_results` attribute.

    Args:
      setter: A callable that takes the kernel results as well as some `*args` and
        `**kwargs` and returns a modified copy of those kernel results.

    Returns:
      new_setter: A wrapped `setter`.


@functools.wraps(getter)<EOL>def _new_getter(kernel_results, *args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>results_stack = []<EOL>while hasattr(kernel_results, '<STR_LIT>'):<EOL><INDENT>results_stack.append(kernel_results)<EOL>kernel_results = kernel_results.inner_results<EOL><DEDENT>return getter(kernel_results, *args, **kwargs)<EOL><DEDENT>return _new_getter<EOL>

Wraps a getter so it applies to the inner-most results in `kernel_results`.

    The wrapped getter unwraps `kernel_results` and returns the return value of
    `getter` called with the first results without an `inner_results` attribute.

    Args:
      getter: A callable that takes Kernel results and returns some value.

    Returns:
      new_getter: A wrapped `getter`.


kernel_stack = []<EOL>while hasattr(kernel, '<STR_LIT>') and '<STR_LIT>' in kernel.parameters:<EOL><INDENT>kernel_stack.append(kernel)<EOL>kernel = kernel.parameters['<STR_LIT>']<EOL><DEDENT>def _recreate_kernel(kernel, parameters):<EOL><INDENT>new_parameters = kernel.parameters.copy()<EOL>new_parameters.update(parameters)<EOL>if '<STR_LIT>' in new_parameters:<EOL><INDENT>new_parameters['<STR_LIT>'] = True<EOL><DEDENT>with deprecation.silence():<EOL><INDENT>return type(kernel)(**new_parameters)<EOL><DEDENT><DEDENT>if hasattr(kernel, '<STR_LIT>'):<EOL><INDENT>kernel = _recreate_kernel(kernel, {})<EOL><DEDENT>for outer_kernel in reversed(kernel_stack):<EOL><INDENT>outer_kernel = _recreate_kernel(outer_kernel, {'<STR_LIT>': kernel})<EOL>kernel = outer_kernel<EOL><DEDENT>return kernel<EOL>

Enables the `store_parameters_in_results` parameter in a chain of kernels.

    This is a temporary utility for use during the transition period of the
    parameter storage methods.

    Args:
      kernel: A TransitionKernel.

    Returns:
      kernel: The same kernel, but recreated with `store_parameters_in_results`
          recursively set to `True` in its parameters and its inner kernels (as
          appropriate).


def _left_doubling_increments(batch_shape, max_doublings, step_size, seed=None,<EOL>name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[batch_shape, max_doublings, step_size]):<EOL><INDENT>step_size = tf.convert_to_tensor(value=step_size)<EOL>dtype = step_size.dtype.base_dtype<EOL>output_shape = tf.concat(([max_doublings + <NUM_LIT:1>], batch_shape), axis=<NUM_LIT:0>)<EOL>expand_left = distributions.Bernoulli(<NUM_LIT:0.5>, dtype=dtype).sample(<EOL>sample_shape=output_shape, seed=seed)<EOL>width_multipliers = tf.cast(<NUM_LIT:2> ** tf.range(<NUM_LIT:0>, max_doublings+<NUM_LIT:1>), dtype=dtype)<EOL>widths_shape = tf.concat(([max_doublings + <NUM_LIT:1>],<EOL>tf.ones_like(batch_shape)), axis=<NUM_LIT:0>)<EOL>width_multipliers = tf.reshape(width_multipliers, shape=widths_shape)<EOL>widths = width_multipliers * step_size<EOL>left_increments = tf.cumsum(widths * expand_left, exclusive=True, axis=<NUM_LIT:0>)<EOL>return left_increments, widths<EOL><DEDENT>

Computes the doubling increments for the left end point.

    The doubling procedure expands an initial interval to find a superset of the
    true slice. At each doubling iteration, the interval width is doubled to
    either the left or the right hand side with equal probability.
    If, initially, the left end point is at `L(0)` and the width of the
    interval is `w(0)`, then the left end point and the width at the
    k-th iteration (denoted L(k) and w(k) respectively) are given by the following
    recursions:

    ```none
    w(k) = 2 * w(k-1)
    L(k) = L(k-1) - w(k-1) * X_k, X_k ~ Bernoulli(0.5)
    or, L(0) - L(k) = w(0) Sum(2^i * X(i+1), 0 <= i < k)
    ```

    This function computes the sequence of `L(0)-L(k)` and `w(k)` for k between 0
    and `max_doublings` independently for each chain.

    Args:
      batch_shape: Positive int32 `tf.Tensor`. The batch shape.
      max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
        doublings to consider.
      step_size: A real `tf.Tensor` with shape compatible with [num_chains].
        The size of the initial interval.
      seed: (Optional) positive int. The random seed. If None, no seed is set.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'find_slice_bounds').

    Returns:
      left_increments: A tensor of shape (max_doublings+1, batch_shape). The
        relative position of the left end point after the doublings.
      widths: A tensor of shape (max_doublings+1, ones_like(batch_shape)). The
        widths of the intervals at each stage of the doubling.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>k = tf.shape(input=x)[<NUM_LIT:0>]<EOL>dtype = x.dtype.base_dtype<EOL>mults = tf.range(<NUM_LIT:2> * k, k, -<NUM_LIT:1>, dtype=dtype)[:, tf.newaxis]<EOL>shifts = tf.range(k, dtype=dtype)[:, tf.newaxis]<EOL>indices = tf.argmax(input=mults * x + shifts, axis=<NUM_LIT:0>, output_type=dtype)<EOL>return indices<EOL><DEDENT>

Finds the index of the optimal set of bounds for each chain.

    For each chain, finds the smallest set of bounds for which both edges lie
    outside the slice. This is equivalent to the point at which a for loop
    implementation (P715 of Neal (2003)) of the algorithm would terminate.

    Performs the following calculation, where i is the number of doublings that
    have been performed and k is the max number of doublings:

    (2 * k - i) * flag + i

    The argmax of the above returns the earliest index where the bounds were
    outside the slice and if there is no such point, the widest bounds.

    Args:
      x: A tensor of shape (max_doublings+1, batch_shape). Type int32, with value
        0 or 1. Indicates if this set of bounds is outside the slice.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'find_slice_bounds').

    Returns:
      indices: A tensor of shape batch_shape. Type int32, with the index of the
        first set of bounds outside the slice and if there are none, the index of
        the widest set.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[x_initial, log_slice_heights, max_doublings, step_size]):<EOL><INDENT>seed_gen = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>x_initial = tf.convert_to_tensor(value=x_initial)<EOL>batch_shape = tf.shape(input=x_initial)<EOL>dtype = step_size.dtype.base_dtype<EOL>left_endpoints = x_initial + step_size * tf.random.uniform(<EOL>batch_shape, minval=-<NUM_LIT:1.0>, maxval=<NUM_LIT:0.0>, dtype=dtype, seed=seed_gen())<EOL>left_increments, widths = _left_doubling_increments(<EOL>batch_shape, max_doublings, step_size, seed=seed_gen())<EOL>left_endpoints -= left_increments<EOL>right_endpoints = left_endpoints + widths<EOL>left_ep_values = tf.map_fn(target_log_prob, left_endpoints)<EOL>right_ep_values = tf.map_fn(target_log_prob, right_endpoints)<EOL>left_ok = left_ep_values < log_slice_heights<EOL>right_ok = right_ep_values < log_slice_heights<EOL>both_ok = left_ok & right_ok<EOL>both_ok_f = tf.reshape(both_ok, [max_doublings + <NUM_LIT:1>, -<NUM_LIT:1>])<EOL>best_interval_idx = _find_best_interval_idx(<EOL>tf.cast(both_ok_f, dtype=tf.int32))<EOL>point_index_gather = tf.stack(<EOL>[best_interval_idx,<EOL>tf.range(tf.size(input=best_interval_idx))],<EOL>axis=<NUM_LIT:1>,<EOL>name='<STR_LIT>')<EOL>left_ep_f = tf.reshape(left_endpoints, [max_doublings + <NUM_LIT:1>, -<NUM_LIT:1>])<EOL>right_ep_f = tf.reshape(right_endpoints, [max_doublings + <NUM_LIT:1>, -<NUM_LIT:1>])<EOL>lower_bounds = tf.reshape(tf.gather_nd(left_ep_f, point_index_gather),<EOL>batch_shape)<EOL>upper_bounds = tf.reshape(tf.gather_nd(right_ep_f, point_index_gather),<EOL>batch_shape)<EOL>both_ok = tf.reduce_any(input_tensor=both_ok, axis=<NUM_LIT:0>)<EOL>return upper_bounds, lower_bounds, both_ok<EOL><DEDENT>

Returns the bounds of the slice at each stage of doubling procedure.

    Precomputes the x coordinates of the left (L) and right (R) endpoints of the
    interval `I` produced in the "doubling" algorithm [Neal 2003][1] P713. Note
    that we simultaneously compute all possible doubling values for each chain,
    for the reason that at small-medium densities, the gains from parallel
    evaluation might cause a speed-up, but this will be benchmarked against the
    while loop implementation.

    Args:
      x_initial: `tf.Tensor` of any shape and any real dtype consumable by
        `target_log_prob`. The initial points.
      target_log_prob: A callable taking a `tf.Tensor` of shape and dtype as
        `x_initial` and returning a tensor of the same shape. The log density of
        the target distribution.
      log_slice_heights: `tf.Tensor` with the same shape as `x_initial` and the
        same dtype as returned by `target_log_prob`. The log of the height of the
        slice for each chain. The values must be bounded above by
        `target_log_prob(x_initial)`.
      max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
        doublings to consider.
      step_size: `tf.Tensor` with same dtype as and shape compatible with
        `x_initial`. The size of the initial interval.
      seed: (Optional) positive int. The random seed. If None, no seed is set.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'find_slice_bounds').

    Returns:
      upper_bounds: A tensor of same shape and dtype as `x_initial`. Slice upper
        bounds for each chain.
      lower_bounds: A tensor of same shape and dtype as `x_initial`. Slice lower
        bounds for each chain.
      both_ok: A tensor of shape `x_initial` and boolean dtype. Indicates if both
        the chosen upper and lower bound lie outside of the slice.

    #### References

    [1]: Radford M. Neal. Slice Sampling. The Annals of Statistics. 2003, Vol 31,
         No. 3 , 705-767.
         https://projecteuclid.org/download/pdf_1/euclid.aos/1056562461