def _sample_with_shrinkage(x_initial, target_log_prob, log_slice_heights,<EOL>step_size, lower_bounds, upper_bounds, seed=None,<EOL>name=None):

with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[x_initial, log_slice_heights, step_size, lower_bounds, upper_bounds]):<EOL><INDENT>seed_gen = distributions.SeedStream(seed, salt='<STR_LIT>')<EOL>found = tf.zeros_like(x_initial, dtype=tf.bool)<EOL>cond = lambda found, *ignored_args: ~tf.reduce_all(input_tensor=found)<EOL>x_next = tf.identity(x_initial)<EOL>x_initial_shape = tf.shape(input=x_initial)<EOL>x_initial_dtype = x_initial.dtype.base_dtype<EOL>def _body(found, left, right, x_next):<EOL><INDENT>"""<STR_LIT>"""<EOL>proportions = tf.random.uniform(<EOL>x_initial_shape, dtype=x_initial_dtype, seed=seed_gen())<EOL>x_proposed = tf.where(~found, left + proportions * (right - left), x_next)<EOL>accept_res = _test_acceptance(x_initial, target_log_prob=target_log_prob,<EOL>decided=found,<EOL>log_slice_heights=log_slice_heights,<EOL>x_proposed=x_proposed, step_size=step_size,<EOL>lower_bounds=left, upper_bounds=right)<EOL>boundary_test = log_slice_heights < target_log_prob(x_proposed)<EOL>can_accept = boundary_test & accept_res<EOL>next_found = found | can_accept<EOL>next_left = tf.where(x_proposed < x_initial, x_proposed, left)<EOL>next_right = tf.where(x_proposed >= x_initial, x_proposed, right)<EOL>return next_found, next_left, next_right, x_proposed<EOL><DEDENT>return tf.while_loop(<EOL>cond=cond,<EOL>body=_body,<EOL>loop_vars=(found, lower_bounds, upper_bounds, x_next))[-<NUM_LIT:1>]<EOL><DEDENT>

Samples from the slice by applying shrinkage for rejected points.

    Implements the one dimensional slice sampling algorithm of Neal (2003), with a
    doubling algorithm (Neal 2003 P715 Fig. 4), which doubles the size of the
    interval at each iteration and shrinkage (Neal 2003 P716 Fig. 5), which
    reduces the width of the slice when a selected point is rejected, by setting
    the relevant bound that that value. Randomly sampled points are checked for
    two criteria: that they lie within the slice and that they pass the
    acceptability check (Neal 2003 P717 Fig. 6), which tests that the new state
    could have generated the previous one.

    Args:
      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 `[]`).
      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.
      log_slice_heights: Tensor of the same shape and dtype as the return value
        of `target_log_prob` when applied to `x_initial`. The log of the height of
        the chosen slice.
      step_size: A tensor of shape and dtype compatible with `x_initial`. The min
        interval size in the doubling algorithm.
      lower_bounds: Tensor of same shape and dtype as `x_initial`. Slice lower
        bounds for each chain.
      upper_bounds: Tensor of same shape and dtype as `x_initial`. Slice upper
        bounds for each chain.
      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:
      x_proposed: A tensor of the same shape and dtype as `x_initial`. The next
        proposed state of the chain.


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.


def eight_schools_joint_log_prob(<EOL>treatment_effects, treatment_stddevs,<EOL>avg_effect, avg_stddev, school_effects_standard):

rv_avg_effect = tfd.Normal(loc=<NUM_LIT:0.>, scale=<NUM_LIT>)<EOL>rv_avg_stddev = tfd.Normal(loc=<NUM_LIT>, scale=<NUM_LIT:1.>)<EOL>rv_school_effects_standard = mvn(<EOL>loc=tf.zeros_like(school_effects_standard),<EOL>scale=tf.ones_like(school_effects_standard))<EOL>rv_treatment_effects = mvn(<EOL>loc=(avg_effect + tf.exp(avg_stddev) * school_effects_standard),<EOL>scale=treatment_stddevs)<EOL>return (<EOL>rv_avg_effect.log_prob(avg_effect) +<EOL>rv_avg_stddev.log_prob(avg_stddev) +<EOL>rv_school_effects_standard.log_prob(school_effects_standard) +<EOL>rv_treatment_effects.log_prob(treatment_effects))<EOL>


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.


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

Helper which uses graph seed if using eager.


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


def _get_drift(step_size_parts, volatility_parts, grads_volatility,<EOL>grads_target_log_prob,<EOL>name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [<EOL>step_size_parts, volatility_parts, grads_volatility, grads_target_log_prob<EOL>]):<EOL><INDENT>drift_parts = []<EOL>for step_size, volatility, grad_volatility, grad_target_log_prob in (<EOL>zip(step_size_parts,<EOL>volatility_parts,<EOL>grads_volatility,<EOL>grads_target_log_prob)):<EOL><INDENT>volatility_squared = tf.square(volatility)<EOL>drift = <NUM_LIT:0.5> * step_size * (volatility_squared * grad_target_log_prob<EOL>+ grad_volatility)<EOL>drift_parts.append(drift)<EOL><DEDENT>return drift_parts<EOL><DEDENT>

