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] and [Betancourt (2018)][2].
  
  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))`.)
  
  #### 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.
  
  [2]: Michael Betancourt. A Conceptual Introduction to Hamiltonian Monte Carlo.
  _arXiv preprint arXiv:1701.02434_, 2018. https://arxiv.org/abs/1701.02434
class_docstr_tok: ["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","]","and","[","Betancourt","(","2018",")","]","[","2","]",".","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",")",")","`",".",")","#","#","#","#","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",".","[","2","]",":","Michael","Betancourt",".","A","Conceptual","Introduction","to","Hamiltonian","Monte","Carlo",".","_arXiv","preprint","arXiv:1701.02434_",",","2018.","https",":","\/\/arxiv.org\/abs\/1701.02434"]
docstr: 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.
    momentum_distribution: A `tfp.distributions.Distribution` instance to draw
      momentum from. Defaults to normal distributions with identity
      covariance.
    max_tree_depth: Maximum depth of the tree implicitly built by NUTS. The
      maximum number of leapfrog steps is bounded by `2**max_tree_depth` i.e.
      the number of nodes in a binary tree `max_tree_depth` nodes deep. The
      default setting of 10 takes up to 1024 leapfrog steps.
    max_energy_diff: Scaler threshold of energy differences at each leapfrog,
      divergence samples are defined as leapfrog steps that exceed this
      threshold. Default to 1000.
    unrolled_leapfrog_steps: The number of leapfrogs to unroll per tree
      expansion step. Applies a direct linear multipler to the maximum
      trajectory length implied by max_tree_depth. Defaults to 1.
    parallel_iterations: The number of iterations allowed to run in parallel.
      It must be a positive integer. See `tf.while_loop` for more details.
    experimental_shard_axis_names: A structure of string names indicating how
      members of the state are sharded.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'PreconditionedNoUTurnSampler').
code: def __init__(
      self,
      target_log_prob_fn,
      step_size,
      momentum_distribution=None,
      max_tree_depth=10,
      max_energy_diff=1000.0,
      unrolled_leapfrog_steps=1,
      parallel_iterations=10,
      experimental_shard_axis_names=None,
      name=None,
  ):
      """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.
        momentum_distribution: A `tfp.distributions.Distribution` instance to draw
          momentum from. Defaults to normal distributions with identity
          covariance.
        max_tree_depth: Maximum depth of the tree implicitly built by NUTS. The
          maximum number of leapfrog steps is bounded by `2**max_tree_depth` i.e.
          the number of nodes in a binary tree `max_tree_depth` nodes deep. The
          default setting of 10 takes up to 1024 leapfrog steps.
        max_energy_diff: Scaler threshold of energy differences at each leapfrog,
          divergence samples are defined as leapfrog steps that exceed this
          threshold. Default to 1000.
        unrolled_leapfrog_steps: The number of leapfrogs to unroll per tree
          expansion step. Applies a direct linear multipler to the maximum
          trajectory length implied by max_tree_depth. Defaults to 1.
        parallel_iterations: The number of iterations allowed to run in parallel.
          It must be a positive integer. See `tf.while_loop` for more details.
        experimental_shard_axis_names: A structure of string names indicating how
          members of the state are sharded.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., 'PreconditionedNoUTurnSampler').
      """
      with tf.name_scope(
          name or "PreconditionedNoUTurnSampler"
      ) as name:
          # Process `max_tree_depth` argument.
          max_tree_depth = tf.get_static_value(max_tree_depth)
          if max_tree_depth is None or max_tree_depth < 1:
              raise ValueError(
                  "max_tree_depth must be known statically and >= 1 but was "
                  "{}".format(max_tree_depth)
              )
          self._max_tree_depth = max_tree_depth
  
          # Compute parameters derived from `max_tree_depth`.
          instruction_array = build_tree_uturn_instruction(
              max_tree_depth, init_memory=-1
          )
          [
              write_instruction_numpy,
              read_instruction_numpy,
          ] = generate_efficient_write_read_instruction(
              instruction_array
          )
  
          # TensorArray version of the read/write instruction need to be created
          # within the function call to be compatible with XLA. Here we store the
          # numpy version of the instruction and convert it to TensorArray later.
          self._write_instruction = write_instruction_numpy
          self._read_instruction = read_instruction_numpy
  
          # Process all other arguments.
          self._target_log_prob_fn = target_log_prob_fn
          self._step_size = step_size
  
          self._parameters = dict(
              target_log_prob_fn=target_log_prob_fn,
              step_size=step_size,
              max_tree_depth=max_tree_depth,
              max_energy_diff=max_energy_diff,
              unrolled_leapfrog_steps=unrolled_leapfrog_steps,
              parallel_iterations=parallel_iterations,
              experimental_shard_axis_names=experimental_shard_axis_names,
              name=name,
          )
          self.momentum_distribution = momentum_distribution
          self._parallel_iterations = parallel_iterations
          self._unrolled_leapfrog_steps = unrolled_leapfrog_steps
          self._name = name
          self._max_energy_diff = max_energy_diff
  
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import preconditioning_utils
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math.generic import log_add_exp
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow_probability.python.mcmc.kernel import TransitionKernel
label_desc: Use image node_id 1 to create a new PreconditionedNoUTurnSampler object from inherited base classes: TransitionKernel with example: obj = PreconditionedNoUTurnSampler(target_log_prob_fn, step_size, momentum_distribution, max_tree_depth, max_energy_diff, unrolled_leapfrog_steps, parallel_iterations, experimental_shard_axis_names, name)