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 _build_tree(
      self,
      z,
      r,
      z_grads,
      log_slice,
      direction,
      tree_depth,
      energy_current,
  ):
      if tree_depth == 0:
          return self._build_basetree(
              z, r, z_grads, log_slice, direction, energy_current
          )
  
      # build the first half of tree
      half_tree = self._build_tree(
          z,
          r,
          z_grads,
          log_slice,
          direction,
          tree_depth - 1,
          energy_current,
      )
      z_proposal = half_tree.z_proposal
      z_proposal_pe = half_tree.z_proposal_pe
      z_proposal_grads = half_tree.z_proposal_grads
  
      # Check conditions to stop doubling. If we meet that condition,
      #     there is no need to build the other tree.
      if half_tree.turning or half_tree.diverging:
          return half_tree
  
      # Else, build remaining half of tree.
      # If we are going to the right, start from the right leaf of the first half.
      if direction == 1:
          z = half_tree.z_right
          r = half_tree.r_right
          z_grads = half_tree.z_right_grads
      else:  # otherwise, start from the left leaf of the first half
          z = half_tree.z_left
          r = half_tree.r_left
          z_grads = half_tree.z_left_grads
      other_half_tree = self._build_tree(
          z,
          r,
          z_grads,
          log_slice,
          direction,
          tree_depth - 1,
          energy_current,
      )
  
      if self.use_multinomial_sampling:
          tree_weight = _logaddexp(
              half_tree.weight, other_half_tree.weight
          )
      else:
          tree_weight = half_tree.weight + other_half_tree.weight
      sum_accept_probs = (
          half_tree.sum_accept_probs + other_half_tree.sum_accept_probs
      )
      num_proposals = (
          half_tree.num_proposals + other_half_tree.num_proposals
      )
      r_sum = {
          site_names: half_tree.r_sum[site_names]
          + other_half_tree.r_sum[site_names]
          for site_names in self.inverse_mass_matrix
      }
  
      # The probability of that proposal belongs to which half of tree
      #     is computed based on the weights of each half.
      if self.use_multinomial_sampling:
          other_half_tree_prob = (
              other_half_tree.weight - tree_weight
          ).exp()
      else:
          # For the special case that the weights of each half are both 0,
          #   we choose the proposal from the first half
          #   (any is fine, because the probability of picking it at the end is 0!).
          other_half_tree_prob = (
              other_half_tree.weight / tree_weight
              if tree_weight > 0
              else scalar_like(tree_weight, 0.0)
          )
      is_other_half_tree = pyro.sample(
          "is_other_half_tree",
          dist.Bernoulli(probs=other_half_tree_prob),
      )
  
      if is_other_half_tree == 1:
          z_proposal = other_half_tree.z_proposal
          z_proposal_pe = other_half_tree.z_proposal_pe
          z_proposal_grads = other_half_tree.z_proposal_grads
  
      # leaves of the full tree are determined by the direction
      if direction == 1:
          z_left = half_tree.z_left
          r_left = half_tree.r_left
          r_left_unscaled = half_tree.r_left_unscaled
          z_left_grads = half_tree.z_left_grads
          z_right = other_half_tree.z_right
          r_right = other_half_tree.r_right
          r_right_unscaled = other_half_tree.r_right_unscaled
          z_right_grads = other_half_tree.z_right_grads
      else:
          z_left = other_half_tree.z_left
          r_left = other_half_tree.r_left
          r_left_unscaled = other_half_tree.r_left_unscaled
          z_left_grads = other_half_tree.z_left_grads
          z_right = half_tree.z_right
          r_right = half_tree.r_right
          r_right_unscaled = half_tree.r_right_unscaled
          z_right_grads = half_tree.z_right_grads
  
      # We already check if first half tree is turning. Now, we check
      #     if the other half tree or full tree are turning.
      turning = other_half_tree.turning or self._is_turning(
          r_left_unscaled, r_right_unscaled, r_sum
      )
  
      # The divergence is checked by the second half tree (the first half is already checked).
      diverging = other_half_tree.diverging
  
      return _TreeInfo(
          z_left,
          r_left,
          r_left_unscaled,
          z_left_grads,
          z_right,
          r_right,
          r_right_unscaled,
          z_right_grads,
          z_proposal,
          z_proposal_pe,
          z_proposal_grads,
          r_sum,
          tree_weight,
          turning,
          diverging,
          sum_accept_probs,
          num_proposals,
      )
  
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 4 for calling the NUTS obj's underlying member method code with example usage: obj._build_tree(z, r, z_grads, log_slice, direction, tree_depth, energy_current) and returns: _TreeInfo, self, half_tree