Compute diffusion drift at the current location `current_state`.

    The drift of the diffusion at is computed as

    ```none
    0.5 * `step_size` * volatility_parts * `target_log_prob_fn(current_state)`
    + `step_size` * `grads_volatility`
    ```

    where `volatility_parts` = `volatility_fn(current_state)**2` and
    `grads_volatility` is a gradient of `volatility_parts` at the `current_state`.

    Args:
      step_size_parts: Python `list` of `Tensor`s representing the step size for
        Euler-Maruyama method. Must broadcast with the shape of
        `volatility_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)`.
      grads_volatility: Python list of `Tensor`s representing the value of the
        gradient of `volatility_parts**2` wrt the state of the chain.
      grads_target_log_prob: Python list of `Tensor`s representing
        gradient of `target_log_prob_fn(*state_parts`) wrt `state_parts`. Must
        have same shape as `volatility_parts`.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'mala_get_drift').

    Returns:
      drift_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`.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [<EOL>current_state_parts, proposed_state_parts, current_volatility_parts,<EOL>proposed_volatility_parts, current_drift_parts, proposed_drift_parts,<EOL>step_size_parts, independent_chain_ndims<EOL>]):<EOL><INDENT>proposed_log_density_parts = []<EOL>dual_log_density_parts = []<EOL>for [<EOL>current_state,<EOL>proposed_state,<EOL>current_volatility,<EOL>proposed_volatility,<EOL>current_drift,<EOL>proposed_drift,<EOL>step_size,<EOL>] in zip(<EOL>current_state_parts,<EOL>proposed_state_parts,<EOL>current_volatility_parts,<EOL>proposed_volatility_parts,<EOL>current_drift_parts,<EOL>proposed_drift_parts,<EOL>step_size_parts,<EOL>):<EOL><INDENT>axis = tf.range(independent_chain_ndims, tf.rank(current_state))<EOL>state_diff = proposed_state - current_state<EOL>current_volatility *= tf.sqrt(step_size)<EOL>proposed_energy = (state_diff - current_drift) / current_volatility<EOL>proposed_volatility *= tf.sqrt(step_size)<EOL>proposed_energy = (<EOL>tf.reduce_sum(<EOL>input_tensor=mcmc_util.safe_sum(<EOL>[tf.math.log(current_volatility),<EOL><NUM_LIT:0.5> * (proposed_energy**<NUM_LIT:2>)]),<EOL>axis=axis))<EOL>proposed_log_density_parts.append(-proposed_energy)<EOL>dual_energy = (state_diff + proposed_drift) / proposed_volatility<EOL>dual_energy = (<EOL>tf.reduce_sum(<EOL>input_tensor=mcmc_util.safe_sum(<EOL>[tf.math.log(proposed_volatility), <NUM_LIT:0.5> * (dual_energy**<NUM_LIT:2>)]),<EOL>axis=axis))<EOL>dual_log_density_parts.append(-dual_energy)<EOL><DEDENT>proposed_log_density_reduce = tf.reduce_sum(<EOL>input_tensor=tf.stack(proposed_log_density_parts, axis=-<NUM_LIT:1>), axis=-<NUM_LIT:1>)<EOL>dual_log_density_reduce = tf.reduce_sum(<EOL>input_tensor=tf.stack(dual_log_density_parts, axis=-<NUM_LIT:1>), axis=-<NUM_LIT:1>)<EOL>return mcmc_util.safe_sum([dual_log_density_reduce,<EOL>-proposed_log_density_reduce])<EOL><DEDENT>

r"""Helper to `kernel` which computes the log acceptance-correction.

    Computes `log_acceptance_correction` as described in `MetropolisHastings`
    class. The proposal density is normal. More specifically,

     ```none
    q(proposed_state | current_state) \sim N(current_state + current_drift,
    step_size * current_volatility**2)

    q(current_state | proposed_state) \sim N(proposed_state + proposed_drift,
    step_size * proposed_volatility**2)
    ```

    The `log_acceptance_correction` is then

    ```none
    log_acceptance_correctio = q(current_state | proposed_state)
    - q(proposed_state | current_state)
    ```

    Args:
      current_state_parts: Python `list` of `Tensor`s representing the value(s) of
        the current state of the chain.
      proposed_state_parts:  Python `list` of `Tensor`s representing the value(s)
        of the proposed state of the chain. Must broadcast with the shape of
        `current_state_parts`.
      current_volatility_parts: Python `list` of `Tensor`s representing the value
        of `volatility_fn(*current_volatility_parts)`. Must broadcast with the
        shape of `current_state_parts`.
      proposed_volatility_parts: Python `list` of `Tensor`s representing the value
        of `volatility_fn(*proposed_volatility_parts)`. Must broadcast with the
        shape of `current_state_parts`
      current_drift_parts: Python `list` of `Tensor`s representing value of the
        drift `_get_drift(*current_state_parts, ..)`. Must broadcast with the
        shape of `current_state_parts`.
      proposed_drift_parts: Python `list` of `Tensor`s representing value of the
        drift `_get_drift(*proposed_drift_parts, ..)`. Must broadcast with the
        shape of `current_state_parts`.
      step_size_parts: Python `list` of `Tensor`s representing the step size for
        Euler-Maruyama method. Must broadcast with the shape of
        `current_state_parts`.
      independent_chain_ndims: Scalar `int` `Tensor` representing the number of
        leftmost `Tensor` dimensions which index independent chains.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'compute_log_acceptance_correction').

    Returns:
      log_acceptance_correction: `Tensor` representing the `log`
        acceptance-correction.  (See docstring for mathematical definition.)


state_parts = list(state) if mcmc_util.is_list_like(state) else [state]<EOL>needs_volatility_fn_gradients = grads_volatility_fn is None<EOL>if volatility_fn_results is None:<EOL><INDENT>volatility_fn_results = volatility_fn(*state_parts)<EOL><DEDENT>volatility_fn_results = (list(volatility_fn_results)<EOL>if mcmc_util.is_list_like(volatility_fn_results)<EOL>else [volatility_fn_results])<EOL>if len(volatility_fn_results) == <NUM_LIT:1>:<EOL><INDENT>volatility_fn_results *= len(state_parts)<EOL><DEDENT>if len(state_parts) != len(volatility_fn_results):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>volatility_fn_results = _maybe_broadcast_volatility(volatility_fn_results,<EOL>state_parts)<EOL>if grads_volatility_fn is None:<EOL><INDENT>[<EOL>_,<EOL>grads_volatility_fn,<EOL>] = diag_jacobian(<EOL>xs=state_parts,<EOL>ys=volatility_fn_results,<EOL>sample_shape=sample_shape,<EOL>parallel_iterations=parallel_iterations,<EOL>fn=volatility_fn)<EOL><DEDENT>if needs_volatility_fn_gradients:<EOL><INDENT>grads_volatility_fn = [<EOL><NUM_LIT> * g * volatility if g is not None else tf.zeros_like(<EOL>fn_arg, dtype=fn_arg.dtype.base_dtype)<EOL>for g, volatility, fn_arg in zip(<EOL>grads_volatility_fn, volatility_fn_results, state_parts)<EOL>]<EOL><DEDENT>return volatility_fn_results, grads_volatility_fn<EOL>

Helper which computes `volatility_fn` results and grads, if needed.


return [v + tf.zeros_like(sp, dtype=sp.dtype.base_dtype)<EOL>for v, sp in zip(volatility_parts, state_parts)]<EOL>

Helper to broadcast `volatility_parts` to the shape of `state_parts`.


state_parts = list(state) if mcmc_util.is_list_like(state) else [state]<EOL>[<EOL>target_log_prob,<EOL>grads_target_log_prob,<EOL>] = mcmc_util.maybe_call_fn_and_grads(<EOL>target_log_prob_fn,<EOL>state_parts,<EOL>target_log_prob,<EOL>grads_target_log_prob)<EOL>[<EOL>volatility_parts,<EOL>grads_volatility,<EOL>] = _maybe_call_volatility_fn_and_grads(<EOL>volatility_fn,<EOL>state_parts,<EOL>volatility,<EOL>grads_volatility_fn,<EOL>distribution_util.prefer_static_shape(target_log_prob),<EOL>parallel_iterations)<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>if diffusion_drift is None:<EOL><INDENT>diffusion_drift_parts = _get_drift(step_sizes, volatility_parts,<EOL>grads_volatility,<EOL>grads_target_log_prob)<EOL><DEDENT>else:<EOL><INDENT>diffusion_drift_parts = (list(diffusion_drift)<EOL>if mcmc_util.is_list_like(diffusion_drift)<EOL>else [diffusion_drift])<EOL>if len(state_parts) != len(diffusion_drift):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT><DEDENT>return [<EOL>state_parts,<EOL>step_sizes,<EOL>target_log_prob,<EOL>grads_target_log_prob,<EOL>volatility_parts,<EOL>grads_volatility,<EOL>diffusion_drift_parts,<EOL>]<EOL>

Helper which processes input args to meet list-like assumptions.


impl = metropolis_hastings.MetropolisHastings(<EOL>inner_kernel=UncalibratedLangevin(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>step_size=step_size,<EOL>volatility_fn=volatility_fn,<EOL>seed=seed,<EOL>parallel_iterations=parallel_iterations,<EOL>name=name),<EOL>seed=seed)<EOL>self._impl = impl<EOL>parameters = impl.inner_kernel.parameters.copy()<EOL>del parameters['<STR_LIT>']<EOL>self._parameters = parameters<EOL>

Initializes MALA transition kernel.

        Args:
          target_log_prob_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it's a list) and returns its
            (possibly unnormalized) log-density under the target distribution.
          step_size: `Tensor` or Python `list` of `Tensor`s representing the step
            size for the leapfrog integrator. Must broadcast with the shape of
            `current_state`. 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_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it's a list) and returns
            volatility value at `current_state`. Should return a `Tensor` or Python
            `list` of `Tensor`s that must broadcast with the shape of
            `current_state` Defaults to the identity function.
          seed: Python integer to seed the random number generator.
          parallel_iterations: the number of coordinates for which the gradients of
            the volatility matrix `volatility_fn` can be computed in parallel.
            Default value: `None` (i.e., no seed).
          name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., 'mala_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.

        Raises:
          ValueError: if there isn't one `step_size` or a list with same length as
            `current_state`.
          TypeError: if `volatility_fn` is not callable.


return self._parameters<EOL>

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


def one_step(self, current_state, previous_kernel_results):

return self._impl.one_step(current_state, previous_kernel_results)<EOL>

Runs one iteration of MALA.

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


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

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


self._seed_stream = distributions.SeedStream(<EOL>seed, salt='<STR_LIT>')<EOL>if volatility_fn is None:<EOL><INDENT>volatility_fn = lambda *args: <NUM_LIT:1.><EOL><DEDENT>if not callable(volatility_fn):<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>type(volatility_fn)))<EOL><DEDENT>self._parameters = dict(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>step_size=step_size,<EOL>volatility_fn=volatility_fn,<EOL>compute_acceptance=tf.convert_to_tensor(value=compute_acceptance),<EOL>seed=seed,<EOL>parallel_iterations=parallel_iterations,<EOL>name=name)<EOL>

Initializes Langevin diffusion transition kernel.

        Args:
          target_log_prob_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it's a list) and returns its
            (possibly unnormalized) log-density under the target distribution.
          step_size: `Tensor` or Python `list` of `Tensor`s representing the step
            size for the leapfrog integrator. Must broadcast with the shape of
            `current_state`. 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_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it's a list) and returns
            volatility value at `current_state`. Should return a `Tensor` or Python
            `list` of `Tensor`s that must broadcast with the shape of
            `current_state` Defaults to the identity function.
          parallel_iterations: the number of coordinates for which the gradients of
            the volatility matrix `volatility_fn` can be computed in parallel.
          compute_acceptance: Python 'bool' indicating whether to compute the
            Metropolis log-acceptance ratio used to construct
            `MetropolisAdjustedLangevinAlgorithm` kernel.
          seed: Python integer to seed the random number generator.
            Default value: `None` (i.e., no seed).
          name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., 'mala_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.

        Raises:
          ValueError: if there isn't one `step_size` or a list with same length as
            `current_state`.
          TypeError: if `volatility_fn` is not callable.


return self._parameters<EOL>


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


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


return []<EOL>

Returns an object with the same type as returned by `one_step(...)[1]`.

        Args:
          init_state: `Tensor` or Python `list` of `Tensor`s representing the
            initial state(s) of the Markov chain(s).

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


Helper which uses graph seed if using TFE.


target = tfd.Normal(dtype(<NUM_LIT:0.>), dtype(<NUM_LIT:1.>))<EOL>def make_kernel_fn(target_log_prob_fn, seed):<EOL><INDENT>return tfp.mcmc.HamiltonianMonteCarlo(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>seed=seed,<EOL>step_size=step_size,<EOL>num_leapfrog_steps=<NUM_LIT:3>)<EOL><DEDENT>remc = tfp.mcmc.ReplicaExchangeMC(<EOL>target_log_prob_fn=tf.function(target.log_prob, autograph=False),<EOL>inverse_temperatures=inverse_temperatures,<EOL>make_kernel_fn=make_kernel_fn,<EOL>seed=_set_seed(<NUM_LIT:1>))<EOL>samples = tfp.mcmc.sample_chain(<EOL>num_results=num_results,<EOL>current_state=target.sample(seed=_set_seed(<NUM_LIT:1>)),<EOL>kernel=remc,<EOL>num_burnin_steps=<NUM_LIT:50>,<EOL>trace_fn=None,<EOL>parallel_iterations=<NUM_LIT:1>)  <EOL>self.assertAllEqual((num_results,), samples.shape)<EOL>return self.evaluate(samples)<EOL>

Sampling from standard normal with REMC.


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

Convenience function to efficiently construct a MultivariateNormalDiag.


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>


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.


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 offset_normal_chains_should_fail_wrapper(<EOL>self, sample_shape, independent_chain_shape, other_shape):

state_shape = sample_shape + independent_chain_shape + other_shape<EOL>state_ = rng.randn(*state_shape)<EOL>offset = np.linspace(<EOL><NUM_LIT:0>, <NUM_LIT:2>, num=np.prod(independent_chain_shape)).reshape([<NUM_LIT:1>] * len(<EOL>sample_shape) + independent_chain_shape + [<NUM_LIT:1>] * len(other_shape))<EOL>state_ += offset<EOL>self.check_results(state_, independent_chain_shape, should_pass=False)<EOL>

Check results with normal chains that are offset from each other.


def default_exchange_proposed_fn_(num_replica, seed=None):<EOL><INDENT>"""<STR_LIT>"""<EOL>seed_stream = distributions.SeedStream(seed, '<STR_LIT>')<EOL>zero_start = tf.random.uniform([], seed=seed_stream()) > <NUM_LIT:0.5><EOL>if num_replica % <NUM_LIT:2> == <NUM_LIT:0>:<EOL><INDENT>def _exchange():<EOL><INDENT>flat_exchange = tf.range(num_replica)<EOL>if num_replica > <NUM_LIT:2>:<EOL><INDENT>start = tf.cast(~zero_start, dtype=tf.int32)<EOL>end = num_replica - start<EOL>flat_exchange = flat_exchange[start:end]<EOL><DEDENT>return tf.reshape(flat_exchange, [tf.size(input=flat_exchange) // <NUM_LIT:2>, <NUM_LIT:2>])<EOL><DEDENT><DEDENT>else:<EOL><INDENT>def _exchange():<EOL><INDENT>start = tf.cast(zero_start, dtype=tf.int32)<EOL>end = num_replica - tf.cast(~zero_start, dtype=tf.int32)<EOL>flat_exchange = tf.range(num_replica)[start:end]<EOL>return tf.reshape(flat_exchange, [tf.size(input=flat_exchange) // <NUM_LIT:2>, <NUM_LIT:2>])<EOL><DEDENT><DEDENT>def _null_exchange():<EOL><INDENT>return tf.reshape(tf.cast([], dtype=tf.int32), shape=[<NUM_LIT:0>, <NUM_LIT:2>])<EOL><DEDENT>return tf.cond(<EOL>pred=tf.random.uniform([], seed=seed_stream()) < prob_exchange,<EOL>true_fn=_exchange,<EOL>false_fn=_null_exchange)<EOL><DEDENT>return default_exchange_proposed_fn_<EOL>

Default exchange proposal function, for replica exchange MC.

    With probability `prob_exchange` propose combinations of replica for exchange.
    When exchanging, create combinations of adjacent replicas in
    [Replica Exchange Monte Carlo](
    https://en.wikipedia.org/wiki/Parallel_tempering)

    ```
    exchange_fn = default_exchange_proposed_fn(prob_exchange=0.5)
    exchange_proposed = exchange_fn(num_replica=3)

    exchange_proposed.eval()
    ==> [[0, 1]]  # 1 exchange, 0 <--> 1

    exchange_proposed.eval()
    ==> []  # 0 exchanges
    ```

    Args:
      prob_exchange: Scalar `Tensor` giving probability that any exchanges will
        be generated.

    Returns:
      default_exchange_proposed_fn_: Python callable which take a number of
        replicas (a Python integer), and return combinations of replicas for
        exchange as an [n, 2] integer `Tensor`, `0 <= n <= num_replica // 2`,
        with *unique* values in the set `{0, ..., num_replica}`.


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.


