experimental/no_u_turn_sampler/nuts.py
def _build_tree(value_and_gradients_fn,
                current_state,
                current_target_log_prob,
                current_grads_target_log_prob,
                current_momentum,
                direction,
                depth,
                step_size,
                log_slice_sample,
                max_simulation_error=1000.,
                seed=None):
  """Builds a tree at a given tree depth and at a given state.

  The `current` state is immediately adjacent to, but outside of,
  the subtrajectory spanned by the returned `forward` and `reverse` states.

  Args:
    value_and_gradients_fn: Python callable which takes an argument like
      `*current_state` and returns a tuple of its (possibly unnormalized)
      log-density under the target distribution and its gradient with respect to
      each state.
    current_state: List of `Tensor`s representing the current states of the
      NUTS trajectory.
    current_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `current_state`.
    current_grads_target_log_prob: List of `Tensor`s representing gradient of
      `current_target_log_prob` with respect to `current_state`. Must have same
      shape as `current_state`.
    current_momentum: List of `Tensor`s representing the momentums of
      `current_state`. Must have same shape as `current_state`.
    direction: int that is either -1 or 1. It determines whether to perform
      leapfrog integration backwards (reverse) or forward in time respectively.
    depth: non-negative int that indicates how deep of a tree to build.
      Each call to `_build_tree` takes `2**depth` leapfrog steps.
    step_size: List of `Tensor`s representing the step sizes for the leapfrog
      integrator. Must have same shape as `current_state`.
    log_slice_sample: The log of an auxiliary slice variable. It is used
      together with `max_simulation_error` to avoid simulating trajectories with
      too much numerical error.
    max_simulation_error: Maximum simulation error to tolerate before
      terminating the trajectory. Simulation error is the
      `log_slice_sample` minus the log-joint probability at the simulated state.
    seed: Integer to seed the random number generator.

  Returns:
    reverse_state: List of `Tensor`s representing the "reverse" states of the
      NUTS trajectory. Has same shape as `current_state`.
    reverse_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `reverse_state`.
    reverse_grads_target_log_prob: List of `Tensor`s representing gradient of
      `reverse_target_log_prob` with respect to `reverse_state`. Has same shape
      as `reverse_state`.
    reverse_momentum: List of `Tensor`s representing the momentums of
      `reverse_state`. Has same shape as `reverse_state`.
    forward_state: List of `Tensor`s representing the "forward" states of the
      NUTS trajectory. Has same shape as `current_state`.
    forward_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `forward_state`.
    forward_grads_target_log_prob: List of `Tensor`s representing gradient of
      `forward_target_log_prob` with respect to `forward_state`. Has same shape
      as `forward_state`.
    forward_momentum: List of `Tensor`s representing the momentums of
      `forward_state`. Has same shape as `forward_state`.
    next_state: List of `Tensor`s representing the next states of the NUTS
      trajectory. Has same shape as `current_state`.
    next_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at `next_state`.
    next_grads_target_log_prob: List of `Tensor`s representing the gradient of
      `next_target_log_prob` with respect to `next_state`.
    num_states: Number of acceptable candidate states in the subtree. A state is
      acceptable if it is "in the slice", that is, if its log-joint probability
      with its momentum is greater than `log_slice_sample`.
    continue_trajectory: bool determining whether to continue the simulation
      trajectory. The trajectory is continued if no U-turns are encountered
      within the built subtree, and if the log-probability accumulation due to
      integration error does not exceed `max_simulation_error`.
  """
  if depth == 0:  # base case
    # Take a leapfrog step. Terminate the tree-building if the simulation
    # error from the leapfrog integrator is too large. States discovered by
    # continuing the simulation are likely to have very low probability.
    [
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
    ] = _leapfrog(
        value_and_gradients_fn=value_and_gradients_fn,
        current_state=current_state,
        current_grads_target_log_prob=current_grads_target_log_prob,
        current_momentum=current_momentum,
        step_size=direction * step_size)
    next_log_joint = _log_joint(next_target_log_prob, next_momentum)
    num_states = tf.cast(next_log_joint > log_slice_sample, dtype=tf.int32)
    continue_trajectory = (next_log_joint >
                           log_slice_sample - max_simulation_error)
    return [
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        num_states,
        continue_trajectory,
    ]

  # Build a tree at the current state.
  seed_stream = tfd.SeedStream(seed, "build_tree")
  [
      reverse_state,
      reverse_target_log_prob,
      reverse_grads_target_log_prob,
      reverse_momentum,
      forward_state,
      forward_target_log_prob,
      forward_grads_target_log_prob,
      forward_momentum,
      next_state,
      next_target_log_prob,
      next_grads_target_log_prob,
      num_states,
      continue_trajectory,
  ] = _build_tree(value_and_gradients_fn=value_and_gradients_fn,
                  current_state=current_state,
                  current_target_log_prob=current_target_log_prob,
                  current_grads_target_log_prob=current_grads_target_log_prob,
                  current_momentum=current_momentum,
                  direction=direction,
                  depth=depth - 1,
                  step_size=step_size,
                  log_slice_sample=log_slice_sample,
                  seed=seed_stream())
  if continue_trajectory:
    # If the just-built subtree did not terminate, build a second subtree at
    # the forward or reverse state, as appropriate.
    if direction < 0:
      [
          reverse_state,
          reverse_target_log_prob,
          reverse_grads_target_log_prob,
          reverse_momentum,
          _,
          _,
          _,
          _,
          far_state,
          far_target_log_prob,
          far_grads_target_log_prob,
          far_num_states,
          far_continue_trajectory,
      ] = _build_tree(
          value_and_gradients_fn=value_and_gradients_fn,
          current_state=reverse_state,
          current_target_log_prob=reverse_target_log_prob,
          current_grads_target_log_prob=reverse_grads_target_log_prob,
          current_momentum=reverse_momentum,
          direction=direction,
          depth=depth - 1,
          step_size=step_size,
          log_slice_sample=log_slice_sample,
          seed=seed_stream())
    else:
      [
          _,
          _,
          _,
          _,
          forward_state,
          forward_target_log_prob,
          forward_grads_target_log_prob,
          forward_momentum,
          far_state,
          far_target_log_prob,
          far_grads_target_log_prob,
          far_num_states,
          far_continue_trajectory,
      ] = _build_tree(
          value_and_gradients_fn=value_and_gradients_fn,
          current_state=forward_state,
          current_target_log_prob=forward_target_log_prob,
          current_grads_target_log_prob=forward_grads_target_log_prob,
          current_momentum=forward_momentum,
          direction=direction,
          depth=depth - 1,
          step_size=step_size,
          log_slice_sample=log_slice_sample,
          seed=seed_stream())

    # Propose either `next_state` (which came from the first subtree and so is
    # nearby) or the new forward/reverse state (which came from the second
    # subtree and so is far away).
    num_states += far_num_states
    accept_far_state = _random_bernoulli(
        [],
        probs=far_num_states / num_states,
        dtype=tf.bool,
        seed=seed_stream())
    if accept_far_state:
      next_state = far_state
      next_target_log_prob = far_target_log_prob
      next_grads_target_log_prob = far_grads_target_log_prob

    # Continue the NUTS trajectory if the far subtree did not terminate either,
    # and if the reverse-most and forward-most states do not exhibit a U-turn.
    has_no_u_turn = tf.logical_and(
        _has_no_u_turn(forward_state, reverse_state, forward_momentum),
        _has_no_u_turn(forward_state, reverse_state, reverse_momentum))
    continue_trajectory = far_continue_trajectory and has_no_u_turn

  return [
      reverse_state,
      reverse_target_log_prob,
      reverse_grads_target_log_prob,
      reverse_momentum,
      forward_state,
      forward_target_log_prob,
      forward_grads_target_log_prob,
      forward_momentum,
      next_state,
      next_target_log_prob,
      next_grads_target_log_prob,
      num_states,
      continue_trajectory,
  ]
