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: 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.  Due
      to limitations of the underlying auto-batching system,
      target_log_prob_fn may be invoked with junk data at some batch indexes,
      which it must process without crashing.  (The results at those indexes
      are ignored).
    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.
    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-1`
      i.e. the number of nodes in a binary tree `max_tree_depth` nodes deep.
      The default setting of 10 takes up to 1023 leapfrog steps.
    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. This
      parameter can be useful for amortizing the auto-batching control flow
      overhead.
    num_trajectories_per_step: Python `int` giving the number of NUTS
      trajectories to run as "one" step.  Setting this higher than 1 may be
      favorable for performance by giving the autobatching system the
      opportunity to batch gradients across consecutive trajectories.  The
      intermediate samples are thinned: only the last sample from the run (in
      each batch member) is returned.
    use_auto_batching: Boolean.  If `False`, do not invoke the auto-batching
      system; operate on batch size 1 only.
    stackless: Boolean.  If `True`, invoke the stackless version of
      the auto-batching system.  Only works in Eager mode.
    backend: Auto-batching backend object. Falls back to a default
      TensorFlowBackend().
    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., 'nuts_kernel').
docstr_tok: ["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",".","Due","to","limitations","of","the","underlying","auto-batching","system",",","target_log_prob_fn","may","be","invoked","with","junk","data","at","some","batch","indexes",",","which","it","must","process","without","crashing",".","(","The","results","at","those","indexes","are","ignored",")",".","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",".","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-1","`","i.e",".","the","number","of","nodes","in","a","binary","tree","`","max_tree_depth","`","nodes","deep",".","The","default","setting","of","10","takes","up","to","1023","leapfrog","steps",".","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",".","This","parameter","can","be","useful","for","amortizing","the","auto-batching","control","flow","overhead",".","num_trajectories_per_step",":","Python","`","int","`","giving","the","number","of","NUTS","trajectories","to","run","as","``","one","''","step",".","Setting","this","higher","than","1","may","be","favorable","for","performance","by","giving","the","autobatching","system","the","opportunity","to","batch","gradients","across","consecutive","trajectories",".","The","intermediate","samples","are","thinned",":","only","the","last","sample","from","the","run","(","in","each","batch","member",")","is","returned",".","use_auto_batching",":","Boolean",".","If","`","False","`",",","do","not","invoke","the","auto-batching","system",";","operate","on","batch","size","1","only",".","stackless",":","Boolean",".","If","`","True","`",",","invoke","the","stackless","version","of","the","auto-batching","system",".","Only","works","in","Eager","mode",".","backend",":","Auto-batching","backend","object",".","Falls","back","to","a","default","TensorFlowBackend","(",")",".","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.",",","'nuts_kernel","'",")","."]
code: def __init__(
      self,
      target_log_prob_fn,
      step_size,
      max_tree_depth=10,
      unrolled_leapfrog_steps=1,
      num_trajectories_per_step=1,
      use_auto_batching=True,
      stackless=False,
      backend=None,
      seed=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.  Due
          to limitations of the underlying auto-batching system,
          target_log_prob_fn may be invoked with junk data at some batch indexes,
          which it must process without crashing.  (The results at those indexes
          are ignored).
        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.
        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-1`
          i.e. the number of nodes in a binary tree `max_tree_depth` nodes deep.
          The default setting of 10 takes up to 1023 leapfrog steps.
        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. This
          parameter can be useful for amortizing the auto-batching control flow
          overhead.
        num_trajectories_per_step: Python `int` giving the number of NUTS
          trajectories to run as "one" step.  Setting this higher than 1 may be
          favorable for performance by giving the autobatching system the
          opportunity to batch gradients across consecutive trajectories.  The
          intermediate samples are thinned: only the last sample from the run (in
          each batch member) is returned.
        use_auto_batching: Boolean.  If `False`, do not invoke the auto-batching
          system; operate on batch size 1 only.
        stackless: Boolean.  If `True`, invoke the stackless version of
          the auto-batching system.  Only works in Eager mode.
        backend: Auto-batching backend object. Falls back to a default
          TensorFlowBackend().
        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., 'nuts_kernel').
      """
      self._parameters = dict(locals())
      del self._parameters["self"]
      self.target_log_prob_fn = target_log_prob_fn
      self.step_size = step_size
      if max_tree_depth < 1:
          raise ValueError(
              "max_tree_depth must be >= 1 but was {}".format(
                  max_tree_depth
              )
          )
      self.max_tree_depth = max_tree_depth
      self.unrolled_leapfrog_steps = unrolled_leapfrog_steps
      self.num_trajectories_per_step = num_trajectories_per_step
      self.use_auto_batching = use_auto_batching
      self.stackless = stackless
      self.backend = backend
      self._seed_stream = SeedStream(seed, "nuts_one_step")
      self.name = "nuts_kernel" if name is None else name
      # TODO(b/125544625): Identify why we need `use_gradient_tape=True`, i.e.,
      # what's different between `tape.gradient` and `tf.gradient`.
      value_and_gradients_fn = lambda *args: gradient.value_and_gradient(  # pylint: disable=g-long-lambda
          self.target_log_prob_fn, args, use_gradient_tape=True
      )
      self.value_and_gradients_fn = _embed_no_none_gradient_check(
          value_and_gradients_fn
      )
      max_tree_edges = max_tree_depth - 1
      self.many_steps, self.autobatch_context = _make_evolve_trajectory(
          self.value_and_gradients_fn,
          max_tree_edges,
          unrolled_leapfrog_steps,
          self._seed_stream,
      )
      self._block_code_cache = {}
  
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 1 to create a new NoUTurnSampler object from inherited base classes: kernel_base with example: obj = NoUTurnSampler(target_log_prob_fn, step_size, max_tree_depth, unrolled_leapfrog_steps, num_trajectories_per_step, use_auto_batching, stackless, backend, seed, name)