def fit_with_hmc(model,<EOL>observed_time_series,<EOL>num_results=<NUM_LIT:100>,<EOL>num_warmup_steps=<NUM_LIT:50>,<EOL>num_leapfrog_steps=<NUM_LIT:15>,<EOL>initial_state=None,<EOL>initial_step_size=None,<EOL>chain_batch_shape=(),<EOL>num_variational_steps=<NUM_LIT>,<EOL>variational_optimizer=None,<EOL>seed=None,<EOL>name=None):
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