def _build_tree(value_and_gradients_fn,
                current_state,
                current_target_log_prob,
                current_grads_target_log_prob,
                current_momentum,
                direction,
                depth,
                step_size,
                log_slice_sample,
                max_simulation_error=1000.,
                seed=None):
  """Builds a tree at a given tree depth and at a given state.

  The `current` state is immediately adjacent to, but outside of,
  the subtrajectory spanned by the returned `forward` and `reverse` states.

  Args:
    value_and_gradients_fn: Python callable which takes an argument like
      `*current_state` and returns a tuple of its (possibly unnormalized)
      log-density under the target distribution and its gradient with respect to
      each state.
    current_state: List of `Tensor`s representing the current states of the
      NUTS trajectory.
    current_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `current_state`.
    current_grads_target_log_prob: List of `Tensor`s representing gradient of
      `current_target_log_prob` with respect to `current_state`. Must have same
      shape as `current_state`.
    current_momentum: List of `Tensor`s representing the momentums of
      `current_state`. Must have same shape as `current_state`.
    direction: int that is either -1 or 1. It determines whether to perform
      leapfrog integration backwards (reverse) or forward in time respectively.
    depth: non-negative int that indicates how deep of a tree to build.
      Each call to `_build_tree` takes `2**depth` leapfrog steps.
    step_size: List of `Tensor`s representing the step sizes for the leapfrog
      integrator. Must have same shape as `current_state`.
    log_slice_sample: The log of an auxiliary slice variable. It is used
      together with `max_simulation_error` to avoid simulating trajectories with
      too much numerical error.
    max_simulation_error: Maximum simulation error to tolerate before
      terminating the trajectory. Simulation error is the
      `log_slice_sample` minus the log-joint probability at the simulated state.
    seed: Integer to seed the random number generator.

  Returns:
    reverse_state: List of `Tensor`s representing the "reverse" states of the
      NUTS trajectory. Has same shape as `current_state`.
    reverse_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `reverse_state`.
    reverse_grads_target_log_prob: List of `Tensor`s representing gradient of
      `reverse_target_log_prob` with respect to `reverse_state`. Has same shape
      as `reverse_state`.
    reverse_momentum: List of `Tensor`s representing the momentums of
      `reverse_state`. Has same shape as `reverse_state`.
    forward_state: List of `Tensor`s representing the "forward" states of the
      NUTS trajectory. Has same shape as `current_state`.
    forward_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `forward_state`.
    forward_grads_target_log_prob: List of `Tensor`s representing gradient of
      `forward_target_log_prob` with respect to `forward_state`. Has same shape
      as `forward_state`.
    forward_momentum: List of `Tensor`s representing the momentums of
      `forward_state`. Has same shape as `forward_state`.
    next_state: List of `Tensor`s representing the next states of the NUTS
      trajectory. Has same shape as `current_state`.
    next_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at `next_state`.
    next_grads_target_log_prob: List of `Tensor`s representing the gradient of
      `next_target_log_prob` with respect to `next_state`.
    num_states: Number of acceptable candidate states in the subtree. A state is
      acceptable if it is "in the slice", that is, if its log-joint probability
      with its momentum is greater than `log_slice_sample`.
    continue_trajectory: bool determining whether to continue the simulation
      trajectory. The trajectory is continued if no U-turns are encountered
      within the built subtree, and if the log-probability accumulation due to
      integration error does not exceed `max_simulation_error`.
  """
  if depth == 0:  # base case
    # Take a leapfrog step. Terminate the tree-building if the simulation
    # error from the leapfrog integrator is too large. States discovered by
    # continuing the simulation are likely to have very low probability.
    [
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
    ] = _leapfrog(
        value_and_gradients_fn=value_and_gradients_fn,
        current_state=current_state,
        current_grads_target_log_prob=current_grads_target_log_prob,
        current_momentum=current_momentum,
        step_size=direction * step_size)
    next_log_joint = _log_joint(next_target_log_prob, next_momentum)
    num_states = tf.cast(next_log_joint > log_slice_sample, dtype=tf.int32)
    continue_trajectory = (next_log_joint >
                           log_slice_sample - max_simulation_error)
    return [
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        num_states,
        continue_trajectory,
    ]

  # Build a tree at the current state.
  seed_stream = tfd.SeedStream(seed, "build_tree")
  [
      reverse_state,
      reverse_target_log_prob,
      reverse_grads_target_log_prob,
      reverse_momentum,
      forward_state,
      forward_target_log_prob,
      forward_grads_target_log_prob,
      forward_momentum,
      next_state,
      next_target_log_prob,
      next_grads_target_log_prob,
      num_states,
      continue_trajectory,
  ] = _build_tree(value_and_gradients_fn=value_and_gradients_fn,
                  current_state=current_state,
                  current_target_log_prob=current_target_log_prob,
                  current_grads_target_log_prob=current_grads_target_log_prob,
                  current_momentum=current_momentum,
                  direction=direction,
                  depth=depth - 1,
                  step_size=step_size,
                  log_slice_sample=log_slice_sample,
                  seed=seed_stream())
  if continue_trajectory:
    # If the just-built subtree did not terminate, build a second subtree at
    # the forward or reverse state, as appropriate.
    if direction < 0:
      [
          reverse_state,
          reverse_target_log_prob,
          reverse_grads_target_log_prob,
          reverse_momentum,
          _,
          _,
          _,
          _,
          far_state,
          far_target_log_prob,
          far_grads_target_log_prob,
          far_num_states,
          far_continue_trajectory,
      ] = _build_tree(
          value_and_gradients_fn=value_and_gradients_fn,
          current_state=reverse_state,
          current_target_log_prob=reverse_target_log_prob,
          current_grads_target_log_prob=reverse_grads_target_log_prob,
          current_momentum=reverse_momentum,
          direction=direction,
          depth=depth - 1,
          step_size=step_size,
          log_slice_sample=log_slice_sample,
          seed=seed_stream())
    else:
      [
          _,
          _,
          _,
          _,
          forward_state,
          forward_target_log_prob,
          forward_grads_target_log_prob,
          forward_momentum,
          far_state,
          far_target_log_prob,
          far_grads_target_log_prob,
          far_num_states,
          far_continue_trajectory,
      ] = _build_tree(
          value_and_gradients_fn=value_and_gradients_fn,
          current_state=forward_state,
          current_target_log_prob=forward_target_log_prob,
          current_grads_target_log_prob=forward_grads_target_log_prob,
          current_momentum=forward_momentum,
          direction=direction,
          depth=depth - 1,
          step_size=step_size,
          log_slice_sample=log_slice_sample,
          seed=seed_stream())

    # Propose either `next_state` (which came from the first subtree and so is
    # nearby) or the new forward/reverse state (which came from the second
    # subtree and so is far away).
    num_states += far_num_states
    accept_far_state = _random_bernoulli(
        [],
        probs=far_num_states / num_states,
        dtype=tf.bool,
        seed=seed_stream())
    if accept_far_state:
      next_state = far_state
      next_target_log_prob = far_target_log_prob
      next_grads_target_log_prob = far_grads_target_log_prob

    # Continue the NUTS trajectory if the far subtree did not terminate either,
    # and if the reverse-most and forward-most states do not exhibit a U-turn.
    has_no_u_turn = tf.logical_and(
        _has_no_u_turn(forward_state, reverse_state, forward_momentum),
        _has_no_u_turn(forward_state, reverse_state, reverse_momentum))
    continue_trajectory = far_continue_trajectory and has_no_u_turn

  return [
      reverse_state,
      reverse_target_log_prob,
      reverse_grads_target_log_prob,
      reverse_momentum,
      forward_state,
      forward_target_log_prob,
      forward_grads_target_log_prob,
      forward_momentum,
      next_state,
      next_target_log_prob,
      next_grads_target_log_prob,
      num_states,
      continue_trajectory,
  ]
