with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>seed = tfd.SeedStream(seed, salt='<STR_LIT>')<EOL>if initial_step_size is None or initial_state is None:<EOL><INDENT>def make_variational():<EOL><INDENT>return build_factored_variational_loss(<EOL>model, observed_time_series,<EOL>init_batch_shape=chain_batch_shape, seed=seed())<EOL><DEDENT>make_variational = tf.compat.v1.make_template('<STR_LIT>',<EOL>make_variational)<EOL>_, variational_distributions = make_variational()<EOL>minimize_op = _minimize_in_graph(<EOL>build_loss_fn=lambda: make_variational()[<NUM_LIT:0>],  <EOL>num_steps=num_variational_steps,<EOL>optimizer=variational_optimizer)<EOL>with tf.control_dependencies([minimize_op]):<EOL><INDENT>if initial_state is None:<EOL><INDENT>initial_state = [tf.stop_gradient(d.sample())<EOL>for d in variational_distributions.values()]<EOL><DEDENT>if initial_step_size is None:<EOL><INDENT>initial_step_size = [<EOL>transformed_q.distribution.stddev()<EOL>for transformed_q in variational_distributions.values()]<EOL><DEDENT><DEDENT><DEDENT>observed_time_series = sts_util.pad_batch_dimension_for_multiple_chains(<EOL>observed_time_series, model, chain_batch_shape=chain_batch_shape)<EOL>samples, kernel_results = mcmc.sample_chain(<EOL>num_results=num_results,<EOL>current_state=initial_state,<EOL>num_burnin_steps=num_warmup_steps,<EOL>kernel=mcmc.SimpleStepSizeAdaptation(<EOL>inner_kernel=mcmc.TransformedTransitionKernel(<EOL>inner_kernel=mcmc.HamiltonianMonteCarlo(<EOL>target_log_prob_fn=model.joint_log_prob(<EOL>observed_time_series),<EOL>step_size=initial_step_size,<EOL>num_leapfrog_steps=num_leapfrog_steps,<EOL>state_gradients_are_stopped=True,<EOL>seed=seed()),<EOL>bijector=[param.bijector for param in model.parameters]),<EOL>num_adaptation_steps=int(num_warmup_steps * <NUM_LIT>),<EOL>adaptation_rate=tf.convert_to_tensor(<EOL>value=<NUM_LIT:0.1>, dtype=initial_state[<NUM_LIT:0>].dtype)),<EOL>parallel_iterations=<NUM_LIT:1> if seed is not None else <NUM_LIT:10>)<EOL>return samples, kernel_results<EOL><DEDENT>

Draw posterior samples using Hamiltonian Monte Carlo (HMC).

    Markov chain Monte Carlo (MCMC) methods are considered the gold standard of
    Bayesian inference; under suitable conditions and in the limit of infinitely
    many draws they generate samples from the true posterior distribution. HMC [1]
    uses gradients of the model's log-density function to propose samples,
    allowing it to exploit posterior geometry. However, it is computationally more
    expensive than variational inference and relatively sensitive to tuning.

    This method attempts to provide a sensible default approach for fitting
    StructuralTimeSeries models using HMC. It first runs variational inference as
    a fast posterior approximation, and initializes the HMC sampler from the
    variational posterior, using the posterior standard deviations to set
    per-variable step sizes (equivalently, a diagonal mass matrix). During the
    warmup phase, it adapts the step size to target an acceptance rate of 0.75,
    which is thought to be in the desirable range for optimal mixing [2].

    Args:
      model: An instance of `StructuralTimeSeries` representing a
        time-series model. This represents a joint distribution over
        time-series and their parameters with batch shape `[b1, ..., bN]`.
      observed_time_series: `float` `Tensor` of shape
        `concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
        `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
        dimension may (optionally) be omitted if `num_timesteps > 1`. May
        optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
        a mask `Tensor` to specify timesteps with missing observations.
      num_results: Integer number of Markov chain draws.
        Default value: `100`.
      num_warmup_steps: Integer number of steps to take before starting to
        collect results. The warmup steps are also used to adapt the step size
        towards a target acceptance rate of 0.75.
        Default value: `50`.
      num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
        for. Total progress per HMC step is roughly proportional to
        `step_size * num_leapfrog_steps`.
        Default value: `15`.
      initial_state: Optional Python `list` of `Tensor`s, one for each model
        parameter, representing the initial state(s) of the Markov chain(s). These
        should have shape `concat([chain_batch_shape, param.prior.batch_shape,
        param.prior.event_shape])`. If `None`, the initial state is set
        automatically using a sample from a variational posterior.
        Default value: `None`.
      initial_step_size: Python `list` of `Tensor`s, one for each model parameter,
        representing the step size for the leapfrog integrator. Must
        broadcast with the shape of `initial_state`. Larger step sizes lead to
        faster progress, but too-large step sizes make rejection exponentially
        more likely. If `None`, the step size is set automatically using the
        standard deviation of a variational posterior.
        Default value: `None`.
      chain_batch_shape: Batch shape (Python `tuple`, `list`, or `int`) of chains
        to run in parallel.
        Default value: `[]` (i.e., a single chain).
      num_variational_steps: Python `int` number of steps to run the variational
        optimization to determine the initial state and step sizes.
        Default value: `150`.
      variational_optimizer: Optional `tf.train.Optimizer` instance to use in
        the variational optimization. If `None`, defaults to
        `tf.train.AdamOptimizer(0.1)`.
        Default value: `None`.
      seed: Python integer to seed the random number generator.
      name: Python `str` name prefixed to ops created by this function.
        Default value: `None` (i.e., 'fit_with_hmc').

    Returns:
      samples: Python `list` of `Tensors` representing posterior samples of model
        parameters, with shapes `[concat([[num_results], chain_batch_shape,
        param.prior.batch_shape, param.prior.event_shape]) for param in
        model.parameters]`.
      kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
        `Tensor`s representing internal calculations made within the HMC sampler.

    #### Examples

    Assume we've built a structural time-series model:

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

    To draw posterior samples using HMC under default settings:

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

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      samples_, kernel_results_ = sess.run((samples, kernel_results))

    print("acceptance rate: {}".format(
      np.mean(kernel_results_.inner_results.is_accepted, axis=0)))
    print("posterior means: {}".format(
      {param.name: np.mean(param_draws, axis=0)
       for (param, param_draws) in zip(model.parameters, samples_)}))
    ```

    We can also run multiple chains. This may help diagnose convergence issues
    and allows us to exploit vectorization to draw samples more quickly, although
    warmup still requires the same number of sequential steps.

    ```python
    from matplotlib import pylab as plt

    samples, kernel_results = tfp.sts.fit_with_hmc(
      model, observed_time_series, chain_batch_shape=[10])

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      samples_, kernel_results_ = sess.run((samples, kernel_results))

    print("acceptance rate: {}".format(
      np.mean(kernel_results_.inner_results.inner_results.is_accepted, axis=0)))

    # Plot the sampled traces for each parameter. If the chains have mixed, their
    # traces should all cover the same region of state space, frequently crossing
    # over each other.
    for (param, param_draws) in zip(model.parameters, samples_):
      if param.prior.event_shape.ndims > 0:
        print("Only plotting traces for scalar parameters, skipping {}".format(
          param.name))
        continue
      plt.figure(figsize=[10, 4])
      plt.title(param.name)
      plt.plot(param_draws)
      plt.ylabel(param.name)
      plt.xlabel("HMC step")

    # Combining the samples from multiple chains into a single dimension allows
    # us to easily pass sampled parameters to downstream forecasting methods.
    combined_samples_ = [np.reshape(param_draws,
                                    [-1] + list(param_draws.shape[2:]))
                         for param_draws in samples_]
    ```

    For greater flexibility, you may prefer to implement your own sampler using
    the TensorFlow Probability primitives in `tfp.mcmc`. The following recipe
    constructs a basic HMC sampler, using a `TransformedTransitionKernel` to
    incorporate constraints on the parameter space.

    ```python
    transformed_hmc_kernel = mcmc.TransformedTransitionKernel(
        inner_kernel=mcmc.SimpleStepSizeAdaptation(
            inner_kernel=mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=model.joint_log_prob(observed_time_series),
                step_size=step_size,
                num_leapfrog_steps=num_leapfrog_steps,
                state_gradients_are_stopped=True,
                seed=seed),
            num_adaptation_steps = int(0.8 * num_warmup_steps)),
        bijector=[param.bijector for param in model.parameters])

    # Initialize from a Uniform[-2, 2] distribution in unconstrained space.
    initial_state = [tfp.sts.sample_uniform_initial_state(
      param, return_constrained=True) for param in model.parameters]

    samples, kernel_results = tfp.mcmc.sample_chain(
      kernel=transformed_hmc_kernel,
      num_results=num_results,
      current_state=initial_state,
      num_burnin_steps=num_warmup_steps)
    ```

    #### References

    [1]: Radford Neal. MCMC Using Hamiltonian Dynamics. _Handbook of Markov Chain
         Monte Carlo_, 2011. https://arxiv.org/abs/1206.1901
    [2]  M.J. Betancourt, Simon Byrne, and Mark Girolami. Optimizing The
         Integrator Step Size for Hamiltonian Monte Carlo.
         https://arxiv.org/abs/1411.6669


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

Convert a numpy array to a TF placeholder.

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

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


return self._observation_noise_scale<EOL>

Standard deviation of the observation noise.


return self._initial_state_prior<EOL>

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


return self._components<EOL>

List of component `StructuralTimeSeries` models.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>if observed_time_series is not None:<EOL><INDENT>observed_mean, observed_stddev, _ = (<EOL>sts_util.empirical_statistics(observed_time_series))<EOL><DEDENT>else:<EOL><INDENT>observed_mean, observed_stddev = <NUM_LIT:0.>, <NUM_LIT:1.><EOL><DEDENT>if observation_noise_scale_prior is None:<EOL><INDENT>observation_noise_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev), scale=<NUM_LIT>)<EOL><DEDENT>if constant_offset is None:<EOL><INDENT>constant_offset = observed_mean<EOL><DEDENT>component_names = [c.name for c in components]<EOL>if len(component_names) != len(set(component_names)):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(component_names))<EOL><DEDENT>components_by_name = collections.OrderedDict(<EOL>[(c.name, c) for c in components])<EOL>parameters = [<EOL>Parameter('<STR_LIT>', observation_noise_scale_prior,<EOL>tfb.Softplus()),<EOL>] + [Parameter(name='<STR_LIT>'.format(component.name, parameter.name),<EOL>prior=parameter.prior,<EOL>bijector=parameter.bijector)<EOL>for component in components for parameter in component.parameters]<EOL>self._components = components<EOL>self._components_by_name = components_by_name<EOL>self._constant_offset = constant_offset<EOL>super(Sum, self).__init__(<EOL>parameters=parameters,<EOL>latent_size=sum(<EOL>[component.latent_size for component in components]),<EOL>name=name)<EOL><DEDENT>

Specify a structural time series model representing a sum of components.

        Args:
          components: Python `list` of one or more StructuralTimeSeries instances.
            These must have unique names.
          constant_offset: optional scalar `float` `Tensor`, or batch of scalars,
            specifying a constant value added to the sum of outputs from the
            component models. This allows the components to model the shifted series
            `observed_time_series - constant_offset`. If `None`, this is set to the
            mean of the provided `observed_time_series`.
            Default value: `None`.
          observation_noise_scale_prior: optional `tfd.Distribution` instance
            specifying a prior on `observation_noise_scale`. If `None`, a heuristic
            default prior is constructed based on the provided
            `observed_time_series`.
            Default value: `None`.
          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. This is
            used to set the constant offset, if not provided, and to construct a
            default heuristic `observation_noise_scale_prior` if not provided. 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` name of this model component; used as `name_scope`
            for ops created by this class.
            Default value: 'Sum'.

        Raises:
          ValueError: if components do not have unique names.


