class_docstr: Runs one step of the No U-Turn Sampler.
  
  The No U-Turn Sampler (NUTS) is an adaptive variant of the Hamiltonian Monte
  Carlo (HMC) method for MCMC.  NUTS adapts the distance traveled in response to
  the curvature of the target density.  Conceptually, one proposal consists of
  reversibly evolving a trajectory through the sample space, continuing until
  that trajectory turns back on itself (hence the name, "No U-Turn").  This
  class implements one random NUTS step from a given
  `current_state`.  Mathematical details and derivations can be found in
  [Hoffman, Gelman (2011)][1].
  
  The `one_step` function can update multiple chains in parallel. It assumes
  that a prefix of leftmost dimensions of `current_state` index independent
  chain states (and are therefore updated independently).  The output of
  `target_log_prob_fn(*current_state)` should sum log-probabilities across all
  event dimensions.  Slices along the rightmost dimensions may have different
  target distributions; for example, `current_state[0][0, ...]` could have a
  different target distribution from `current_state[0][1, ...]`.  These
  semantics are governed by `target_log_prob_fn(*current_state)`. (The number of
  independent chains is `tf.size(target_log_prob_fn(*current_state))`.)
  
  TODO(axch): Examples (e.g., a la HMC).  For them to be sensible, need to
  pick sensible step sizes, or implement step size adaptation, or both.
  
  #### References
  
  [1] Matthew D. Hoffman, Andrew Gelman.  The No-U-Turn Sampler: Adaptively
  Setting Path Lengths in Hamiltonian Monte Carlo.  2011.
  https://arxiv.org/pdf/1111.4246.pdf.
docstr: Runs one iteration of the No U-Turn Sampler.
  
  Args:
    current_state: `Tensor` or Python `list` of `Tensor`s representing the
      current state(s) of the Markov chain(s). The first `r` dimensions index
      independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
    previous_kernel_results: `collections.namedtuple` containing `Tensor`s
      representing values from previous calls to this function (or from the
      `bootstrap_results` function.)
  
  Returns:
    next_state: `Tensor` or Python list of `Tensor`s representing the state(s)
      of the Markov chain(s) after taking `self.num_trajectories_per_step`
      steps. Has same type and shape as `current_state`.
    kernel_results: `collections.namedtuple` of internal calculations used to
      advance the chain.
docstr_tok: ["Runs","one","iteration","of","the","No","U-Turn","Sampler",".","Args",":","current_state",":","`","Tensor","`","or","Python","`","list","`","of","`","Tensor","`","s","representing","the","current","state","(","s",")","of","the","Markov","chain","(","s",")",".","The","first","`","r","`","dimensions","index","independent","chains",",","`","r","=","tf.rank","(","target_log_prob_fn","(","*","current_state",")",")","`",".","previous_kernel_results",":","`","collections.namedtuple","`","containing","`","Tensor","`","s","representing","values","from","previous","calls","to","this","function","(","or","from","the","`","bootstrap_results","`","function",".",")","Returns",":","next_state",":","`","Tensor","`","or","Python","list","of","`","Tensor","`","s","representing","the","state","(","s",")","of","the","Markov","chain","(","s",")","after","taking","`","self.num_trajectories_per_step","`","steps",".","Has","same","type","and","shape","as","`","current_state","`",".","kernel_results",":","`","collections.namedtuple","`","of","internal","calculations","used","to","advance","the","chain","."]
code: def one_step(self, current_state, previous_kernel_results):
      """Runs one iteration of the No U-Turn Sampler.
  
      Args:
        current_state: `Tensor` or Python `list` of `Tensor`s representing the
          current state(s) of the Markov chain(s). The first `r` dimensions index
          independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
        previous_kernel_results: `collections.namedtuple` containing `Tensor`s
          representing values from previous calls to this function (or from the
          `bootstrap_results` function.)
  
      Returns:
        next_state: `Tensor` or Python list of `Tensor`s representing the state(s)
          of the Markov chain(s) after taking `self.num_trajectories_per_step`
          steps. Has same type and shape as `current_state`.
        kernel_results: `collections.namedtuple` of internal calculations used to
          advance the chain.
      """
      if self.stackless and not tf.executing_eagerly():
          raise ValueError(
              "Cannot use stackless auto-batching in graph mode."
          )
      current_target_log_prob = previous_kernel_results.target_log_prob
      current_grads_log_prob = (
          previous_kernel_results.grads_target_log_prob
      )
      leapfrogs_taken = previous_kernel_results.leapfrogs_taken
      leapfrogs_computed = previous_kernel_results.leapfrogs_computed
      with tf.name_scope(self.name):
          unwrap_state_list = False
          with tf.name_scope("initialize"):
              if not tf.nest.is_nested(current_state):
                  unwrap_state_list = True
                  current_state = [current_state]
              current_state = [
                  tf.convert_to_tensor(value=s) for s in current_state
              ]
              step_size = self.step_size
              if not tf.nest.is_nested(step_size):
                  step_size = [step_size]
              step_size = [
                  tf.convert_to_tensor(value=s) for s in step_size
              ]
              if len(step_size) == 1:
                  step_size = step_size * len(current_state)
              if len(step_size) != len(current_state):
                  raise ValueError(
                      "Expected either one step size or {} (size of "
                      "`current_state`), but found {}".format(
                          len(current_state), len(step_size)
                      )
                  )
  
          num_steps = tf.constant(
              [self.num_trajectories_per_step], dtype=tf.int64
          )
          if self.backend is None:
              if self._seed_stream() is not None:
                  # The user wanted reproducible results; limit the parallel iterations
                  backend = ab.TensorFlowBackend(
                      while_parallel_iterations=1
                  )
              else:
                  backend = ab.TensorFlowBackend()
          else:
              backend = self.backend
          # The `dry_run` and `max_stack_depth` arguments are added by the
          # @ctx.batch decorator, confusing pylint.
          # pylint: disable=unexpected-keyword-arg
          (
              (
                  next_state,
                  next_target_log_prob,
                  next_grads_target_log_prob,
              ),
              new_leapfrogs,
          ) = self.many_steps(
              num_steps,
              current_state,
              current_target_log_prob,
              current_grads_log_prob,
              step_size,
              tf.zeros_like(leapfrogs_taken),  # leapfrogs
              dry_run=not self.use_auto_batching,
              stackless=self.stackless,
              backend=backend,
              max_stack_depth=self.max_tree_depth + 4,
              block_code_cache=self._block_code_cache,
          )
  
          if unwrap_state_list:
              next_state = next_state[0]
          return next_state, NUTSKernelResults(
              next_target_log_prob,
              next_grads_target_log_prob,
              leapfrogs_taken + new_leapfrogs,
              leapfrogs_computed
              + tf.math.reduce_max(input_tensor=new_leapfrogs),
          )
  
imports: import collections
  import functools
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python import random
  from tensorflow_probability.python.experimental import auto_batching
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.util.seed_stream import SeedStream
label_desc: Use image node_id 4 for calling the NoUTurnSampler obj's underlying member method code with example usage: obj.one_step(current_state, previous_kernel_results) and returns: next_state,  NUTSKernelResults