Builds a tree at a given tree depth and at a given state.

  The `current` state is immediately adjacent to, but outside of,
  the subtrajectory spanned by the returned `forward` and `reverse` states.

  Args:
    value_and_gradients_fn: Python callable which takes an argument like
      `*current_state` and returns a tuple of its (possibly unnormalized)
      log-density under the target distribution and its gradient with respect to
      each state.
    current_state: List of `Tensor`s representing the current states of the
      NUTS trajectory.
    current_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `current_state`.
    current_grads_target_log_prob: List of `Tensor`s representing gradient of
      `current_target_log_prob` with respect to `current_state`. Must have same
      shape as `current_state`.
    current_momentum: List of `Tensor`s representing the momentums of
      `current_state`. Must have same shape as `current_state`.
    direction: int that is either -1 or 1. It determines whether to perform
      leapfrog integration backwards (reverse) or forward in time respectively.
    depth: non-negative int that indicates how deep of a tree to build.
      Each call to `_build_tree` takes `2**depth` leapfrog steps.
    step_size: List of `Tensor`s representing the step sizes for the leapfrog
      integrator. Must have same shape as `current_state`.
    log_slice_sample: The log of an auxiliary slice variable. It is used
      together with `max_simulation_error` to avoid simulating trajectories with
      too much numerical error.
    max_simulation_error: Maximum simulation error to tolerate before
      terminating the trajectory. Simulation error is the
      `log_slice_sample` minus the log-joint probability at the simulated state.
    seed: Integer to seed the random number generator.

  Returns:
    reverse_state: List of `Tensor`s representing the "reverse" states of the
      NUTS trajectory. Has same shape as `current_state`.
    reverse_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `reverse_state`.
    reverse_grads_target_log_prob: List of `Tensor`s representing gradient of
      `reverse_target_log_prob` with respect to `reverse_state`. Has same shape
      as `reverse_state`.
    reverse_momentum: List of `Tensor`s representing the momentums of
      `reverse_state`. Has same shape as `reverse_state`.
    forward_state: List of `Tensor`s representing the "forward" states of the
      NUTS trajectory. Has same shape as `current_state`.
    forward_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `forward_state`.
    forward_grads_target_log_prob: List of `Tensor`s representing gradient of
      `forward_target_log_prob` with respect to `forward_state`. Has same shape
      as `forward_state`.
    forward_momentum: List of `Tensor`s representing the momentums of
      `forward_state`. Has same shape as `forward_state`.
    next_state: List of `Tensor`s representing the next states of the NUTS
      trajectory. Has same shape as `current_state`.
    next_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at `next_state`.
    next_grads_target_log_prob: List of `Tensor`s representing the gradient of
      `next_target_log_prob` with respect to `next_state`.
    num_states: Number of acceptable candidate states in the subtree. A state is
      acceptable if it is "in the slice", that is, if its log-joint probability
      with its momentum is greater than `log_slice_sample`.
    continue_trajectory: bool determining whether to continue the simulation
      trajectory. The trajectory is continued if no U-turns are encountered
      within the built subtree, and if the log-probability accumulation due to
      integration error does not exceed `max_simulation_error`.