if hasattr(kernel_results, field_name):<EOL><INDENT>return getattr(kernel_results, field_name)<EOL><DEDENT>if hasattr(kernel_results, '<STR_LIT>'):<EOL><INDENT>return getattr(kernel_results.accepted_results, field_name)<EOL><DEDENT>raise TypeError('<STR_LIT>' % (field_name, kernel_results))<EOL>

field_name from kernel_results or kernel_results.accepted_results.


inverse_temperatures = tf.convert_to_tensor(<EOL>value=inverse_temperatures, name='<STR_LIT>')<EOL>inverse_temperatures.shape.assert_is_fully_defined()<EOL>inverse_temperatures.shape.assert_has_rank(<NUM_LIT:1>)<EOL>self._seed_stream = distributions.SeedStream(seed, salt=name)<EOL>self._seeded_mcmc = seed is not None<EOL>self._parameters = dict(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>inverse_temperatures=inverse_temperatures,<EOL>num_replica=tf.compat.dimension_value(inverse_temperatures.shape[<NUM_LIT:0>]),<EOL>exchange_proposed_fn=exchange_proposed_fn,<EOL>seed=seed,<EOL>name=name)<EOL>self.replica_kernels = []<EOL>for i in range(self.num_replica):<EOL><INDENT>self.replica_kernels.append(<EOL>make_kernel_fn(<EOL>target_log_prob_fn=_replica_log_prob_fn(inverse_temperatures[i],<EOL>target_log_prob_fn),<EOL>seed=self._seed_stream()))<EOL><DEDENT>

