experimental/no_u_turn_sampler/nuts.py
def kernel(target_log_prob_fn,
           current_state,
           step_size,
           seed=None,
           current_target_log_prob=None,
           current_grads_target_log_prob=None,
           name=None):
  """Simulates a No-U-Turn Sampler (NUTS) trajectory.

  Args:
    target_log_prob_fn: Python callable which takes an argument like
      `*current_state` and returns its (possibly unnormalized) log-density under
      the target distribution.
    current_state: List of `Tensor`s representing the states to simulate from.
    step_size: List of `Tensor`s representing the step sizes for the leapfrog
      integrator. Must have same shape as `current_state`.
    seed: Integer to seed the random number generator.
    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`.
    name: A name for the operation.

  Returns:
    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`.

  Raises:
    NotImplementedError: If the execution mode is not eager.
  """
  if not tf.executing_eagerly():
    raise NotImplementedError("`kernel` is only available in Eager mode.")

  with tf.compat.v1.name_scope(
      name,
      default_name="nuts_kernel",
      values=[
          current_state, step_size, seed, current_target_log_prob,
          current_grads_target_log_prob
      ]):
    with tf.compat.v1.name_scope("initialize"):
      current_state = [tf.convert_to_tensor(value=s) for s in current_state]
      step_size = [tf.convert_to_tensor(value=s) for s in step_size]
      value_and_gradients_fn = lambda *args: tfp.math.value_and_gradient(  # pylint: disable=g-long-lambda
          target_log_prob_fn, args)
      value_and_gradients_fn = _embed_no_none_gradient_check(
          value_and_gradients_fn)
      if (current_target_log_prob is None or
          current_grads_target_log_prob is None):
        (current_target_log_prob,
         current_grads_target_log_prob) = value_and_gradients_fn(*current_state)

      seed_stream = tfd.SeedStream(seed, "nuts_kernel")
      current_momentum = []
      for state_tensor in current_state:
        momentum_tensor = tf.random.normal(
            shape=tf.shape(input=state_tensor),
            dtype=state_tensor.dtype,
            seed=seed_stream())
        current_momentum.append(momentum_tensor)

      # Draw a slice variable u ~ Uniform(0, p(initial state, initial
      # momentum)) and compute log u. For numerical stability, we perform this
      # in log space where log u = log (u' * p(...)) = log u' + log
      # p(...) and u' ~ Uniform(0, 1).
      log_slice_sample = tf.math.log(tf.random.uniform([], seed=seed_stream()))
      log_slice_sample += _log_joint(current_target_log_prob,
                                     current_momentum)

      # Initialize loop variables. It comprises a collection of information
      # about a "reverse" state, a collection of information about a "forward"
      # state, a collection of information about the next state,
      # the trajectory's tree depth, the number of candidate states, and
      # whether to continue the trajectory.
      reverse_state = current_state
      reverse_target_log_prob = current_target_log_prob
      reverse_grads_target_log_prob = current_grads_target_log_prob
      reverse_momentum = current_momentum
      forward_state = current_state
      forward_target_log_prob = current_target_log_prob
      forward_grads_target_log_prob = current_grads_target_log_prob
      forward_momentum = current_momentum
      next_state = current_state
      next_target_log_prob = current_target_log_prob
      next_grads_target_log_prob = current_grads_target_log_prob
      depth = 0
      num_states = 1
      continue_trajectory = True

    while continue_trajectory:
      # Grow the No-U-Turn Sampler trajectory by choosing a random direction and
      # simulating Hamiltonian dynamics in that direction. This extends either
      # the forward or reverse state.
      direction = tfp.math.random_rademacher([], seed=seed_stream())
      if direction < 0:
        [
            reverse_state,
            reverse_target_log_prob,
            reverse_grads_target_log_prob,
            reverse_momentum,
            _,
            _,
            _,
            _,
            next_state_in_subtree,
            next_target_log_prob_in_subtree,
            next_grads_target_log_prob_in_subtree,
            num_states_in_subtree,
            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,
            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,
            next_state_in_subtree,
            next_target_log_prob_in_subtree,
            next_grads_target_log_prob_in_subtree,
            num_states_in_subtree,
            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,
            step_size=step_size,
            log_slice_sample=log_slice_sample,
            seed=seed_stream())

      if continue_trajectory:
        # If the built tree did not terminate, accept the tree's next state
        # with a certain probability.
        accept_state_in_subtree = _random_bernoulli(
            [],
            probs=tf.minimum(1., num_states_in_subtree / num_states),
            dtype=tf.bool,
            seed=seed_stream())
        if accept_state_in_subtree:
          next_state = next_state_in_subtree
          next_target_log_prob = next_target_log_prob_in_subtree
          next_grads_target_log_prob = next_grads_target_log_prob_in_subtree

      # Continue the NUTS trajectory if the tree-building did not terminate, 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 = continue_trajectory and has_no_u_turn
      num_states += num_states_in_subtree
      depth += 1

    return next_state, next_target_log_prob, next_grads_target_log_prob
