class_docstr: No-U-Turn Sampler kernel, which provides an efficient and convenient way
  to run Hamiltonian Monte Carlo. The number of steps taken by the
  integrator is dynamically adjusted on each call to ``sample`` to ensure
  an optimal length for the Hamiltonian trajectory [1]. As such, the samples
  generated will typically have lower autocorrelation than those generated
  by the :class:`~pyro.infer.mcmc.HMC` kernel. Optionally, the NUTS kernel
  also provides the ability to adapt step size during the warmup phase.
  
  Refer to the `baseball example <https://github.com/pyro-ppl/pyro/blob/dev/examples/baseball.py>`_
  to see how to do Bayesian inference in Pyro using NUTS.
  
  **References**
  
  [1] `The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo`,
      Matthew D. Hoffman, and Andrew Gelman.
  [2] `A Conceptual Introduction to Hamiltonian Monte Carlo`,
      Michael Betancourt
  [3] `Slice Sampling`,
      Radford M. Neal
  
  :param model: Python callable containing Pyro primitives.
  :param potential_fn: Python callable calculating potential energy with input
      is a dict of real support parameters.
  :param float step_size: Determines the size of a single step taken by the
      verlet integrator while computing the trajectory using Hamiltonian
      dynamics. If not specified, it will be set to 1.
  :param bool adapt_step_size: A flag to decide if we want to adapt step_size
      during warm-up phase using Dual Averaging scheme.
  :param bool adapt_mass_matrix: A flag to decide if we want to adapt mass
      matrix during warm-up phase using Welford scheme.
  :param bool full_mass: A flag to decide if mass matrix is dense or diagonal.
  :param bool use_multinomial_sampling: A flag to decide if we want to sample
      candidates along its trajectory using "multinomial sampling" or using
      "slice sampling". Slice sampling is used in the original NUTS paper [1],
      while multinomial sampling is suggested in [2]. By default, this flag is
      set to True. If it is set to `False`, NUTS uses slice sampling.
  :param dict transforms: Optional dictionary that specifies a transform
      for a sample site with constrained support to unconstrained space. The
      transform should be invertible, and implement `log_abs_det_jacobian`.
      If not specified and the model has sites with constrained support,
      automatic transformations will be applied, as specified in
      :mod:`torch.distributions.constraint_registry`.
  :param int max_plate_nesting: Optional bound on max number of nested
      :func:`pyro.plate` contexts. This is required if model contains
      discrete sample sites that can be enumerated over in parallel.
  :param bool jit_compile: Optional parameter denoting whether to use
      the PyTorch JIT to trace the log density computation, and use this
      optimized executable trace in the integrator.
  :param dict jit_options: A dictionary contains optional arguments for
      :func:`torch.jit.trace` function.
  :param bool ignore_jit_warnings: Flag to ignore warnings from the JIT
      tracer when ``jit_compile=True``. Default is False.
  :param float target_accept_prob: Target acceptance probability of step size
      adaptation scheme. Increasing this value will lead to a smaller step size,
      so the sampling will be slower but more robust. Default to 0.8.
  :param int max_tree_depth: Max depth of the binary tree created during the doubling
      scheme of NUTS sampler. Default to 10.
  :param callable init_strategy: A per-site initialization function.
      See :ref:`autoguide-initialization` section for available functions.
  
  Example:
  
      >>> true_coefs = torch.tensor([1., 2., 3.])
      >>> data = torch.randn(2000, 3)
      >>> dim = 3
      >>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample()
      >>>
      >>> def model(data):
      ...     coefs_mean = torch.zeros(dim)
      ...     coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(3)))
      ...     y = pyro.sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels)
      ...     return y
      >>>
      >>> nuts_kernel = NUTS(model, adapt_step_size=True)
      >>> mcmc = MCMC(nuts_kernel, num_samples=500, warmup_steps=300)
      >>> mcmc.run(data)
      >>> mcmc.get_samples()['beta'].mean(0)  # doctest: +SKIP
      tensor([ 0.9221,  1.9464,  2.9228])
code: def __init__(
      self,
      model=None,
      potential_fn=None,
      step_size=1,
      adapt_step_size=True,
      adapt_mass_matrix=True,
      full_mass=False,
      use_multinomial_sampling=True,
      transforms=None,
      max_plate_nesting=None,
      jit_compile=False,
      jit_options=None,
      ignore_jit_warnings=False,
      target_accept_prob=0.8,
      max_tree_depth=10,
      init_strategy=init_to_uniform,
  ):
      super().__init__(
          model,
          potential_fn,
          step_size,
          adapt_step_size=adapt_step_size,
          adapt_mass_matrix=adapt_mass_matrix,
          full_mass=full_mass,
          transforms=transforms,
          max_plate_nesting=max_plate_nesting,
          jit_compile=jit_compile,
          jit_options=jit_options,
          ignore_jit_warnings=ignore_jit_warnings,
          target_accept_prob=target_accept_prob,
          init_strategy=init_strategy,
      )
      self.use_multinomial_sampling = use_multinomial_sampling
      self._max_tree_depth = max_tree_depth
      # There are three conditions to stop doubling process:
      #     + Tree is becoming too big.
      #     + The trajectory is making a U-turn.
      #     + The probability of the states becoming negligible: p(z, r) << u,
      # here u is the "slice" variable introduced at the `self.sample(...)` method.
      # Denote E_p = -log p(z, r), E_u = -log u, the third condition is equivalent to
      #     sliced_energy := E_p - E_u > some constant =: max_sliced_energy.
      # This also suggests the notion "diverging" in the implemenation:
      #     when the energy E_p diverges from E_u too much, we stop doubling.
      # Here, as suggested in [1], we set dE_max = 1000.
      self._max_sliced_energy = 1000
  
file_path: pyro/pyro/infer/mcmc/nuts.py
imports: from collections import namedtuple
  import pyro
  import pyro.distributions
  from pyro.distributions.util import scalar_like
  from pyro.infer.autoguide import init_to_uniform
  from pyro.infer.mcmc.hmc import HMC
  from pyro.ops.integrator import potential_grad, velocity_verlet
  from pyro.util import optional, torch_isnan
label_desc: Use image node_id 1 to create a new NUTS object from inherited base classes: HMC with example: obj = NUTS(model, potential_fn, step_size, adapt_step_size, adapt_mass_matrix, full_mass, use_multinomial_sampling, transforms, max_plate_nesting, jit_compile, jit_options, ignore_jit_warnings, target_accept_prob, max_tree_depth, init_strategy)