return self._num_seasons<EOL>


return self._num_steps_per_season<EOL>


return self._constrain_mean_effect_to_zero<EOL>

Whether to constrain the mean effect to zero.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>values=[drift_scale, observation_noise_scale]) as name:<EOL><INDENT>dtype = initial_state_prior.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_steps_per_season = np.squeeze(np.asarray(num_steps_per_season))<EOL>if num_steps_per_season.ndim == <NUM_LIT:0>:  <EOL><INDENT>num_steps_per_season = np.tile(num_steps_per_season, num_seasons)<EOL><DEDENT>elif ((num_steps_per_season.ndim <= <NUM_LIT:2>)  <EOL>and (num_steps_per_season.shape[-<NUM_LIT:1>] != num_seasons)):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>num_seasons, num_steps_per_season.shape))<EOL><DEDENT>is_last_day_of_season = build_is_last_day_of_season(num_steps_per_season)<EOL>[<EOL>effects_to_residuals,<EOL>residuals_to_effects<EOL>] = build_effects_to_residuals_matrix(num_seasons, dtype=dtype)<EOL>seasonal_transition_matrix = build_seasonal_transition_matrix(<EOL>num_seasons=num_seasons,<EOL>is_last_day_of_season=is_last_day_of_season,<EOL>dtype=dtype,<EOL>basis_change_matrix=effects_to_residuals,<EOL>basis_change_matrix_inv=residuals_to_effects)<EOL>seasonal_transition_noise = build_constrained_seasonal_transition_noise(<EOL>drift_scale, num_seasons, is_last_day_of_season)<EOL>observation_matrix = tf.concat(<EOL>[tf.ones([<NUM_LIT:1>, <NUM_LIT:1>], dtype=dtype),<EOL>tf.zeros([<NUM_LIT:1>, num_seasons-<NUM_LIT:1>], dtype=dtype)], axis=-<NUM_LIT:1>)<EOL>observation_matrix = tf.matmul(observation_matrix, residuals_to_effects)<EOL>self._drift_scale = drift_scale<EOL>self._observation_noise_scale = observation_noise_scale<EOL>self._num_seasons = num_seasons<EOL>self._num_steps_per_season = num_steps_per_season<EOL>super(ConstrainedSeasonalStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=seasonal_transition_matrix,<EOL>transition_noise=seasonal_transition_noise,<EOL>observation_matrix=observation_matrix,<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=observation_noise_scale[..., tf.newaxis]),<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>

Build a seasonal effect state space model with a zero-sum constraint.

        {seasonal_init_args}


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


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.


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.


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


return self._name<EOL>


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.


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


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


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.


return self._parameters<EOL>

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


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


def slice_sampler_one_dim(target_log_prob, x_initial, step_size=<NUM_LIT>,<EOL>max_doublings=<NUM_LIT:30>, seed=None, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[x_initial, step_size, max_doublings]):<EOL><INDENT>x_initial = tf.convert_to_tensor(value=x_initial)<EOL>dtype = x_initial.dtype.base_dtype<EOL>log_slice_heights = target_log_prob(x_initial) - tf.random.gamma(<EOL>tf.shape(input=x_initial), alpha=<NUM_LIT:1>, dtype=dtype, seed=seed)<EOL>upper_bounds, lower_bounds, bounds_satisfied = slice_bounds_by_doubling(<EOL>x_initial, target_log_prob, log_slice_heights, max_doublings, step_size,<EOL>seed=seed)<EOL>retval = _sample_with_shrinkage(x_initial, target_log_prob=target_log_prob,<EOL>log_slice_heights=log_slice_heights,<EOL>step_size=step_size,<EOL>lower_bounds=lower_bounds,<EOL>upper_bounds=upper_bounds, seed=seed)<EOL>return (retval, target_log_prob(retval), bounds_satisfied,<EOL>upper_bounds, lower_bounds)<EOL><DEDENT>

For a given x position in each Markov chain, returns the next x.

    Applies the one dimensional slice sampling algorithm as defined in Neal (2003)
    to an input tensor x of shape (num_chains,) where num_chains is the number of
    simulataneous Markov chains, and returns the next tensor x of shape
    (num_chains,) when these chains are evolved by the slice sampling algorithm.

    Args:
      target_log_prob: Callable accepting a tensor like `x_initial` and returning
        a tensor containing the log density at that point of the same shape.
      x_initial: A tensor of any shape. The initial positions of the chains. This
        function assumes that all the dimensions of `x_initial` are batch
        dimensions (i.e. the event shape is `[]`).
      step_size: A tensor of shape and dtype compatible with `x_initial`. The min
        interval size in the doubling algorithm.
      max_doublings: Scalar tensor of dtype `tf.int32`. The maximum number of
        doublings to try to find the slice bounds.
      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:
      retval: A tensor of the same shape and dtype as `x_initial`. The next state
        of the Markov chain.
      next_target_log_prob: The target log density evaluated at `retval`.
      bounds_satisfied: A tensor of bool dtype and shape batch dimensions.
      upper_bounds: Tensor of the same shape and dtype as `x_initial`. The upper
        bounds for the slice found.
      lower_bounds: Tensor of the same shape and dtype as `x_initial`. The lower
        bounds for the slice found.


<EOL>return tfd.Independent(tfd.Normal(*args, **kwargs),<EOL>reinterpreted_batch_ndims=<NUM_LIT:1>)<EOL>

Convenience function to efficiently construct a MultivariateNormalDiag.


num_schools = <NUM_LIT:8><EOL>treatment_effects = tf.constant(<EOL>[<NUM_LIT>, <NUM_LIT:8>, -<NUM_LIT:3>, <NUM_LIT:7>, -<NUM_LIT:1>, <NUM_LIT:1>, <NUM_LIT>, <NUM_LIT:12>],<EOL>dtype=np.float32,<EOL>name='<STR_LIT>')<EOL>treatment_stddevs = tf.constant(<EOL>[<NUM_LIT:15>, <NUM_LIT:10>, <NUM_LIT:16>, <NUM_LIT:11>, <NUM_LIT:9>, <NUM_LIT:11>, <NUM_LIT:10>, <NUM_LIT>],<EOL>dtype=np.float32,<EOL>name='<STR_LIT>')<EOL>def unnormalized_posterior_log_prob(<EOL>avg_effect, avg_stddev, school_effects_standard):<EOL><INDENT>"""<STR_LIT>"""<EOL>return eight_schools_joint_log_prob(<EOL>treatment_effects, treatment_stddevs,<EOL>avg_effect, avg_stddev, school_effects_standard)<EOL><DEDENT>if tf.executing_eagerly():<EOL><INDENT>sample_chain = tf.function(tfp.mcmc.sample_chain)<EOL><DEDENT>else:<EOL><INDENT>sample_chain = tfp.mcmc.sample_chain<EOL><DEDENT>def computation():<EOL><INDENT>"""<STR_LIT>"""<EOL>_, kernel_results = sample_chain(<EOL>num_results=num_results,<EOL>num_burnin_steps=num_burnin_steps,<EOL>current_state=(<EOL>tf.zeros([], name='<STR_LIT>'),<EOL>tf.zeros([], name='<STR_LIT>'),<EOL>tf.ones([num_schools], name='<STR_LIT>'),<EOL>),<EOL>kernel=tfp.mcmc.HamiltonianMonteCarlo(<EOL>target_log_prob_fn=unnormalized_posterior_log_prob,<EOL>step_size=step_size,<EOL>num_leapfrog_steps=num_leapfrog_steps))<EOL>return kernel_results.is_accepted<EOL><DEDENT>is_accepted_tensor = computation()<EOL>if not tf.executing_eagerly():<EOL><INDENT>session = tf.compat.v1.Session()<EOL>session.run(is_accepted_tensor)<EOL><DEDENT>start_time = time.time()<EOL>if tf.executing_eagerly():<EOL><INDENT>is_accepted = computation()<EOL><DEDENT>else:<EOL><INDENT>is_accepted = session.run(is_accepted_tensor)<EOL><DEDENT>wall_time = time.time() - start_time<EOL>num_accepted = np.sum(is_accepted)<EOL>acceptance_rate = np.float32(num_accepted) / np.float32(num_results)<EOL>return dict(<EOL>iters=(num_results + num_burnin_steps) * num_leapfrog_steps,<EOL>extras={'<STR_LIT>': acceptance_rate},<EOL>wall_time=wall_time)<EOL>

Runs HMC on the eight-schools unnormalized posterior.


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


initial_draws = np.random.normal(size=[<NUM_LIT:30>, <NUM_LIT:2>, <NUM_LIT:1>])<EOL>x_ph = tf.compat.v1.placeholder_with_default(<EOL>np.float32(initial_draws), shape=initial_draws.shape, name='<STR_LIT>')<EOL>self._ais_gets_correct_log_normalizer(x_ph, independent_chain_ndims)<EOL>