def kernel(target_log_prob_fn,
           current_state,
           step_size,
           seed=None,
           current_target_log_prob=None,
           current_grads_target_log_prob=None,
           name=None):
  """Simulates a No-U-Turn Sampler (NUTS) trajectory.

  Args:
    target_log_prob_fn: Python callable which takes an argument like
      `*current_state` and returns its (possibly unnormalized) log-density under
      the target distribution.
    current_state: List of `Tensor`s representing the states to simulate from.
    step_size: List of `Tensor`s representing the step sizes for the leapfrog
      integrator. Must have same shape as `current_state`.
    seed: Integer to seed the random number generator.
    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`.
    name: A name for the operation.

  Returns:
    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`.

  Raises:
    NotImplementedError: If the execution mode is not eager.
  """
  if not tf.executing_eagerly():
    raise NotImplementedError("`kernel` is only available in Eager mode.")

  with tf.compat.v1.name_scope(
      name,
      default_name="nuts_kernel",
      values=[
          current_state, step_size, seed, current_target_log_prob,
          current_grads_target_log_prob
      ]):
    with tf.compat.v1.name_scope("initialize"):
      current_state = [tf.convert_to_tensor(value=s) for s in current_state]
      step_size = [tf.convert_to_tensor(value=s) for s in step_size]
      value_and_gradients_fn = lambda *args: tfp.math.value_and_gradient(  # pylint: disable=g-long-lambda
          target_log_prob_fn, args)
      value_and_gradients_fn = _embed_no_none_gradient_check(
          value_and_gradients_fn)
      if (current_target_log_prob is None or
          current_grads_target_log_prob is None):
        (current_target_log_prob,
         current_grads_target_log_prob) = value_and_gradients_fn(*current_state)

      seed_stream = tfd.SeedStream(seed, "nuts_kernel")
      current_momentum = []
      for state_tensor in current_state:
        momentum_tensor = tf.random.normal(
            shape=tf.shape(input=state_tensor),
            dtype=state_tensor.dtype,
            seed=seed_stream())
        current_momentum.append(momentum_tensor)

      # Draw a slice variable u ~ Uniform(0, p(initial state, initial
      # momentum)) and compute log u. For numerical stability, we perform this
      # in log space where log u = log (u' * p(...)) = log u' + log
      # p(...) and u' ~ Uniform(0, 1).
      log_slice_sample = tf.math.log(tf.random.uniform([], seed=seed_stream()))
      log_slice_sample += _log_joint(current_target_log_prob,
                                     current_momentum)

      # Initialize loop variables. It comprises a collection of information
      # about a "reverse" state, a collection of information about a "forward"
      # state, a collection of information about the next state,
      # the trajectory's tree depth, the number of candidate states, and
      # whether to continue the trajectory.
      reverse_state = current_state
      reverse_target_log_prob = current_target_log_prob
      reverse_grads_target_log_prob = current_grads_target_log_prob
      reverse_momentum = current_momentum
      forward_state = current_state
      forward_target_log_prob = current_target_log_prob
      forward_grads_target_log_prob = current_grads_target_log_prob
      forward_momentum = current_momentum
      next_state = current_state
      next_target_log_prob = current_target_log_prob
      next_grads_target_log_prob = current_grads_target_log_prob
      depth = 0
      num_states = 1
      continue_trajectory = True

    while continue_trajectory:
      # Grow the No-U-Turn Sampler trajectory by choosing a random direction and
      # simulating Hamiltonian dynamics in that direction. This extends either
      # the forward or reverse state.
      direction = tfp.math.random_rademacher([], seed=seed_stream())
      if direction < 0:
        [
            reverse_state,
            reverse_target_log_prob,
            reverse_grads_target_log_prob,
            reverse_momentum,
            _,
            _,
            _,
            _,
            next_state_in_subtree,
            next_target_log_prob_in_subtree,
            next_grads_target_log_prob_in_subtree,
            num_states_in_subtree,
            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,
            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,
            next_state_in_subtree,
            next_target_log_prob_in_subtree,
            next_grads_target_log_prob_in_subtree,
            num_states_in_subtree,
            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,
            step_size=step_size,
            log_slice_sample=log_slice_sample,
            seed=seed_stream())

      if continue_trajectory:
        # If the built tree did not terminate, accept the tree's next state
        # with a certain probability.
        accept_state_in_subtree = _random_bernoulli(
            [],
            probs=tf.minimum(1., num_states_in_subtree / num_states),
            dtype=tf.bool,
            seed=seed_stream())
        if accept_state_in_subtree:
          next_state = next_state_in_subtree
          next_target_log_prob = next_target_log_prob_in_subtree
          next_grads_target_log_prob = next_grads_target_log_prob_in_subtree

      # Continue the NUTS trajectory if the tree-building did not terminate, 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 = continue_trajectory and has_no_u_turn
      num_states += num_states_in_subtree
      depth += 1

    return next_state, next_target_log_prob, next_grads_target_log_prob
Simulates a No-U-Turn Sampler (NUTS) trajectory.

  Args:
    target_log_prob_fn: Python callable which takes an argument like
      `*current_state` and returns its (possibly unnormalized) log-density under
      the target distribution.
    current_state: List of `Tensor`s representing the states to simulate from.
    step_size: List of `Tensor`s representing the step sizes for the leapfrog
      integrator. Must have same shape as `current_state`.
    seed: Integer to seed the random number generator.
    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`.
    name: A name for the operation.

  Returns:
    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`.

  Raises:
    NotImplementedError: If the execution mode is not eager.