Instantiates this object.

        Args:
          target_log_prob_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it's a list) and returns its
            (possibly unnormalized) log-density under the target distribution.
          inverse_temperatures: `1D` `Tensor of inverse temperatures to perform
            samplings with each replica. Must have statically known `shape`.
            `inverse_temperatures[0]` produces the states returned by samplers,
            and is typically == 1.
          make_kernel_fn: Python callable which takes target_log_prob_fn and seed
            args and returns a TransitionKernel instance.
          exchange_proposed_fn: Python callable which take a number of replicas, and
            return combinations of replicas for exchange.
          seed: Python integer to seed the random number generator.
            Default value: `None` (i.e., no seed).
          name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., "remc_kernel").

        Raises:
          ValueError: `inverse_temperatures` doesn't have statically known 1D shape.


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


with tf.compat.v1.name_scope(<EOL>name=mcmc_util.make_name(self.name, '<STR_LIT>', '<STR_LIT>'),<EOL>values=[init_state]):<EOL><INDENT>replica_results = [<EOL>self.replica_kernels[i].bootstrap_results(init_state)<EOL>for i in range(self.num_replica)<EOL>]<EOL>init_state_parts = (<EOL>list(init_state)<EOL>if mcmc_util.is_list_like(init_state) else [init_state])<EOL>replica_states = [[<EOL>tf.convert_to_tensor(value=s) for s in init_state_parts<EOL>] for i in range(self.num_replica)]<EOL>if not mcmc_util.is_list_like(init_state):<EOL><INDENT>replica_states = [s[<NUM_LIT:0>] for s in replica_states]<EOL><DEDENT>return ReplicaExchangeMCKernelResults(<EOL>replica_states=replica_states,<EOL>replica_results=replica_results,<EOL>sampled_replica_states=replica_states,<EOL>sampled_replica_results=replica_results,<EOL>)<EOL><DEDENT>

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

        Args:
          init_state: `Tensor` or Python `list` of `Tensor`s representing the
            initial state(s) of the Markov chain(s).

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


if step_counter is None and num_adaptation_steps is not None:<EOL><INDENT>step_counter = tf.compat.v1.get_variable(<EOL>name='<STR_LIT>',<EOL>initializer=np.array(-<NUM_LIT:1>, dtype=np.int32),<EOL>dtype=tf.int32,<EOL>trainable=False,<EOL>use_resource=True)<EOL><DEDENT>def step_size_simple_update_fn(step_size_var, kernel_results):<EOL><INDENT>"""<STR_LIT>"""<EOL>if kernel_results is None:<EOL><INDENT>if mcmc_util.is_list_like(step_size_var):<EOL><INDENT>return [tf.identity(ss) for ss in step_size_var]<EOL><DEDENT>return tf.identity(step_size_var)<EOL><DEDENT>log_n = tf.math.log(<EOL>tf.cast(<EOL>tf.size(input=kernel_results.log_accept_ratio),<EOL>kernel_results.log_accept_ratio.dtype))<EOL>log_mean_accept_ratio = tf.reduce_logsumexp(<EOL>input_tensor=tf.minimum(kernel_results.log_accept_ratio, <NUM_LIT:0.>)) - log_n<EOL>adjustment = tf.where(<EOL>log_mean_accept_ratio < tf.cast(<EOL>tf.math.log(target_rate), log_mean_accept_ratio.dtype),<EOL>-decrement_multiplier / (<NUM_LIT:1.> + decrement_multiplier),<EOL>increment_multiplier)<EOL>def build_assign_op():<EOL><INDENT>if mcmc_util.is_list_like(step_size_var):<EOL><INDENT>return [<EOL>ss.assign_add(ss * tf.cast(adjustment, ss.dtype))<EOL>for ss in step_size_var<EOL>]<EOL><DEDENT>return step_size_var.assign_add(<EOL>step_size_var * tf.cast(adjustment, step_size_var.dtype))<EOL><DEDENT>if num_adaptation_steps is None:<EOL><INDENT>return build_assign_op()<EOL><DEDENT>else:<EOL><INDENT>with tf.control_dependencies([step_counter.assign_add(<NUM_LIT:1>)]):<EOL><INDENT>return tf.cond(<EOL>pred=step_counter < num_adaptation_steps,<EOL>true_fn=build_assign_op,<EOL>false_fn=lambda: step_size_var)<EOL><DEDENT><DEDENT><DEDENT>return step_size_simple_update_fn<EOL>

Create a function implementing a step-size update policy.

    The simple policy increases or decreases the `step_size_var` based on the
    average of `exp(minimum(0., log_accept_ratio))`. It is based on
    [Section 4.2 of Andrieu and Thoms (2008)](
    https://people.eecs.berkeley.edu/~jordan/sail/readings/andrieu-thoms.pdf).

    The `num_adaptation_steps` argument is set independently of any burnin
    for the overall chain. In general, adaptation prevents the chain from
    reaching a stationary distribution, so obtaining consistent samples requires
    `num_adaptation_steps` be set to a value [somewhat smaller](
    http://andrewgelman.com/2017/12/15/burn-vs-warm-iterative-simulation-algorithms/#comment-627745)
    than the number of burnin steps. However, it may sometimes be helpful to set
    `num_adaptation_steps` to a larger value during development in order to
    inspect the behavior of the chain during adaptation.

    Args:
      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. If `None`, the step size is adapted
        on every step (note this breaks stationarity of the chain!).
      target_rate: Scalar `Tensor` representing desired `accept_ratio`.
        Default value: `0.75` (i.e., [center of asymptotically optimal
        rate](https://arxiv.org/abs/1411.6669)).
      decrement_multiplier: `Tensor` representing amount to downscale current
        `step_size`.
        Default value: `0.01`.
      increment_multiplier: `Tensor` representing amount to upscale current
        `step_size`.
        Default value: `0.01`.
      step_counter: Scalar `int` `Variable` specifying the current step. The step
        size is adapted iff `step_counter < num_adaptation_steps`.
        Default value: if `None`, an internal variable
          `step_size_adaptation_step_counter` is created and initialized to `-1`.

    Returns:
      step_size_simple_update_fn: Callable that takes args
        `step_size_var, kernel_results` and returns updated step size(s).


<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', [<EOL>independent_chain_ndims, step_sizes, current_momentum_parts,<EOL>current_state_parts, current_target_log_prob,<EOL>current_target_log_prob_grad_parts<EOL>]):<EOL><INDENT>proposed_momentum_parts = [<EOL>v + <NUM_LIT:0.5> * tf.cast(eps, v.dtype) * g<EOL>for v, eps, g<EOL>in zip(current_momentum_parts,<EOL>step_sizes,<EOL>current_target_log_prob_grad_parts)]<EOL>proposed_state_parts = [<EOL>x + tf.cast(eps, v.dtype) * v<EOL>for x, eps, v<EOL>in zip(current_state_parts,<EOL>step_sizes,<EOL>proposed_momentum_parts)]<EOL>if state_gradients_are_stopped:<EOL><INDENT>proposed_state_parts = [tf.stop_gradient(x) for x in proposed_state_parts]<EOL><DEDENT>[<EOL>proposed_target_log_prob,<EOL>proposed_target_log_prob_grad_parts,<EOL>] = mcmc_util.maybe_call_fn_and_grads(<EOL>target_log_prob_fn,<EOL>proposed_state_parts)<EOL>if not proposed_target_log_prob.dtype.is_floating:<EOL><INDENT>raise TypeError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if any(g is None for g in proposed_target_log_prob_grad_parts):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>current_state_parts,<EOL>proposed_state_parts,<EOL>proposed_target_log_prob_grad_parts))<EOL><DEDENT>proposed_momentum_parts = [<EOL>v + <NUM_LIT:0.5> * tf.cast(eps, v.dtype) * g<EOL>for v, eps, g<EOL>in zip(proposed_momentum_parts,<EOL>step_sizes,<EOL>proposed_target_log_prob_grad_parts)]<EOL>return [<EOL>proposed_momentum_parts,<EOL>proposed_state_parts,<EOL>proposed_target_log_prob,<EOL>proposed_target_log_prob_grad_parts,<EOL>]<EOL><DEDENT>

Applies `num_leapfrog_steps` of the leapfrog integrator.

    Assumes a simple quadratic kinetic energy function: `0.5 ||momentum||**2`.

    #### Examples:

    ##### Simple quadratic potential.

    ```python
    import matplotlib.pyplot as plt
    %matplotlib inline
    import numpy as np
    import tensorflow as tf
    from tensorflow_probability.python.mcmc.hmc import _leapfrog_integrator_one_step  # pylint: disable=line-too-long
    tfd = tfp.distributions

    dims = 10
    num_iter = int(1e3)
    dtype = np.float32

    position = tf.placeholder(np.float32)
    momentum = tf.placeholder(np.float32)

    target_log_prob_fn = tfd.MultivariateNormalDiag(
        loc=tf.zeros(dims, dtype)).log_prob

    def _leapfrog_one_step(*args):
      # Closure representing computation done during each leapfrog step.
      return _leapfrog_integrator_one_step(
          target_log_prob_fn=target_log_prob_fn,
          independent_chain_ndims=0,
          step_sizes=[0.1],
          current_momentum_parts=args[0],
          current_state_parts=args[1],
          current_target_log_prob=args[2],
          current_target_log_prob_grad_parts=args[3])

    # Do leapfrog integration.
    [
        [next_momentum],
        [next_position],
        next_target_log_prob,
        next_target_log_prob_grad_parts,
    ] = tf.while_loop(
        cond=lambda *args: True,
        body=_leapfrog_one_step,
        loop_vars=[
          [momentum],
          [position],
          target_log_prob_fn(position),
          tf.gradients(target_log_prob_fn(position), position),
        ],
        maximum_iterations=3)

    momentum_ = np.random.randn(dims).astype(dtype)
    position_ = np.random.randn(dims).astype(dtype)
    positions = np.zeros([num_iter, dims], dtype)

    with tf.Session() as sess:
      for i in xrange(num_iter):
        position_, momentum_ = sess.run(
            [next_momentum, next_position],
            feed_dict={position: position_, momentum: momentum_})
        positions[i] = position_

    plt.plot(positions[:, 0]);  # Sinusoidal.
    ```

    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `*current_state_parts` and returns its (possibly unnormalized) log-density
        under the target distribution.
      independent_chain_ndims: Scalar `int` `Tensor` representing the number of
        leftmost `Tensor` dimensions which index independent chains.
      step_sizes: Python `list` of `Tensor`s representing the step size for the
        leapfrog integrator. Must broadcast with the shape of
        `current_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.
      current_momentum_parts: Tensor containing the value(s) of the momentum
        variable(s) to update.
      current_state_parts: Python `list` of `Tensor`s representing the current
        state(s) of the Markov chain(s). The first `independent_chain_ndims` of
        the `Tensor`(s) index different chains.
      current_target_log_prob: `Tensor` representing the value of
        `target_log_prob_fn(*current_state_parts)`. The only reason to specify
        this argument is to reduce TF graph size.
      current_target_log_prob_grad_parts: Python list of `Tensor`s representing
        gradient of `target_log_prob_fn(*current_state_parts`) wrt
        `current_state_parts`. Must have same shape as `current_state_parts`. The
        only reason to specify this argument is to reduce TF graph size.
      state_gradients_are_stopped: Python `bool` indicating that the proposed new
        state be run through `tf.stop_gradient`. This is particularly useful when
        combining optimization over samples from the HMC chain.
        Default value: `False` (i.e., do not apply `stop_gradient`).
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'hmc_leapfrog_integrator').

    Returns:
      proposed_momentum_parts: Updated value of the momentum.
      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`.
      proposed_target_log_prob: `Tensor` representing the value of
        `target_log_prob_fn` at `next_state`.
      proposed_target_log_prob_grad_parts: Gradient of `proposed_target_log_prob`
        wrt `next_state`.

    Raises:
      ValueError: if `len(momentum_parts) != len(state_parts)`.
      ValueError: if `len(state_parts) != len(step_sizes)`.
      ValueError: if `len(state_parts) != len(grads_target_log_prob)`.
      TypeError: if `not target_log_prob.dtype.is_floating`.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[independent_chain_ndims, current_momentums, proposed_momentums]):<EOL><INDENT>log_current_kinetic, log_proposed_kinetic = [], []<EOL>for current_momentum, proposed_momentum in zip(<EOL>current_momentums, proposed_momentums):<EOL><INDENT>axis = tf.range(independent_chain_ndims, tf.rank(current_momentum))<EOL>log_current_kinetic.append(_log_sum_sq(current_momentum, axis))<EOL>log_proposed_kinetic.append(_log_sum_sq(proposed_momentum, axis))<EOL><DEDENT>current_kinetic = <NUM_LIT:0.5> * tf.exp(<EOL>tf.reduce_logsumexp(<EOL>input_tensor=tf.stack(log_current_kinetic, axis=-<NUM_LIT:1>), axis=-<NUM_LIT:1>))<EOL>proposed_kinetic = <NUM_LIT:0.5> * tf.exp(<EOL>tf.reduce_logsumexp(<EOL>input_tensor=tf.stack(log_proposed_kinetic, axis=-<NUM_LIT:1>), axis=-<NUM_LIT:1>))<EOL>return mcmc_util.safe_sum([current_kinetic, -proposed_kinetic])<EOL><DEDENT>

Helper to `kernel` which computes the log acceptance-correction.

    A sufficient but not necessary condition for the existence of a stationary
    distribution, `p(x)`, is "detailed balance", i.e.:

    ```none
    p(x'|x) p(x) = p(x|x') p(x')
    ```

    In the Metropolis-Hastings algorithm, a state is proposed according to
    `g(x'|x)` and accepted according to `a(x'|x)`, hence
    `p(x'|x) = g(x'|x) a(x'|x)`.

    Inserting this into the detailed balance equation implies:

    ```none
        g(x'|x) a(x'|x) p(x) = g(x|x') a(x|x') p(x')
    ==> a(x'|x) / a(x|x') = p(x') / p(x) [g(x|x') / g(x'|x)]    (*)
    ```

    One definition of `a(x'|x)` which satisfies (*) is:

    ```none
    a(x'|x) = min(1, p(x') / p(x) [g(x|x') / g(x'|x)])
    ```

    (To see that this satisfies (*), notice that under this definition only at
    most one `a(x'|x)` and `a(x|x') can be other than one.)

    We call the bracketed term the "acceptance correction".

    In the case of UncalibratedHMC, the log acceptance-correction is not the log
    proposal-ratio. UncalibratedHMC augments the state-space with momentum, z.
    Assuming a standard Gaussian distribution for momentums, the chain eventually
    converges to:

    ```none
    p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
    ```

    Relating this back to Metropolis-Hastings parlance, for HMC we have:

    ```none
    p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
    g([x, z] | [x', z']) = g([x', z'] | [x, z])
    ```

    In other words, the MH bracketed term is `1`. However, because we desire to
    use a general MH framework, we can place the momentum probability ratio inside
    the metropolis-correction factor thus getting an acceptance probability:

    ```none
                         target_prob(x')
    accept_prob(x'|x) = -----------------  [exp(-0.5 z**2) / exp(-0.5 z'**2)]
                         target_prob(x)
    ```

    (Note: we actually need to handle the kinetic energy change at each leapfrog
    step, but this is the idea.)

    Args:
      current_momentums: `Tensor` representing the value(s) of the current
        momentum(s) of the state (parts).
      proposed_momentums: `Tensor` representing the value(s) of the proposed
        momentum(s) of the state (parts).
      independent_chain_ndims: Scalar `int` `Tensor` representing the number of
        leftmost `Tensor` dimensions which index independent chains.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'compute_log_acceptance_correction').

    Returns:
      log_acceptance_correction: `Tensor` representing the `log`
        acceptance-correction.  (See docstring for mathematical definition.)


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>if state_gradients_are_stopped:<EOL><INDENT>state_parts = [tf.stop_gradient(x) for x in state_parts]<EOL><DEDENT>target_log_prob, grads_target_log_prob = mcmc_util.maybe_call_fn_and_grads(<EOL>target_log_prob_fn,<EOL>state_parts,<EOL>target_log_prob,<EOL>grads_target_log_prob)<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>grads_target_log_prob,<EOL>]<EOL>

Helper which processes input args to meet list-like assumptions.


return tf.reduce_logsumexp(<EOL>input_tensor=<NUM_LIT> * tf.math.log(tf.abs(x)), axis=axis)<EOL>


if step_size_update_fn and store_parameters_in_results:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>impl = metropolis_hastings.MetropolisHastings(<EOL>inner_kernel=UncalibratedHamiltonianMonteCarlo(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>step_size=step_size,<EOL>num_leapfrog_steps=num_leapfrog_steps,<EOL>state_gradients_are_stopped=state_gradients_are_stopped,<EOL>seed=seed,<EOL>name='<STR_LIT>' if name is None else name,<EOL>store_parameters_in_results=store_parameters_in_results),<EOL>seed=seed)<EOL>parameters = impl.inner_kernel.parameters.copy()<EOL>parameters['<STR_LIT>'] = step_size_update_fn<EOL>self._impl = impl<EOL>self._parameters = parameters<EOL>

Initializes this transition kernel.

        Args:
          target_log_prob_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it's a list) and returns its
            (possibly unnormalized) log-density under the target distribution.
          step_size: `Tensor` or Python `list` of `Tensor`s representing the step
            size for the leapfrog integrator. Must broadcast with the shape of
            `current_state`. 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.
          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`.
          state_gradients_are_stopped: Python `bool` indicating that the proposed
            new state be run through `tf.stop_gradient`. This is particularly useful
            when combining optimization over samples from the HMC chain.
            Default value: `False` (i.e., do not apply `stop_gradient`).
          step_size_update_fn: Python `callable` taking current `step_size`
            (typically a `tf.Variable`) and `kernel_results` (typically
            `collections.namedtuple`) and returns updated step_size (`Tensor`s).
            Default value: `None` (i.e., do not update `step_size` automatically).
          seed: Python integer to seed the random number generator.
          store_parameters_in_results: If `True`, then `step_size` and
            `num_leapfrog_steps` are written to and read from eponymous fields in
            the kernel results objects returned from `one_step` and
            `bootstrap_results`. This allows wrapper kernels to adjust those
            parameters on the fly. This is incompatible with `step_size_update_fn`,
            which must be set to `None`.
          name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., 'hmc_kernel').


return self._impl.inner_kernel.step_size<EOL>

Returns the step_size parameter.

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

        Returns:
          step_size: A floating point `Tensor` or a list of such `Tensors`.


return self._impl.inner_kernel.num_leapfrog_steps<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.


previous_step_size_assign = (<EOL>[] if self.step_size_update_fn is None<EOL>else (previous_kernel_results.extra.step_size_assign<EOL>if mcmc_util.is_list_like(<EOL>previous_kernel_results.extra.step_size_assign)<EOL>else [previous_kernel_results.extra.step_size_assign]))<EOL>with tf.control_dependencies(previous_step_size_assign):<EOL><INDENT>next_state, kernel_results = self._impl.one_step(<EOL>current_state, previous_kernel_results)<EOL>if self.step_size_update_fn is not None:<EOL><INDENT>step_size_assign = self.step_size_update_fn(  <EOL>self.step_size, kernel_results)<EOL>kernel_results = kernel_results._replace(<EOL>extra=HamiltonianMonteCarloExtraKernelResults(<EOL>step_size_assign=step_size_assign))<EOL><DEDENT>return next_state, kernel_results<EOL><DEDENT>

Runs one iteration of Hamiltonian Monte Carlo.

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


kernel_results = self._impl.bootstrap_results(init_state)<EOL>if self.step_size_update_fn is not None:<EOL><INDENT>step_size_assign = self.step_size_update_fn(self.step_size, None)  <EOL>kernel_results = kernel_results._replace(<EOL>extra=HamiltonianMonteCarloExtraKernelResults(<EOL>step_size_assign=step_size_assign))<EOL><DEDENT>return kernel_results<EOL>

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


if seed is not None and tf.executing_eagerly():<EOL><INDENT>raise NotImplementedError('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>self._seed_stream = distributions.SeedStream(seed, '<STR_LIT>')<EOL>if not store_parameters_in_results:<EOL><INDENT>mcmc_util.warn_if_parameters_are_not_simple_tensors(<EOL>dict(step_size=step_size, num_leapfrog_steps=num_leapfrog_steps))<EOL><DEDENT>self._parameters = dict(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>step_size=step_size,<EOL>num_leapfrog_steps=num_leapfrog_steps,<EOL>state_gradients_are_stopped=state_gradients_are_stopped,<EOL>seed=seed,<EOL>name=name,<EOL>store_parameters_in_results=store_parameters_in_results,<EOL>)<EOL>

Initializes this transition kernel.

        Args:
          target_log_prob_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it's a list) and returns its
            (possibly unnormalized) log-density under the target distribution.
          step_size: `Tensor` or Python `list` of `Tensor`s representing the step
            size for the leapfrog integrator. Must broadcast with the shape of
            `current_state`. 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.
          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`.
          state_gradients_are_stopped: Python `bool` indicating that the proposed
            new state be run through `tf.stop_gradient`. This is particularly useful
            when combining optimization over samples from the HMC chain.
            Default value: `False` (i.e., do not apply `stop_gradient`).
          seed: Python integer to seed the random number generator.
          store_parameters_in_results: If `True`, then `step_size` and
            `num_leapfrog_steps` are written to and read from eponymous fields in
            the kernel results objects returned from `one_step` and
            `bootstrap_results`. This allows wrapper kernels to adjust those
            parameters on the fly.
          name: Python `str` name prefixed to Ops created by this function.
            Default value: `None` (i.e., 'hmc_kernel').


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


return self._parameters<EOL>


return getattr(kernel_results, '<STR_LIT>', None) is not None<EOL>

Returns `True` if `target_log_prob` is a member of input.


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

if inner_kernel.is_calibrated:<EOL><INDENT>warnings.warn('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>self._seed_stream = distributions.SeedStream(<EOL>seed, '<STR_LIT>')<EOL>self._parameters = dict(<EOL>inner_kernel=inner_kernel,<EOL>seed=seed,<EOL>name=name)<EOL>

Instantiates this object.

        Args:
          inner_kernel: `TransitionKernel`-like object which has
            `collections.namedtuple` `kernel_results` and which contains a
            `target_log_prob` member and optionally a `log_acceptance_correction`
            member.
          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., "mh_kernel").

        Returns:
          metropolis_hastings_kernel: Instance of `TransitionKernel` which wraps the
            input transition kernel with the Metropolis-Hastings algorithm.


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>[<EOL>proposed_state,<EOL>proposed_results,<EOL>] = self.inner_kernel.one_step(<EOL>current_state,<EOL>previous_kernel_results.accepted_results)<EOL>if (not has_target_log_prob(proposed_results) or<EOL>not has_target_log_prob(previous_kernel_results.accepted_results)):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>to_sum = [proposed_results.target_log_prob,<EOL>-previous_kernel_results.accepted_results.target_log_prob]<EOL>try:<EOL><INDENT>if (not mcmc_util.is_list_like(<EOL>proposed_results.log_acceptance_correction)<EOL>or proposed_results.log_acceptance_correction):<EOL><INDENT>to_sum.append(proposed_results.log_acceptance_correction)<EOL><DEDENT><DEDENT>except AttributeError:<EOL><INDENT>warnings.warn('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>log_accept_ratio = mcmc_util.safe_sum(<EOL>to_sum, name='<STR_LIT>')<EOL>log_uniform = tf.math.log(<EOL>tf.random.uniform(<EOL>shape=tf.shape(input=proposed_results.target_log_prob),<EOL>dtype=proposed_results.target_log_prob.dtype.base_dtype,<EOL>seed=self._seed_stream()))<EOL>is_accepted = log_uniform < log_accept_ratio<EOL>next_state = mcmc_util.choose(<EOL>is_accepted,<EOL>proposed_state,<EOL>current_state,<EOL>name='<STR_LIT>')<EOL>kernel_results = MetropolisHastingsKernelResults(<EOL>accepted_results=mcmc_util.choose(<EOL>is_accepted,<EOL>proposed_results,<EOL>previous_kernel_results.accepted_results,<EOL>name='<STR_LIT>'),<EOL>is_accepted=is_accepted,<EOL>log_accept_ratio=log_accept_ratio,<EOL>proposed_state=proposed_state,<EOL>proposed_results=proposed_results,<EOL>extra=[],<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.

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


with tf.compat.v1.name_scope(<EOL>name=mcmc_util.make_name(self.name, '<STR_LIT>', '<STR_LIT>'),<EOL>values=[init_state]):<EOL><INDENT>pkr = self.inner_kernel.bootstrap_results(init_state)<EOL>if not has_target_log_prob(pkr):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>')<EOL><DEDENT>x = pkr.target_log_prob<EOL>return MetropolisHastingsKernelResults(<EOL>accepted_results=pkr,<EOL>is_accepted=tf.ones_like(x, dtype=tf.bool),<EOL>log_accept_ratio=tf.zeros_like(x),<EOL>proposed_state=init_state,<EOL>proposed_results=pkr,<EOL>extra=[],<EOL>)<EOL><DEDENT>

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

        Args:
          init_state: `Tensor` or Python `list` of `Tensor`s representing the
            initial state(s) of the Markov chain(s).

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

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


if not kernel.is_calibrated:<EOL><INDENT>warnings.warn("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>with tf.compat.v1.name_scope(<EOL>name, "<STR_LIT>",<EOL>[num_results, num_burnin_steps, num_steps_between_results]):<EOL><INDENT>num_results = tf.convert_to_tensor(<EOL>value=num_results, dtype=tf.int32, name="<STR_LIT>")<EOL>num_burnin_steps = tf.convert_to_tensor(<EOL>value=num_burnin_steps, dtype=tf.int32, name="<STR_LIT>")<EOL>num_steps_between_results = tf.convert_to_tensor(<EOL>value=num_steps_between_results,<EOL>dtype=tf.int32,<EOL>name="<STR_LIT>")<EOL>current_state = tf.nest.map_structure(<EOL>lambda x: tf.convert_to_tensor(value=x, name="<STR_LIT>"),<EOL>current_state)<EOL>if previous_kernel_results is None:<EOL><INDENT>previous_kernel_results = kernel.bootstrap_results(current_state)<EOL><DEDENT>if trace_fn is None:<EOL><INDENT>trace_fn = lambda *args: ()<EOL>no_trace = True<EOL><DEDENT>else:<EOL><INDENT>no_trace = False<EOL><DEDENT>if trace_fn is sample_chain.__defaults__[<NUM_LIT:4>]:<EOL><INDENT>warnings.warn("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>def _trace_scan_fn(state_and_results, num_steps):<EOL><INDENT>next_state, current_kernel_results = mcmc_util.smart_for_loop(<EOL>loop_num_iter=num_steps,<EOL>body_fn=kernel.one_step,<EOL>initial_loop_vars=list(state_and_results),<EOL>parallel_iterations=parallel_iterations)<EOL>return next_state, current_kernel_results<EOL><DEDENT>(_, final_kernel_results), (all_states, trace) = mcmc_util.trace_scan(<EOL>loop_fn=_trace_scan_fn,<EOL>initial_state=(current_state, previous_kernel_results),<EOL>elems=tf.one_hot(<EOL>indices=<NUM_LIT:0>,<EOL>depth=num_results,<EOL>on_value=<NUM_LIT:1> + num_burnin_steps,<EOL>off_value=<NUM_LIT:1> + num_steps_between_results,<EOL>dtype=tf.int32),<EOL>trace_fn=lambda state_and_results: (state_and_results[<NUM_LIT:0>],<EOL>trace_fn(*state_and_results)),<EOL>parallel_iterations=parallel_iterations)<EOL>if return_final_kernel_results:<EOL><INDENT>return CheckpointableStatesAndTrace(<EOL>all_states=all_states,<EOL>trace=trace,<EOL>final_kernel_results=final_kernel_results)<EOL><DEDENT>else:<EOL><INDENT>if no_trace:<EOL><INDENT>return all_states<EOL><DEDENT>else:<EOL><INDENT>return StatesAndTrace(all_states=all_states, trace=trace)<EOL><DEDENT><DEDENT><DEDENT>

Implements Markov chain Monte Carlo via repeated `TransitionKernel` steps.

    This function samples from an Markov chain at `current_state` and whose
    stationary distribution is governed by the supplied `TransitionKernel`
    instance (`kernel`).

    This function can sample from multiple chains, in parallel. (Whether or not
    there are multiple chains is dictated by the `kernel`.)

    The `current_state` can be represented as a single `Tensor` or a `list` of
    `Tensors` which collectively represent the current state.

    Since MCMC states are correlated, it is sometimes desirable to produce
    additional intermediate states, and then discard them, ending up with a set of
    states with decreased autocorrelation.  See [Owen (2017)][1]. Such "thinning"
    is made possible by setting `num_steps_between_results > 0`. The chain then
    takes `num_steps_between_results` extra steps between the steps that make it
    into the results. The extra steps are never materialized (in calls to
    `sess.run`), and thus do not increase memory requirements.

    Warning: when setting a `seed` in the `kernel`, ensure that `sample_chain`'s
    `parallel_iterations=1`, otherwise results will not be reproducible.

    In addition to returning the chain state, this function supports tracing of
    auxiliary variables used by the kernel. The traced values are selected by
    specifying `trace_fn`. By default, all kernel results are traced but in the
    future the default will be changed to no results being traced, so plan
    accordingly. See below for some examples of this feature.

    Args:
      num_results: Integer number of Markov chain draws.
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s).
      previous_kernel_results: A `Tensor` or a nested collection of `Tensor`s
        representing internal calculations made within the previous call to this
        function (or as returned by `bootstrap_results`).
      kernel: An instance of `tfp.mcmc.TransitionKernel` which implements one step
        of the Markov chain.
      num_burnin_steps: Integer number of chain steps to take before starting to
        collect results.
        Default value: 0 (i.e., no burn-in).
      num_steps_between_results: Integer number of chain steps between collecting
        a result. Only one out of every `num_steps_between_samples + 1` steps is
        included in the returned results.  The number of returned chain states is
        still equal to `num_results`.  Default value: 0 (i.e., no thinning).
      trace_fn: A callable that takes in the current chain state and the previous
        kernel results and return a `Tensor` or a nested collection of `Tensor`s
        that is then traced along with the chain state.
      return_final_kernel_results: If `True`, then the final kernel results are
        returned alongside the chain state and the trace specified by the
        `trace_fn`.
      parallel_iterations: The number of iterations allowed to run in parallel. It
        must be a positive integer. See `tf.while_loop` for more details.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "mcmc_sample_chain").

    Returns:
      checkpointable_states_and_trace: if `return_final_kernel_results` is
        `True`. The return value is an instance of
        `CheckpointableStatesAndTrace`.
      all_states: if `return_final_kernel_results` is `False` and `trace_fn` is
        `None`. The return value is a `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` but with a prepended
        `num_results`-size dimension.
      states_and_trace: if `return_final_kernel_results` is `False` and
        `trace_fn` is not `None`. The return value is an instance of
        `StatesAndTrace`.

    #### Examples

    ##### Sample from a diagonal-variance Gaussian.

    I.e.,

    ```none
    for i=1..n:
      x[i] ~ MultivariateNormal(loc=0, scale=diag(true_stddev))  # likelihood
    ```

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

    dims = 10
    true_stddev = np.sqrt(np.linspace(1., 3., dims))
    likelihood = tfd.MultivariateNormalDiag(loc=0., scale_diag=true_stddev)

    states = tfp.mcmc.sample_chain(
        num_results=1000,
        num_burnin_steps=500,
        current_state=tf.zeros(dims),
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=likelihood.log_prob,
          step_size=0.5,
          num_leapfrog_steps=2),
        trace_fn=None)

    sample_mean = tf.reduce_mean(states, axis=0)
    # ==> approx all zeros

    sample_stddev = tf.sqrt(tf.reduce_mean(
        tf.squared_difference(states, sample_mean),
        axis=0))
    # ==> approx equal true_stddev
    ```

    ##### Sampling from factor-analysis posteriors with known factors.

    I.e.,

    ```none
    # prior
    w ~ MultivariateNormal(loc=0, scale=eye(d))
    for i=1..n:
      # likelihood
      x[i] ~ Normal(loc=w^T F[i], scale=1)
    ```

    where `F` denotes factors.

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

    # Specify model.
    def make_prior(dims):
      return tfd.MultivariateNormalDiag(
          loc=tf.zeros(dims))

    def make_likelihood(weights, factors):
      return tfd.MultivariateNormalDiag(
          loc=tf.matmul(weights, factors, adjoint_b=True))

    def joint_log_prob(num_weights, factors, x, w):
      return (make_prior(num_weights).log_prob(w) +
              make_likelihood(w, factors).log_prob(x))

    def unnormalized_log_posterior(w):
      # Posterior is proportional to: `p(W, X=x | factors)`.
      return joint_log_prob(num_weights, factors, x, w)

    # Setup data.
    num_weights = 10 # == d
    num_factors = 40 # == n
    num_chains = 100

    weights = make_prior(num_weights).sample(1)
    factors = tf.random_normal([num_factors, num_weights])
    x = make_likelihood(weights, factors).sample()

    # Sample from Hamiltonian Monte Carlo Markov Chain.

    # Get `num_results` samples from `num_chains` independent chains.
    chains_states, kernels_results = tfp.mcmc.sample_chain(
        num_results=1000,
        num_burnin_steps=500,
        current_state=tf.zeros([num_chains, num_weights], name='init_weights'),
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=unnormalized_log_posterior,
          step_size=0.1,
          num_leapfrog_steps=2))

    # Compute sample stats.
    sample_mean = tf.reduce_mean(chains_states, axis=[0, 1])
    # ==> approx equal to weights

    sample_var = tf.reduce_mean(
        tf.squared_difference(chains_states, sample_mean),
        axis=[0, 1])
    # ==> less than 1
    ```

    ##### Custom tracing functions.

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

    likelihood = tfd.Normal(loc=0., scale=1.)

    def sample_chain(trace_fn):
      return tfp.mcmc.sample_chain(
        num_results=1000,
        num_burnin_steps=500,
        current_state=0.,
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=likelihood.log_prob,
          step_size=0.5,
          num_leapfrog_steps=2),
        trace_fn=trace_fn)

    def trace_log_accept_ratio(states, previous_kernel_results):
      return previous_kernel_results.log_accept_ratio

    def trace_everything(states, previous_kernel_results):
      return previous_kernel_results

    _, log_accept_ratio = sample_chain(trace_fn=trace_log_accept_ratio)
    _, kernel_results = sample_chain(trace_fn=trace_everything)

    acceptance_prob = tf.exp(tf.minimum(log_accept_ratio_, 0.))
    # Equivalent to, but more efficient than:
    acceptance_prob = tf.exp(tf.minimum(kernel_results.log_accept_ratio_, 0.))
    ```

    #### References

    [1]: Art B. Owen. Statistically efficient thinning of a Markov chain sampler.
         _Technical Report_, 2017.
         http://statweb.stanford.edu/~owen/reports/bestthinning.pdf


Validates that target_accept_prob is in (0, 1).


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


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

Helper which uses graph seed if using TFE.