Tests that AIS yields reasonable estimates of normalizers.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [<EOL>random_draw_parts, state_parts, drift_parts, step_size_parts,<EOL>volatility_parts<EOL>]):<EOL><INDENT>proposed_state_parts = []<EOL>for random_draw, state, drift, step_size, volatility in zip(<EOL>random_draw_parts,<EOL>state_parts,<EOL>drift_parts,<EOL>step_size_parts,<EOL>volatility_parts):<EOL><INDENT>proposal = state + drift + volatility * tf.sqrt(step_size) * random_draw<EOL>proposed_state_parts.append(proposal)<EOL><DEDENT>return proposed_state_parts<EOL><DEDENT>

Applies one step of Euler-Maruyama method.

    Generates proposal of the form:

    ```python
    tfd.Normal(loc=state_parts + _get_drift(state_parts, ...),
               scale=tf.sqrt(step_size * volatility_fn(current_state)))
    ```

    `_get_drift(state_parts, ..)` is a diffusion drift value at `state_parts`.

    Args:
      random_draw_parts: Python `list` of `Tensor`s containing the value(s) of the
        random perturbation variable(s). Must broadcast with the shape of
        `state_parts`.
      state_parts: Python `list` of `Tensor`s representing the current
        state(s) of the Markov chain(s).
      drift_parts: Python `list` of `Tensor`s representing value of the drift
        `_get_drift(*state_parts, ..)`. Must broadcast with the shape of
        `state_parts`.
      step_size_parts: Python `list` of `Tensor`s representing the step size for
        the Euler-Maruyama method. Must broadcast with the shape of
        `state_parts`.  Larger step sizes lead to faster progress, but
        too-large step sizes make rejection exponentially more likely. When
        possible, it's often helpful to match per-variable step sizes to the
        standard deviations of the target distribution in each variable.
      volatility_parts: Python `list` of `Tensor`s representing the value of
        `volatility_fn(*state_parts)`. Must broadcast with the shape of
        `state_parts`.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'mala_euler_method').

    Returns:
      proposed_state_parts: Tensor or Python list of `Tensor`s representing the
        state(s) of the Markov chain(s) at each result step. Has same shape as
        input `current_state_parts`.


return self._parameters<EOL>


return self._impl.bootstrap_results(init_state)<EOL>

Creates initial `previous_kernel_results` using a supplied `state`.


raise NotImplementedError()<EOL>

Returns `True` if Markov chain converges to specified distribution.

        `TransitionKernel`s which are "uncalibrated" are often calibrated by
        composing them with the `tfp.mcmc.MetropolisHastings` `TransitionKernel`.


@abc.abstractmethod<EOL><INDENT>def one_step(self, current_state, previous_kernel_results):<DEDENT>

raise NotImplementedError()<EOL>

Takes one step of the TransitionKernel.

        Must be overridden by subclasses.

        Args:
          current_state: `Tensor` or Python `list` of `Tensor`s representing the
            current state(s) of the Markov chain(s).
          previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
            `list` of `Tensor`s representing internal calculations made within the
            previous call to this function (or as returned by `bootstrap_results`).

        Returns:
          next_state: `Tensor` or Python `list` of `Tensor`s representing the
            next state(s) of the Markov chain(s).
          kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
            `Tensor`s representing internal calculations made within this function.


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

Helper which uses graph seed if using TFE.


if not tf.executing_eagerly():<EOL><INDENT>tf.compat.v1.reset_default_graph()<EOL><DEDENT>count_data = tf.cast(<EOL>tf.concat(<EOL>[tfd.Poisson(rate=<NUM_LIT>).sample(<NUM_LIT>),<EOL>tfd.Poisson(rate=<NUM_LIT>).sample(<NUM_LIT>)],<EOL>axis=<NUM_LIT:0>),<EOL>dtype=tf.float32)<EOL>if tf.executing_eagerly():<EOL><INDENT>count_data = count_data.numpy()<EOL><DEDENT>else:<EOL><INDENT>with tf.compat.v1.Session():<EOL><INDENT>count_data = count_data.eval()<EOL><DEDENT><DEDENT>def unnormalized_log_posterior(lambda1, lambda2, tau):<EOL><INDENT>return text_messages_joint_log_prob(count_data, lambda1, lambda2, tau)<EOL><DEDENT>if tf.executing_eagerly():<EOL><INDENT>sample_chain = tf.function(tfp.mcmc.sample_chain)<EOL><DEDENT>else:<EOL><INDENT>sample_chain = tfp.mcmc.sample_chain<EOL><DEDENT>step_size = tf.compat.v2.Variable(<EOL>name='<STR_LIT>',<EOL>initial_value=tf.constant(<NUM_LIT>, dtype=tf.float32),<EOL>trainable=False)<EOL>def computation():<EOL><INDENT>"""<STR_LIT>"""<EOL>initial_chain_state = [<EOL>tf.constant(count_data.mean(), name='<STR_LIT>'),<EOL>tf.constant(count_data.mean(), name='<STR_LIT>'),<EOL>tf.constant(<NUM_LIT:0.5>, name='<STR_LIT>'),<EOL>]<EOL>unconstraining_bijectors = [<EOL>tfp.bijectors.Exp(),       <EOL>tfp.bijectors.Exp(),       <EOL>tfp.bijectors.Sigmoid(),   <EOL>]<EOL>_, kernel_results = sample_chain(<EOL>num_results=num_results,<EOL>num_burnin_steps=num_burnin_steps,<EOL>current_state=initial_chain_state,<EOL>kernel=tfp.mcmc.TransformedTransitionKernel(<EOL>inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(<EOL>target_log_prob_fn=unnormalized_log_posterior,<EOL>num_leapfrog_steps=num_leapfrog_steps,<EOL>step_size=step_size,<EOL>step_size_update_fn=<EOL>tfp.mcmc.make_simple_step_size_update_policy(num_burnin_steps),<EOL>state_gradients_are_stopped=True),<EOL>bijector=unconstraining_bijectors))<EOL>return kernel_results.inner_results.is_accepted<EOL><DEDENT>is_accepted_tensor = computation()<EOL>if not tf.executing_eagerly():<EOL><INDENT>session = tf.compat.v1.Session()<EOL>session.run(tf.compat.v1.global_variables_initializer())<EOL>session.run(is_accepted_tensor)<EOL><DEDENT>start_time = time.time()<EOL>if tf.executing_eagerly():<EOL><INDENT>is_accepted = computation()<EOL><DEDENT>else:<EOL><INDENT>is_accepted = session.run(is_accepted_tensor)<EOL><DEDENT>wall_time = time.time() - start_time<EOL>num_accepted = np.sum(is_accepted)<EOL>acceptance_rate = np.float32(num_accepted) / np.float32(num_results)<EOL>return dict(<EOL>iters=(num_results + num_burnin_steps) * num_leapfrog_steps,<EOL>extras={'<STR_LIT>': acceptance_rate},<EOL>wall_time=wall_time)<EOL>

Runs HMC on the text-messages unnormalized posterior.


def text_messages_joint_log_prob(count_data, lambda_1, lambda_2, tau):

alpha = (<NUM_LIT:1.> / tf.reduce_mean(input_tensor=count_data))<EOL>rv_lambda = tfd.Exponential(rate=alpha)<EOL>rv_tau = tfd.Uniform()<EOL>lambda_ = tf.gather(<EOL>[lambda_1, lambda_2],<EOL>indices=tf.cast(<EOL>tau * tf.cast(tf.size(input=count_data), dtype=tf.float32) <= tf.cast(<EOL>tf.range(tf.size(input=count_data)), dtype=tf.float32),<EOL>dtype=tf.int32))<EOL>rv_observation = tfd.Poisson(rate=lambda_)<EOL>return (rv_lambda.log_prob(lambda_1) + rv_lambda.log_prob(lambda_2) +<EOL>rv_tau.log_prob(tau) +<EOL>tf.reduce_sum(input_tensor=rv_observation.log_prob(count_data)))<EOL>


state_shape = sample_shape + independent_chain_shape + other_shape<EOL>state_ = rng.randn(*state_shape).astype(dtype)<EOL>if other_shape:<EOL><INDENT>state_ *= rng.rand(*other_shape).astype(dtype)<EOL><DEDENT>self.check_results(state_, independent_chain_shape, should_pass=True)<EOL>


def _get_exchanged_states(self, old_states, exchange_proposed,<EOL>exchange_proposed_n, sampled_replica_states,<EOL>sampled_replica_results):

with tf.compat.v1.name_scope('<STR_LIT>'):<EOL><INDENT>target_log_probs = []<EOL>for replica in range(self.num_replica):<EOL><INDENT>replica_log_prob = _get_field(sampled_replica_results[replica],<EOL>'<STR_LIT>')<EOL>inverse_temp = self.inverse_temperatures[replica]<EOL>target_log_probs.append(replica_log_prob / inverse_temp)<EOL><DEDENT>target_log_probs = tf.stack(target_log_probs, axis=<NUM_LIT:0>)<EOL>dtype = target_log_probs.dtype<EOL>num_state_parts = len(sampled_replica_states[<NUM_LIT:0>])<EOL>exchanged_states = [<EOL>tf.TensorArray(<EOL>dtype,<EOL>size=self.num_replica,<EOL>dynamic_size=False,<EOL>tensor_array_name='<STR_LIT>',<EOL>element_shape=sampled_replica_states[<NUM_LIT:0>][k].shape)<EOL>for k in range(num_state_parts)<EOL>]<EOL>sample_shape = tf.concat(<EOL>([self.num_replica // <NUM_LIT:2>], tf.shape(input=target_log_probs)[<NUM_LIT:1>:]),<EOL>axis=<NUM_LIT:0>)<EOL>log_uniforms = tf.math.log(<EOL>tf.random.uniform(<EOL>shape=sample_shape, dtype=dtype, seed=self._seed_stream()))<EOL>def _swap(is_exchange_accepted, x, y):<EOL><INDENT>"""<STR_LIT>"""<EOL>with tf.compat.v1.name_scope('<STR_LIT>'):<EOL><INDENT>new_x = mcmc_util.choose(is_exchange_accepted, y, x)<EOL>new_y = mcmc_util.choose(is_exchange_accepted, x, y)<EOL><DEDENT>return new_x, new_y<EOL><DEDENT>def cond(i, unused_exchanged_states):<EOL><INDENT>return i < exchange_proposed_n<EOL><DEDENT>def body(i, exchanged_states):<EOL><INDENT>"""<STR_LIT>"""<EOL>m, n = tf.unstack(exchange_proposed[i])<EOL>temp_diff = self.inverse_temperatures[m] - self.inverse_temperatures[n]<EOL>log_accept_ratio = mcmc_util.safe_sum(<EOL>[-temp_diff * target_log_probs[m], temp_diff * target_log_probs[n]])<EOL>is_exchange_accepted = log_uniforms[i] < log_accept_ratio<EOL>for k in range(num_state_parts):<EOL><INDENT>new_m, new_n = _swap(is_exchange_accepted, old_states[k].read(m),<EOL>old_states[k].read(n))<EOL>exchanged_states[k] = exchanged_states[k].write(m, new_m)<EOL>exchanged_states[k] = exchanged_states[k].write(n, new_n)<EOL><DEDENT>return i + <NUM_LIT:1>, exchanged_states<EOL><DEDENT>return tf.while_loop(<EOL>cond=cond, body=body, loop_vars=[tf.constant(<NUM_LIT:0>),<EOL>exchanged_states])[<NUM_LIT:1>]<EOL><DEDENT>

Get list of TensorArrays holding exchanged states, and zeros.


def one_step(self, current_state, previous_kernel_results):

<EOL>with tf.compat.v1.name_scope(<EOL>name=mcmc_util.make_name(self.name, '<STR_LIT>', '<STR_LIT>'),<EOL>values=[current_state, previous_kernel_results]):<EOL><INDENT>sampled_replica_states, sampled_replica_results = zip(*[<EOL>rk.one_step(previous_kernel_results.replica_states[i],<EOL>previous_kernel_results.replica_results[i])<EOL>for i, rk in enumerate(self.replica_kernels)<EOL>])<EOL>sampled_replica_states = list(sampled_replica_states)<EOL>sampled_replica_results = list(sampled_replica_results)<EOL>states_are_lists = mcmc_util.is_list_like(sampled_replica_states[<NUM_LIT:0>])<EOL>if not states_are_lists:<EOL><INDENT>sampled_replica_states = [[s] for s in sampled_replica_states]<EOL><DEDENT>num_state_parts = len(sampled_replica_states[<NUM_LIT:0>])<EOL>dtype = sampled_replica_states[<NUM_LIT:0>][<NUM_LIT:0>].dtype<EOL>old_states = [<EOL>tf.TensorArray(<EOL>dtype,<EOL>size=self.num_replica,<EOL>dynamic_size=False,<EOL>clear_after_read=False,<EOL>tensor_array_name='<STR_LIT>',<EOL>element_shape=sampled_replica_states[<NUM_LIT:0>][k].shape)<EOL>for k in range(num_state_parts)<EOL>]<EOL>for k in range(num_state_parts):<EOL><INDENT>for i in range(self.num_replica):<EOL><INDENT>old_states[k] = old_states[k].write(i, sampled_replica_states[i][k])<EOL><DEDENT><DEDENT>exchange_proposed = self.exchange_proposed_fn(<EOL>self.num_replica, seed=self._seed_stream())<EOL>exchange_proposed_n = tf.shape(input=exchange_proposed)[<NUM_LIT:0>]<EOL>exchanged_states = self._get_exchanged_states(<EOL>old_states, exchange_proposed, exchange_proposed_n,<EOL>sampled_replica_states, sampled_replica_results)<EOL>no_exchange_proposed, _ = tf.compat.v1.setdiff1d(<EOL>tf.range(self.num_replica), tf.reshape(exchange_proposed, [-<NUM_LIT:1>]))<EOL>exchanged_states = self._insert_old_states_where_no_exchange_was_proposed(<EOL>no_exchange_proposed, old_states, exchanged_states)<EOL>next_replica_states = []<EOL>for i in range(self.num_replica):<EOL><INDENT>next_replica_states_i = []<EOL>for k in range(num_state_parts):<EOL><INDENT>next_replica_states_i.append(exchanged_states[k].read(i))<EOL><DEDENT>next_replica_states.append(next_replica_states_i)<EOL><DEDENT>if not states_are_lists:<EOL><INDENT>next_replica_states = [s[<NUM_LIT:0>] for s in next_replica_states]<EOL>sampled_replica_states = [s[<NUM_LIT:0>] for s in sampled_replica_states]<EOL><DEDENT>next_replica_results = [<EOL>rk.bootstrap_results(state)<EOL>for rk, state in zip(self.replica_kernels, next_replica_states)<EOL>]<EOL>next_state = next_replica_states[<NUM_LIT:0>]  <EOL>kernel_results = ReplicaExchangeMCKernelResults(<EOL>replica_states=next_replica_states,<EOL>replica_results=next_replica_results,<EOL>sampled_replica_states=sampled_replica_states,<EOL>sampled_replica_results=sampled_replica_results,<EOL>)<EOL>return next_state, kernel_results<EOL><DEDENT>

Takes one step of the TransitionKernel.

        Args:
          current_state: `Tensor` or Python `list` of `Tensor`s representing the
            current state(s) of the Markov chain(s).
          previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
            `list` of `Tensor`s representing internal calculations made within the
            previous call to this function (or as returned by `bootstrap_results`).

        Returns:
          next_state: `Tensor` or Python `list` of `Tensor`s representing the
            next state(s) of the Markov chain(s).
          kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
            `Tensor`s representing internal calculations made within this function.
            This inculdes replica states.


def _replica_log_prob_fn_(*x):<EOL><INDENT>return inverse_temperature * target_log_prob_fn(*x)<EOL><DEDENT>return _replica_log_prob_fn_<EOL>

Return a log probability function made considering temperature.


return self._parameters['<STR_LIT>']<EOL>

Returns the num_leapfrog_steps parameter.

        If `store_parameters_in_results` argument to the initializer was set to
        `True`, this only returns the value of the `num_leapfrog_steps` placed in
        the kernel results by the `bootstrap_results` method. The actual
        `num_leapfrog_steps` in that situation is governed by the
        `previous_kernel_results` argument to `one_step` method.

        Returns:
          num_leapfrog_steps: An integer `Tensor`.


return self._parameters<EOL>

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


inner_kernel = mcmc_util.enable_store_parameters_in_results(inner_kernel)<EOL>with tf.compat.v1.name_scope(<EOL>mcmc_util.make_name(name, '<STR_LIT>', '<STR_LIT>'),<EOL>values=[target_accept_prob, adaptation_rate,<EOL>num_adaptation_steps]) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([target_accept_prob, adaptation_rate],<EOL>tf.float32)<EOL>target_accept_prob = tf.convert_to_tensor(<EOL>value=target_accept_prob, dtype=dtype, name='<STR_LIT>')<EOL>adaptation_rate = tf.convert_to_tensor(<EOL>value=adaptation_rate, dtype=dtype, name='<STR_LIT>')<EOL>num_adaptation_steps = tf.convert_to_tensor(<EOL>value=num_adaptation_steps,<EOL>dtype=tf.int32,<EOL>name='<STR_LIT>')<EOL>target_accept_prob = _maybe_validate_target_accept_prob(<EOL>target_accept_prob, validate_args)<EOL><DEDENT>self._parameters = dict(<EOL>inner_kernel=inner_kernel,<EOL>num_adaptation_steps=num_adaptation_steps,<EOL>target_accept_prob=target_accept_prob,<EOL>adaptation_rate=adaptation_rate,<EOL>step_size_setter_fn=step_size_setter_fn,<EOL>step_size_getter_fn=step_size_getter_fn,<EOL>log_accept_prob_getter_fn=log_accept_prob_getter_fn,<EOL>name=name,<EOL>)<EOL>

Creates the step size adaptation kernel.

        The default setter_fn and the getter_fn callbacks assume that the inner
        kernel produces kernel results structurally the same as the
        `HamiltonianMonteCarlo` kernel.

        Args:
          inner_kernel: `TransitionKernel`-like object.
          num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
            during which to adjust the step size. This may be greater, less than, or
            equal to the number of burnin steps.
          target_accept_prob: A floating point `Tensor` representing desired
            acceptance probability. Must be a positive number less than 1. This can
            either be a scalar, or have shape [num_chains]. Default value: `0.75`
            (the [center of asymptotically optimal rate for HMC][1]).
          adaptation_rate: `Tensor` representing amount to scale the current
            `step_size`.
          step_size_setter_fn: A callable with the signature
            `(kernel_results, new_step_size) -> new_kernel_results` where
            `kernel_results` are the results of the `inner_kernel`, `new_step_size`
            is a `Tensor` or a nested collection of `Tensor`s with the same
            structure as returned by the `step_size_getter_fn`, and
            `new_kernel_results` are a copy of `kernel_results` with the step
            size(s) set.
          step_size_getter_fn: A callable with the signature
            `(kernel_results) -> step_size` where `kernel_results` are the results
            of the `inner_kernel`, and `step_size` is a floating point `Tensor` or a
            nested collection of such `Tensor`s.
          log_accept_prob_getter_fn: A callable with the signature
            `(kernel_results) -> log_accept_prob` where `kernel_results` are the
            results of the `inner_kernel`, and `log_accept_prob` is a floating point
            `Tensor`. `log_accept_prob` can either be a scalar, or have shape
            [num_chains]. If it's the latter, `step_size` should also have the same
            leading dimension.
          validate_args: Python `bool`. When `True` kernel parameters are checked
            for validity. When `False` invalid inputs may silently render incorrect
            outputs.
          name: Python `str` name prefixed to Ops created by this class. Default:
            'simple_step_size_adaptation'.

        #### References

        [1]: Betancourt, M. J., Byrne, S., & Girolami, M. (2014). _Optimizing The
             Integrator Step Size for Hamiltonian Monte Carlo_.
             http://arxiv.org/abs/1411.6669


return self._impl.inner_kernel.parameters<EOL>


def _prepare_args(target_log_prob_fn, state, step_size,<EOL>target_log_prob=None, maybe_expand=False,<EOL>description='<STR_LIT>'):

state_parts = list(state) if mcmc_util.is_list_like(state) else [state]<EOL>state_parts = [<EOL>tf.convert_to_tensor(value=s, name='<STR_LIT>') for s in state_parts<EOL>]<EOL>target_log_prob = _maybe_call_fn(<EOL>target_log_prob_fn,<EOL>state_parts,<EOL>target_log_prob,<EOL>description)<EOL>step_sizes = (list(step_size) if mcmc_util.is_list_like(step_size)<EOL>else [step_size])<EOL>step_sizes = [<EOL>tf.convert_to_tensor(<EOL>value=s, name='<STR_LIT>', dtype=target_log_prob.dtype)<EOL>for s in step_sizes<EOL>]<EOL>if len(step_sizes) == <NUM_LIT:1>:<EOL><INDENT>step_sizes *= len(state_parts)<EOL><DEDENT>if len(state_parts) != len(step_sizes):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>def maybe_flatten(x):<EOL><INDENT>return x if maybe_expand or mcmc_util.is_list_like(state) else x[<NUM_LIT:0>]<EOL><DEDENT>return [<EOL>maybe_flatten(state_parts),<EOL>maybe_flatten(step_sizes),<EOL>target_log_prob<EOL>]<EOL>

Processes input args to meet list-like assumptions.


self._seed_stream = distributions.SeedStream(<EOL>seed, salt='<STR_LIT>')<EOL>self._parameters = dict(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>step_size=step_size,<EOL>max_doublings=max_doublings,<EOL>seed=seed,<EOL>name=name)<EOL>

Initializes this transition kernel.

        Args:
          target_log_prob_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it is a list) and returns its
            (possibly unnormalized) log-density under the target distribution.
          step_size: Scalar or `tf.Tensor` with same dtype as and shape compatible
            with `x_initial`. The size of the initial interval.
          max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
          doublings to consider.
          seed: Python integer to seed the random number generator.
          name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., 'slice_sampler_kernel').

        Returns:
          next_state: Tensor or Python list of `Tensor`s representing the state(s)
            of the Markov chain(s) at each result step. Has same shape as
            `current_state`.
          kernel_results: `collections.namedtuple` of internal calculations used to
            advance the chain.


Runs one iteration of Slice Sampler.

        Args:
          current_state: `Tensor` or Python `list` of `Tensor`s representing the
            current state(s) of the Markov chain(s). The first `r` dimensions
            index independent chains,
            `r = tf.rank(target_log_prob_fn(*current_state))`.
          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.

        Raises:
          ValueError: if there isn't one `step_size` or a list with same length as
            `current_state`.
          TypeError: if `not target_log_prob.dtype.is_floating`.


def tridiag(d, diag_value, offdiag_value):

diag_mat = tf.eye(d) * (diag_value - offdiag_value)<EOL>three_bands = tf.linalg.band_part(tf.fill([d, d], offdiag_value), <NUM_LIT:1>, <NUM_LIT:1>)<EOL>return diag_mat + three_bands<EOL>

d x d matrix with given value on diag, and one super/sub diag.


<EOL>logu = np.float128(logu)<EOL>n = logu.shape[<NUM_LIT:0>]<EOL>u = np.exp(logu)<EOL>loogeoavg_u = []  <EOL>for j in range(n):<EOL><INDENT>loogeoavg_u.append(np.exp(np.mean(<EOL>[logu[i, ...] for i in range(n) if i != j],<EOL>axis=<NUM_LIT:0>)))<EOL><DEDENT>loogeoavg_u = np.array(loogeoavg_u)<EOL>loosum_u = []  <EOL>for j in range(n):<EOL><INDENT>loosum_u.append(np.sum(<EOL>[u[i, ...] for i in range(n) if i != j],<EOL>axis=<NUM_LIT:0>))<EOL><DEDENT>loosum_u = np.array(loosum_u)<EOL>log_sooavg_u = np.log(loosum_u + loogeoavg_u) - np.log(n)<EOL>log_avg_u = np.log(np.mean(u, axis=<NUM_LIT:0>))<EOL>return log_avg_u, log_sooavg_u<EOL>

Numpy implementation of `csiszar_vimco_helper`.


def arithmetic_geometric(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>y = tf.nn.softplus(logu) - <NUM_LIT:0.5> * logu<EOL>if self_normalized:<EOL><INDENT>y -= np.log(<NUM_LIT>).astype(logu.dtype.as_numpy_dtype)<EOL><DEDENT>return (<NUM_LIT:1.> + tf.exp(logu)) * y<EOL><DEDENT>

The Arithmetic-Geometric Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True` the Arithmetic-Geometric Csiszar-function is:

    ```none
    f(u) = (1 + u) log( (1 + u) / sqrt(u) ) - (1 + u) log(2)
    ```

    When `self_normalized = False` the `(1 + u) log(2)` term is omitted.

    Observe that as an f-Divergence, this Csiszar-function implies:

    ```none
    D_f[p, q] = KL[m, p] + KL[m, q]
    m(x) = 0.5 p(x) + 0.5 q(x)
    ```

    In a sense, this divergence is the "reverse" of the Jensen-Shannon
    f-Divergence.

    This Csiszar-function induces a symmetric f-Divergence, i.e.,
    `D_f[p, q] = D_f[q, p]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      arithmetic_geometric_of_u: `float`-like `Tensor` of the
        Csiszar-function evaluated at `u = exp(logu)`.


def modified_gan(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>y = tf.nn.softplus(logu) - logu<EOL>if self_normalized:<EOL><INDENT>y += <NUM_LIT:0.5> * tf.math.expm1(logu)<EOL><DEDENT>return y<EOL><DEDENT>

The Modified-GAN Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True` the modified-GAN (Generative/Adversarial
    Network) Csiszar-function is:

    ```none
    f(u) = log(1 + u) - log(u) + 0.5 (u - 1)
    ```

    When `self_normalized = False` the `0.5 (u - 1)` is omitted.

    The unmodified GAN Csiszar-function is identical to Jensen-Shannon (with
    `self_normalized = False`).

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      chi_square_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>return <NUM_LIT:0.5> * tf.math.expm1(logu) * logu<EOL><DEDENT>

The Jeffreys Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    The Jeffreys Csiszar-function is:

    ```none
    f(u) = 0.5 ( u log(u) - log(u) )
         = 0.5 kl_forward + 0.5 kl_reverse
         = symmetrized_csiszar_function(kl_reverse)
         = symmetrized_csiszar_function(kl_forward)
    ```

    This Csiszar-function induces a symmetric f-Divergence, i.e.,
    `D_f[p, q] = D_f[q, p]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      jeffreys_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


def kl_reverse(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>return amari_alpha(logu, alpha=<NUM_LIT:0.>, self_normalized=self_normalized)<EOL><DEDENT>

The reverse Kullback-Leibler Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True`, the KL-reverse Csiszar-function is:

    ```none
    f(u) = -log(u) + (u - 1)
    ```

    When `self_normalized = False` the `(u - 1)` term is omitted.

    Observe that as an f-Divergence, this Csiszar-function implies:

    ```none
    D_f[p, q] = KL[q, p]
    ```

    The KL is "reverse" because in maximum likelihood we think of minimizing `q`
    as in `KL[p, q]`.

    Warning: when self_normalized = True` this function makes non-log-space
    calculations and may therefore be numerically unstable for `|logu| >> 0`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      kl_reverse_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
        `u = exp(logu)`.

    Raises:
      TypeError: if `self_normalized` is `None` or a `Tensor`.


def t_power(logu, t, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu, t]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>t = tf.convert_to_tensor(value=t, dtype=logu.dtype.base_dtype, name="<STR_LIT:t>")<EOL>fu = tf.math.expm1(t * logu)<EOL>if self_normalized:<EOL><INDENT>fu -= t * tf.math.expm1(logu)<EOL><DEDENT>fu *= tf.where(tf.logical_and(<NUM_LIT:0.> < t, t < <NUM_LIT:1.>),<EOL>-tf.ones_like(t),<EOL>tf.ones_like(t))<EOL>return fu<EOL><DEDENT>

The T-Power Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True` the T-Power Csiszar-function is:

    ```none
    f(u) = s [ u**t - 1 - t(u - 1) ]
    s = { -1   0 < t < 1
        { +1   otherwise
    ```

    When `self_normalized = False` the `- t(u - 1)` term is omitted.

    This is similar to the `amari_alpha` Csiszar-function, with the associated
    divergence being the same up to factors depending only on `t`.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      t:  `Tensor` of same `dtype` as `logu` and broadcastable shape.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      t_power_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
        at `u = exp(logu)`.


def jensen_shannon(logu, self_normalized=False, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [logu]):<EOL><INDENT>logu = tf.convert_to_tensor(value=logu, name="<STR_LIT>")<EOL>npdt = logu.dtype.as_numpy_dtype<EOL>y = tf.nn.softplus(logu)<EOL>if self_normalized:<EOL><INDENT>y -= np.log(<NUM_LIT:2>).astype(npdt)<EOL><DEDENT>return tf.exp(logu) * logu - (<NUM_LIT:1.> + tf.exp(logu)) * y<EOL><DEDENT>

The Jensen-Shannon Csiszar-function in log-space.

    A Csiszar-function is a member of,

    ```none
    F = { f:R_+ to R : f convex }.
    ```

    When `self_normalized = True`, the Jensen-Shannon Csiszar-function is:

    ```none
    f(u) = u log(u) - (1 + u) log(1 + u) + (u + 1) log(2)
    ```

    When `self_normalized = False` the `(u + 1) log(2)` term is omitted.

    Observe that as an f-Divergence, this Csiszar-function implies:

    ```none
    D_f[p, q] = KL[p, m] + KL[q, m]
    m(x) = 0.5 p(x) + 0.5 q(x)
    ```

    In a sense, this divergence is the "reverse" of the Arithmetic-Geometric
    f-Divergence.

    This Csiszar-function induces a symmetric f-Divergence, i.e.,
    `D_f[p, q] = D_f[q, p]`.

    Warning: this function makes non-log-space calculations and may therefore be
    numerically unstable for `|logu| >> 0`.

    For more information, see:
      Lin, J. "Divergence measures based on the Shannon entropy." IEEE Trans.
      Inf. Th., 37, 145-151, 1991.

    Args:
      logu: `float`-like `Tensor` representing `log(u)` from above.
      self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
        `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
        when `p, q` are unnormalized measures.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      jensen_shannon_of_u: `float`-like `Tensor` of the Csiszar-function
        evaluated at `u = exp(logu)`.


sizes = [tf.cast(tf.size(input=x), dtype=tf.float64) for x in population]<EOL>seed_stream = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>force_crossover_group = distributions.Categorical(sizes).sample(<EOL>[population_size, <NUM_LIT:1>], seed=seed_stream())<EOL>recombinants = []<EOL>for i, population_part in enumerate(population):<EOL><INDENT>pop_part_flat = tf.reshape(population_part, [population_size, -<NUM_LIT:1>])<EOL>mutant_part_flat = tf.reshape(mutants[i], [population_size, -<NUM_LIT:1>])<EOL>part_size = tf.size(input=population_part) // population_size<EOL>force_crossovers = tf.one_hot(<EOL>tf.random.uniform([population_size],<EOL>minval=<NUM_LIT:0>,<EOL>maxval=part_size,<EOL>dtype=tf.int32,<EOL>seed=seed_stream()),<EOL>part_size,<EOL>on_value=True,<EOL>off_value=False,<EOL>dtype=tf.bool)  <EOL>group_mask = tf.math.equal(force_crossover_group, i)<EOL>force_crossovers &= group_mask<EOL>do_binary_crossover = tf.random.uniform(<EOL>[population_size, part_size],<EOL>dtype=crossover_prob.dtype.base_dtype,<EOL>seed=seed_stream()) < crossover_prob<EOL>do_binary_crossover |= force_crossovers<EOL>recombinant_flat = tf.where(<EOL>do_binary_crossover,<EOL>x=mutant_part_flat,<EOL>y=pop_part_flat)<EOL>recombinant = tf.reshape(recombinant_flat, tf.shape(input=population_part))<EOL>recombinants.append(recombinant)<EOL><DEDENT>return recombinants<EOL>

Performs recombination by binary crossover for the current population.

    Let v_i denote the i'th component of the member v and m_i the corresponding
    component of the mutant vector corresponding to v. Then the crossed over
    vector w_i is determined by setting w_i =
    (m_i with probability=crossover_prob else v_i). In addition, DE requires that
    at least one of the components is crossed over (otherwise we end
    up with no change). This is done by choosing on index say k randomly where
    a force crossover is performed (i.e. w_k = m_k). This is the scheme
    implemented in this function.

    Args:
      population: A Python list of `Tensor`s where each `Tensor` in the list
        must be of rank at least 1 and all the elements must have a common
        first dimension. The base population to cross over.
      population_size: A scalar integer `Tensor`. The number of elements in the
        population (i.e. size of the first dimension of any member of
        `population`).
      mutants: A Python list of `Tensor`s with the same structure as `population`.
        The mutated population.
      crossover_prob: A postive real scalar `Tensor` bounded above by 1.0. The
        probability of a crossover being performed for each axis.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.

    Returns:
      A list of `Tensor`s of the same structure, dtype and shape as `population`.
      The recombined population.


return tf.math.reduce_all(input_tensor=tf.math.is_inf(population_values))<EOL>

Checks if all the population values are NaN/infinite.


n = distribution_util.prefer_static_rank(mat)<EOL>perm = tf.range(n)<EOL>perm = tf.concat([perm[:-<NUM_LIT:2>], [perm[-<NUM_LIT:1>], perm[-<NUM_LIT:2>]]], axis=<NUM_LIT:0>)<EOL>return tf.transpose(a=mat, perm=perm)<EOL>

Transpose a possibly batched matrix.

    Args:
      mat: A `tf.Tensor` of shape `[..., n, m]`.

    Returns:
      A tensor of shape `[..., m, n]` with matching batch dimensions.


<EOL>should_update = ~next_state.converged & ~next_state.failed<EOL>gradient_delta = next_state.objective_gradient - prev_state.objective_gradient<EOL>position_delta = next_state.position - prev_state.position<EOL>normalization_factor = tf.reduce_sum(<EOL>input_tensor=gradient_delta * position_delta, axis=-<NUM_LIT:1>)<EOL>should_update = should_update & ~tf.equal(normalization_factor, <NUM_LIT:0>)<EOL>def _do_update_inv_hessian():<EOL><INDENT>next_inv_hessian = _bfgs_inv_hessian_update(<EOL>gradient_delta, position_delta, normalization_factor,<EOL>prev_state.inverse_hessian_estimate)<EOL>return bfgs_utils.update_fields(<EOL>next_state,<EOL>inverse_hessian_estimate=tf.where(should_update,<EOL>next_inv_hessian,<EOL>prev_state.inverse_hessian_estimate))<EOL><DEDENT>return prefer_static.cond(<EOL>tf.reduce_any(input_tensor=should_update),<EOL>_do_update_inv_hessian,<EOL>lambda: next_state)<EOL>

Update the BGFS state by computing the next inverse hessian estimate.


return tf.matmul(tf.expand_dims(t1, axis=-<NUM_LIT:1>), tf.expand_dims(t2, axis=-<NUM_LIT:2>))<EOL>

Computes the outer product of two possibly batched vectors.

    Args:
      t1: A `tf.Tensor` of shape `[..., n]`.
      t2: A `tf.Tensor` of shape `[..., m]`.

    Returns:
      A tensor of shape `[..., n, m]` with matching batch dimensions, let's call
      it `r`, whose components are:

      ```None
        r[..., i, j] = t1[..., i] * t2[..., j]
      ```


failed = ~is_finite(initial_left, initial_right)<EOL>needs_bisect = (initial_right.df < <NUM_LIT:0>) & (initial_right.f > f_lim)<EOL>bisect_args = _IntermediateResult(<EOL>iteration=tf.convert_to_tensor(value=<NUM_LIT:0>),<EOL>stopped=failed | ~needs_bisect,<EOL>failed=failed,<EOL>num_evals=tf.convert_to_tensor(value=<NUM_LIT:0>),<EOL>left=initial_left,<EOL>right=initial_right)<EOL>return _bisect(value_and_gradients_function, bisect_args, f_lim)<EOL>

Bisects an interval and updates to satisfy opposite slope conditions.

    Corresponds to the step U3 in [Hager and Zhang (2006)][2].

    Args:
      value_and_gradients_function: A Python callable that accepts a real scalar
        tensor and returns a namedtuple containing the value filed `f` of the
        function and its derivative value field `df` at that point.
        Alternatively, the function may representthe batching of `n` such line
        functions (e.g. projecting a single multivariate objective function along
        `n` distinct directions at once) accepting n points as input, i.e. a
        tensor of shape [n], and return a tuple of two tensors of shape [n], the
        function values and the corresponding derivatives at the input points.
      initial_left: Return value of value_and_gradients_function at the left end
        point of the current bracketing interval.
      initial_right: Return value of value_and_gradients_function at the right end
        point of the current bracketing interval.
      f_lim: real `Tensor` of shape [n]. The function value threshold for
        the approximate Wolfe conditions to be checked for each batch member.

    Returns:
      A namedtuple containing the following fields:
        iteration: An int32 scalar `Tensor`. The number of iterations performed.
          Bounded above by `max_iterations` parameter.
        stopped: A boolean scalar `Tensor`. True if the bisect algorithm
          terminated.
        failed: A scalar boolean tensor. Indicates whether the objective function
          failed to produce a finite value.
        num_evals: A scalar int32 tensor. The number of value and gradients
          function evaluations.
        left: Return value of value_and_gradients_function at the left end
          point of the bracketing interval found.
        right: Return value of value_and_gradients_function at the right end
          point of the bracketing interval found.


if isinstance(tval, tf.Tensor):<EOL><INDENT>return tf.where(cond, tval, fval)<EOL><DEDENT>elif isinstance(tval, tuple):<EOL><INDENT>cls = type(tval)<EOL>return cls(*(val_where(cond, t, f) for t, f in zip(tval, fval)))<EOL><DEDENT>else:<EOL><INDENT>raise Exception(TypeError)<EOL><DEDENT>


def _bisect(value_and_gradients_function, initial_args, f_lim):

def _loop_cond(curr):<EOL><INDENT>return ~tf.reduce_all(input_tensor=curr.stopped)<EOL><DEDENT>def _loop_body(curr):<EOL><INDENT>"""<STR_LIT>"""<EOL>mid = value_and_gradients_function((curr.left.x + curr.right.x) / <NUM_LIT:2>)<EOL>failed = (curr.failed | ~is_finite(mid) |<EOL>tf.equal(mid.x, curr.left.x) | tf.equal(mid.x, curr.right.x))<EOL>to_update = ~(curr.stopped | failed)<EOL>update_left = (mid.df < <NUM_LIT:0>) & (mid.f <= f_lim)<EOL>left = val_where(to_update & update_left, mid, curr.left)<EOL>right = val_where(to_update & ~update_left, mid, curr.right)<EOL>stopped = curr.stopped | failed | (right.df >= <NUM_LIT:0>)<EOL>return [_IntermediateResult(<EOL>iteration=curr.iteration,<EOL>stopped=stopped,<EOL>failed=failed,<EOL>num_evals=curr.num_evals + <NUM_LIT:1>,<EOL>left=left,<EOL>right=right)]<EOL><DEDENT>return tf.while_loop(<EOL>cond=_loop_cond, body=_loop_body, loop_vars=[initial_args])[<NUM_LIT:0>]<EOL>

Actual implementation of bisect given initial_args in a _BracketResult.


bracket_result = hzl.bracket(value_and_gradients_function, init_interval,<EOL>f_lim, max_iterations, expansion_param)<EOL>converged = init_interval.converged | _very_close(<EOL>bracket_result.left.x, bracket_result.right.x)<EOL>exhausted_iterations = ~converged & tf.greater_equal(<EOL>bracket_result.iteration, max_iterations)<EOL>line_search_args = HagerZhangLineSearchResult(<EOL>converged=converged,<EOL>failed=bracket_result.failed | exhausted_iterations,<EOL>iterations=bracket_result.iteration,<EOL>func_evals=bracket_result.num_evals,<EOL>left=bracket_result.left,<EOL>right=bracket_result.right)<EOL>return _line_search_after_bracketing(<EOL>value_and_gradients_function, line_search_args, init_interval.left,<EOL>f_lim, max_iterations, sufficient_decrease_param, curvature_param,<EOL>shrinkage_param)<EOL>

Brackets the minimum and performs a line search.

    Args:
      value_and_gradients_function: A Python callable that accepts a real scalar
        tensor and returns a namedtuple with the fields 'x', 'f', and 'df' that
        correspond to scalar tensors of real dtype containing the point at which
        the function was evaluated, the value of the function, and its
        derivative at that point. The other namedtuple fields, if present,
        should be tensors or sequences (possibly nested) of tensors.
        In usual optimization application, this function would be generated by
        projecting the multivariate objective function along some specific
        direction. The direction is determined by some other procedure but should
        be a descent direction (i.e. the derivative of the projected univariate
        function must be negative at 0.).
        Alternatively, the function may represent the batching of `n` such line
        functions (e.g. projecting a single multivariate objective function along
        `n` distinct directions at once) accepting n points as input, i.e. a
        tensor of shape [n], and the fields 'x', 'f' and 'df' in the returned
        namedtuple should each be a tensor of shape [n], with the corresponding
        input points, function values, and derivatives at those input points.
      init_interval: Instance of `HagerZhangLineSearchResults` containing
        the initial line search interval. The gradient of init_interval.left must
        be negative (i.e. must be a descent direction), while init_interval.right
        must be positive and finite.
      f_lim: Scalar `Tensor` of float dtype.
      max_iterations: Positive scalar `Tensor` of integral dtype. The maximum
        number of iterations to perform in the line search. The number of
        iterations used to bracket the minimum are also counted against this
        parameter.
      shrinkage_param: Scalar positive Tensor of real dtype. Must be less than
        `1.`. Corresponds to the parameter `gamma` in [Hager and Zhang (2006)][2].
      expansion_param: Scalar positive `Tensor` of real dtype. Must be greater
        than `1.`. Used to expand the initial interval in case it does not bracket
        a minimum. Corresponds to `rho` in [Hager and Zhang (2006)][2].
      sufficient_decrease_param: Positive scalar `Tensor` of real dtype.
        Bounded above by the curvature param. Corresponds to `delta` in the
        terminology of [Hager and Zhang (2006)][2].
      curvature_param: Positive scalar `Tensor` of real dtype. Bounded above
        by `1.`. Corresponds to 'sigma' in the terminology of
        [Hager and Zhang (2006)][2].

    Returns:
      A namedtuple containing the following fields.
        converged: Boolean `Tensor` of shape [n]. Whether a point satisfying
          Wolfe/Approx wolfe was found.
        failed: Boolean `Tensor` of shape [n]. Whether line search failed e.g.
          if either the objective function or the gradient are not finite at
          an evaluation point.
        iterations: Scalar int32 `Tensor`. Number of line search iterations made.
        func_evals: Scalar int32 `Tensor`. Number of function evaluations made.
        left: A namedtuple, as returned by value_and_gradients_function,
          of the left end point of the updated bracketing interval.
        right: A namedtuple, as returned by value_and_gradients_function,
          of the right end point of the updated bracketing interval.


midpoint = (search_interval.left.x + search_interval.right.x) / <NUM_LIT:2><EOL>val_mid = value_and_gradients_function(midpoint)<EOL>is_valid_mid = hzl.is_finite(val_mid)<EOL>still_active = active & is_valid_mid<EOL>new_failed = active & ~is_valid_mid<EOL>next_inteval = search_interval._replace(<EOL>failed=search_interval.failed | new_failed,<EOL>func_evals=search_interval.func_evals + <NUM_LIT:1>)<EOL>def _apply_update():<EOL><INDENT>update_result = hzl.update(<EOL>value_and_gradients_function, next_inteval.left, next_inteval.right,<EOL>val_mid, f_lim, active=still_active)<EOL>return HagerZhangLineSearchResult(<EOL>converged=next_inteval.converged,<EOL>failed=next_inteval.failed | update_result.failed,<EOL>iterations=next_inteval.iterations + update_result.iteration,<EOL>func_evals=next_inteval.func_evals + update_result.num_evals,<EOL>left=update_result.left,<EOL>right=update_result.right)<EOL><DEDENT>return prefer_static.cond(<EOL>tf.reduce_any(input_tensor=still_active),<EOL>_apply_update,<EOL>lambda: next_inteval)<EOL>

Performs bisection and updates the interval.


<EOL>iter_max = np.ceil(-np.log2(_machine_eps(val_c_input.x.dtype)))<EOL>def _cond(i, val_c, to_fix):<EOL><INDENT>del val_c  <EOL>return (i < iter_max) & tf.reduce_any(input_tensor=to_fix)<EOL><DEDENT>def _body(i, val_c, to_fix):<EOL><INDENT>next_c = tf.where(to_fix, val_c.x * step_size_shrink_param, val_c.x)<EOL>next_val_c = value_and_gradients_function(next_c)<EOL>still_to_fix = to_fix & ~hzl.is_finite(next_val_c)<EOL>return (i + <NUM_LIT:1>, next_val_c, still_to_fix)<EOL><DEDENT>to_fix = active & ~hzl.is_finite(val_c_input)<EOL>return tf.while_loop(<EOL>cond=_cond, body=_body, loop_vars=(<NUM_LIT:0>, val_c_input, to_fix))<EOL>

Shrinks the input step size until the value and grad become finite.


def _replace_worst_with_reflected():<EOL><INDENT>next_simplex = _replace_at_index(simplex, worst_index, reflected)<EOL>next_objective_values = _replace_at_index(objective_values, worst_index,<EOL>objective_at_reflected)<EOL>return False, next_simplex, next_objective_values, <NUM_LIT:0><EOL><DEDENT>return _replace_worst_with_reflected<EOL>

Creates the condition function pair for a reflection to be accepted.


n_points = tf.shape(input=arg_batch)[<NUM_LIT:0>]<EOL>if batch_evaluate_objective:<EOL><INDENT>return objective_function(arg_batch), n_points<EOL><DEDENT>return tf.map_fn(objective_function, arg_batch), n_points<EOL>

Evaluates the objective function on a batch of points.

    If `batch_evaluate_objective` is True, returns
    `objective function(arg_batch)` else it maps the `objective_function`
    across the `arg_batch`.

    Args:
      objective_function: A Python callable that accepts a single `Tensor` of
        rank 'R > 1' and any shape 's' and returns a scalar `Tensor` of real dtype
        containing the value of the function at that point. If
        `batch a `Tensor` of shape `[batch_size] + s ` where `batch_size` is the
        size of the batch of args. In this case, the expected return value is a
        `Tensor` of shape `[batch_size]`.
      arg_batch: A `Tensor` of real dtype. The batch of arguments at which to
        evaluate the `objective_function`. If `batch_evaluate_objective` is False,
        `arg_batch` will be unpacked along the zeroth axis and the
        `objective_function` will be applied to each element.
      batch_evaluate_objective: `bool`. Whether the `objective_function` can
        evaluate a batch of arguments at once.

    Returns:
      A tuple containing:
        objective_values: A `Tensor` of real dtype and shape `[batch_size]`.
          The value of the objective function evaluated at the supplied
          `arg_batch`.
        num_evaluations: An `int32` scalar `Tensor`containing the number of
          points on which the objective function was evaluated (i.e `batch_size`).


<EOL>best_vertex = simplex[best_index]<EOL>shrunk_simplex = best_vertex + shrinkage * (simplex - best_vertex)<EOL>objective_at_shrunk_simplex, evals = _evaluate_objective_multiple(<EOL>objective_function,<EOL>shrunk_simplex,<EOL>batch_evaluate_objective)<EOL>return (False,<EOL>shrunk_simplex,<EOL>objective_at_shrunk_simplex,<EOL>evals)<EOL>

Shrinks the simplex around the best vertex.


objective_convergence = tf.abs(worst_objective -<EOL>best_objective) < func_tolerance<EOL>simplex_degeneracy = tf.reduce_max(<EOL>input_tensor=tf.abs(simplex - best_vertex)) < position_tolerance<EOL>return objective_convergence | simplex_degeneracy<EOL>

Returns True if the simplex has converged.

    If the simplex size is smaller than the `position_tolerance` or the variation
    of the function value over the vertices of the simplex is smaller than the
    `func_tolerance` return True else False.

    Args:
      simplex: `Tensor` of real dtype. The simplex to test for convergence. For
        more details, see the docstring for `initial_simplex` argument
        of `minimize`.
      best_vertex: `Tensor` of real dtype and rank one less than `simplex`. The
        vertex with the best (i.e. smallest) objective value.
      best_objective: Scalar `Tensor` of real dtype. The best (i.e. smallest)
        value of the objective function at a vertex.
      worst_objective: Scalar `Tensor` of same dtype as `best_objective`. The
        worst (i.e. largest) value of the objective function at a vertex.
      func_tolerance: Scalar positive `Tensor`. The tolerance for the variation
        of the objective function value over the simplex. If the variation over
        the simplex vertices is below this threshold, convergence is True.
      position_tolerance: Scalar positive `Tensor`. The algorithm stops if the
        lengths (under the supremum norm) of edges connecting to the best vertex
        are below this threshold.

    Returns:
      has_converged: A scalar boolean `Tensor` indicating whether the algorithm
        is deemed to have converged.


initial_simplex = tf.convert_to_tensor(value=initial_simplex)<EOL>num_vertices = tf.shape(input=initial_simplex)[<NUM_LIT:0>]<EOL>dim = num_vertices - <NUM_LIT:1><EOL>num_evaluations = <NUM_LIT:0><EOL>if objective_at_initial_simplex is None:<EOL><INDENT>objective_at_initial_simplex, n_evals = _evaluate_objective_multiple(<EOL>objective_function, initial_simplex, batch_evaluate_objective)<EOL>num_evaluations += n_evals<EOL><DEDENT>objective_at_initial_simplex = tf.convert_to_tensor(<EOL>value=objective_at_initial_simplex)<EOL>return (dim,<EOL>num_vertices,<EOL>initial_simplex,<EOL>objective_at_initial_simplex,<EOL>num_evaluations)<EOL>

Evaluates the objective function at the specified initial simplex.


return tf.broadcast_to(<EOL>tf.convert_to_tensor(value=value, dtype=target.dtype),<EOL>distribution_util.prefer_static_shape(target)[:-<NUM_LIT:1>])<EOL>

Broadcast a value to match the batching dimensions of a target.

    If necessary the value is converted into a tensor. Both value and target
    should be of the same dtype.

    Args:
      value: A value to broadcast.
      target: A `Tensor` of shape [b1, ..., bn, d].

    Returns:
      A `Tensor` of shape [b1, ..., bn] and same dtype as the target.


grad_converged = norm(next_gradient, dims=<NUM_LIT:1>) <= grad_tolerance<EOL>x_converged = norm(next_position - current_position, dims=<NUM_LIT:1>) <= x_tolerance<EOL>f_converged = (norm(next_objective - current_objective, dims=<NUM_LIT:0>) <=<EOL>f_relative_tolerance * current_objective)<EOL>return grad_converged | x_converged | f_converged<EOL>

Checks if the algorithm satisfies the convergence criteria.


if control_inputs:<EOL><INDENT>with tf.control_dependencies(control_inputs):<EOL><INDENT>f0, df0 = value_and_gradients_function(initial_position)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>f0, df0 = value_and_gradients_function(initial_position)<EOL><DEDENT>converged = norm(df0, dims=<NUM_LIT:1>) < grad_tolerance<EOL>return dict(<EOL>converged=converged,<EOL>failed=tf.zeros_like(converged),  <EOL>num_iterations=tf.convert_to_tensor(value=<NUM_LIT:0>),<EOL>num_objective_evaluations=tf.convert_to_tensor(value=<NUM_LIT:1>),<EOL>position=initial_position,<EOL>objective_value=f0,<EOL>objective_gradient=df0)<EOL>

Returns a dictionary to populate the initial state of the search procedure.

    Performs an initial convergence check and the first evaluation of the
    objective function.

    Args:
      value_and_gradients_function: A Python callable that accepts a tensor and
        returns a tuple of two tensors: the objective function value and its
        derivative.
      initial_position: The starting point of the search procedure.
      grad_tolerance: The gradient tolerance for the procedure.
      control_inputs: Optional ops used to assert the validity of inputs, these
        are added as control dependencies to execute before the objective
        function is evaluated for the first time.

    Returns:
      An dictionary with values for the following keys:
        converged: True if the convergence check finds that the initial position
          is already an argmin of the objective function.
        failed: Initialized to False.
        num_objective_evaluations: Initialized to 1.
        position: Initialized to the initial position.
        objective_value: Initialized to the value of the objective function at
          the initial position.
        objective_gradient: Initialized to the gradient of the objective
          function at the initial position.


return tf.reduce_all(input_tensor=converged | failed)<EOL>

Condition to stop when all batch members have converged or failed.


rng = np.random.RandomState(<NUM_LIT:0>)<EOL>num_points = np.prod(grid_spec.shape)<EOL>grid = np.linspace(grid_spec.min, grid_spec.max, num=num_points).astype(dtype)<EOL>grid_spacing = (grid_spec.max - grid_spec.min) / num_points<EOL>grid += <NUM_LIT:0.1> * grid_spacing * rng.randn(*grid.shape)<EOL>grid = np.sort(grid)<EOL>return np.reshape(grid, grid_spec.shape)<EOL>

Returns a uniform grid + noise, reshaped to shape argument.


<EOL>p = tf.constant(<EOL>np.array([<EOL><NUM_LIT:0.>,<EOL>np.exp(-<NUM_LIT>),<EOL>np.exp(-<NUM_LIT>),<EOL><NUM_LIT:1.> - np.exp(-<NUM_LIT>),<EOL><NUM_LIT:1.> - np.exp(-<NUM_LIT>),<EOL><NUM_LIT:1.>,<EOL>]).astype(dtype))<EOL>_, grads = value_and_gradient(special_math.ndtri, p)<EOL>self.assertAllFinite(self.evaluate(grads[<NUM_LIT:0>]))<EOL>

Verifies that ndtri has finite gradients at interesting points.


if seed is None:<EOL><INDENT>return None<EOL><DEDENT>string = (str(seed) + salt).encode("<STR_LIT:utf-8>")<EOL>return int(hashlib.md5(string).hexdigest()[:<NUM_LIT:8>], <NUM_LIT:16>) & <NUM_LIT><EOL>

Generate a new seed, from the given seed and salt.


if not _is_known_dtype(dt):<EOL><INDENT>raise TypeError("<STR_LIT>".format(dt.name))<EOL><DEDENT>if dt.is_floating:<EOL><INDENT>return int(<NUM_LIT:2>**(np.finfo(dt.as_numpy_dtype).nmant + <NUM_LIT:1>))<EOL><DEDENT>if dt.is_integer:<EOL><INDENT>return np.iinfo(dt.as_numpy_dtype).max<EOL><DEDENT>if dt.base_dtype == tf.bool:<EOL><INDENT>return int(<NUM_LIT:1>)<EOL><DEDENT>raise TypeError("<STR_LIT>".format(dt.name))<EOL>

Helper returning the largest integer exactly representable by dtype.


def pad(x, axis, front=False, back=False, value=<NUM_LIT:0>, count=<NUM_LIT:1>, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>value = tf.convert_to_tensor(value=value, dtype=x.dtype, name="<STR_LIT:value>")<EOL>count = tf.convert_to_tensor(value=count, name="<STR_LIT:count>")<EOL>if not dtype_util.is_integer(count.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>".format(<EOL>dtype_util.name(count.dtype)))<EOL><DEDENT>if not front and not back:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>ndims = (<EOL>tensorshape_util.rank(x.shape)<EOL>if tensorshape_util.rank(x.shape) is not None else tf.rank(<EOL>x, name="<STR_LIT>"))<EOL>axis = tf.convert_to_tensor(value=axis, name="<STR_LIT>")<EOL>axis_ = tf.get_static_value(axis)<EOL>if axis_ is not None:<EOL><INDENT>axis = axis_<EOL>if axis < <NUM_LIT:0>:<EOL><INDENT>axis = ndims + axis<EOL><DEDENT>count_ = tf.get_static_value(count)<EOL>if axis_ >= <NUM_LIT:0> or tensorshape_util.rank(x.shape) is not None:<EOL><INDENT>head = x.shape[:axis]<EOL>mid_dim_value = tf.compat.dimension_value(x.shape[axis])<EOL>if count_ is None or mid_dim_value is None:<EOL><INDENT>middle = tf.TensorShape(None)<EOL><DEDENT>else:<EOL><INDENT>middle = tf.TensorShape(mid_dim_value + count_ * (front + back))<EOL><DEDENT>tail = x.shape[axis + <NUM_LIT:1>:]<EOL>final_shape = head.concatenate(middle.concatenate(tail))<EOL><DEDENT>else:<EOL><INDENT>final_shape = None<EOL><DEDENT><DEDENT>else:<EOL><INDENT>axis = tf.where(axis < <NUM_LIT:0>, ndims + axis, axis)<EOL>final_shape = None<EOL><DEDENT>x = tf.pad(<EOL>tensor=x,<EOL>paddings=tf.one_hot(<EOL>indices=tf.stack([axis if front else -<NUM_LIT:1>, axis if back else -<NUM_LIT:1>]),<EOL>depth=ndims,<EOL>axis=<NUM_LIT:0>,<EOL>on_value=count,<EOL>dtype=tf.int32),<EOL>constant_values=value)<EOL>if final_shape is not None:<EOL><INDENT>tensorshape_util.set_shape(x, final_shape)<EOL><DEDENT>return x<EOL><DEDENT>

Pads `value` to the front and/or back of a `Tensor` dim, `count` times.

    Args:
      x: `Tensor` input.
      axis: Scalar `int`-like `Tensor` representing the single dimension to pad.
        (Negative indexing is supported.)
      front: Python `bool`; if `True` the beginning of the `axis` dimension is
        padded with `value`, `count` times. If `False` no front padding is made.
      back: Python `bool`; if `True` the end of the `axis` dimension is padded
        with `value`, `count` times. If `False` no end padding is made.
      value: Scalar `int`-like `Tensor` representing the actual value added to the
        front and/or back of the `axis` dimension of `x`.
      count: Scalar `int`-like `Tensor` representing number of elements added to
        the front and/or back of the `axis` dimension of `x`. E.g., if `front =
        back = True` then `2 * count` elements are added.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      pad: The padded version of input `x`.

    Raises:
      ValueError: if both `front` and `back` are `False`.
      TypeError: if `count` is not `int`-like.


with tf.name_scope(name):<EOL><INDENT>def make_shape_tensor(x):<EOL><INDENT>return tf.convert_to_tensor(value=x, name="<STR_LIT>", dtype=tf.int32)<EOL><DEDENT>def get_tensor_shape(s):<EOL><INDENT>if isinstance(s, tf.TensorShape):<EOL><INDENT>return s<EOL><DEDENT>s_ = tf.get_static_value(make_shape_tensor(s))<EOL>if s_ is not None:<EOL><INDENT>return tf.TensorShape(s_)<EOL><DEDENT>return None<EOL><DEDENT>def get_shape_tensor(s):<EOL><INDENT>if not isinstance(s, tf.TensorShape):<EOL><INDENT>return make_shape_tensor(s)<EOL><DEDENT>if tensorshape_util.is_fully_defined(s):<EOL><INDENT>return make_shape_tensor(tensorshape_util.as_list(s))<EOL><DEDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>shape1_ = get_tensor_shape(shape1)<EOL>shape2_ = get_tensor_shape(shape2)<EOL>if shape1_ is not None and shape2_ is not None:<EOL><INDENT>return tf.broadcast_static_shape(shape1_, shape2_)<EOL><DEDENT>shape1_ = get_shape_tensor(shape1)<EOL>shape2_ = get_shape_tensor(shape2)<EOL>return tf.broadcast_dynamic_shape(shape1_, shape2_)<EOL><DEDENT>

Convenience function which statically broadcasts shape when possible.

    Args:
      shape1:  `1-D` integer `Tensor`.  Already converted to tensor!
      shape2:  `1-D` integer `Tensor`.  Already converted to tensor!
      name:  A string name to prepend to created ops.

    Returns:
      The broadcast shape, either as `TensorShape` (if broadcast can be done
        statically), or as a `Tensor`.


a = tf.convert_to_tensor(value=a, name="<STR_LIT:a>")<EOL>b = tf.convert_to_tensor(value=b, name="<STR_LIT:b>")<EOL>def all_shapes_equal():<EOL><INDENT>return tf.reduce_all(<EOL>input_tensor=tf.equal(<EOL>tf.concat([tf.shape(input=a), tf.shape(input=b)], <NUM_LIT:0>),<EOL>tf.concat([tf.shape(input=b), tf.shape(input=a)], <NUM_LIT:0>)))<EOL><DEDENT>return tf.cond(<EOL>pred=tf.equal(tf.rank(a), tf.rank(b)),<EOL>true_fn=all_shapes_equal,<EOL>false_fn=lambda: tf.constant(False))<EOL>

Returns whether a and b have the same dynamic shape.

    Args:
      a: `Tensor`
      b: `Tensor`

    Returns:
      `bool` `Tensor` representing if both tensors have the same shape.


axis = tf.convert_to_tensor(value=axis, name="<STR_LIT>")<EOL>rank = tf.convert_to_tensor(value=rank, name="<STR_LIT>")<EOL>axis_ = tf.get_static_value(axis)<EOL>rank_ = tf.get_static_value(rank)<EOL>if axis_ is not None and rank_ is not None:<EOL><INDENT>is_scalar = axis_.ndim == <NUM_LIT:0><EOL>if is_scalar:<EOL><INDENT>axis_ = [axis_]<EOL><DEDENT>positive_axis = []<EOL>for a_ in axis_:<EOL><INDENT>if a_ < <NUM_LIT:0>:<EOL><INDENT>positive_axis.append(rank_ + a_)<EOL><DEDENT>else:<EOL><INDENT>positive_axis.append(a_)<EOL><DEDENT><DEDENT>if is_scalar:<EOL><INDENT>positive_axis = positive_axis[<NUM_LIT:0>]<EOL><DEDENT>return tf.convert_to_tensor(value=positive_axis, dtype=axis.dtype)<EOL><DEDENT>return tf.where(axis < <NUM_LIT:0>, rank + axis, axis)<EOL>

Make (possibly negatively indexed) `axis` argument non-negative.