def kernel(target_log_prob_fn,<EOL>current_state,<EOL>step_size,<EOL>seed=None,<EOL>current_target_log_prob=None,<EOL>current_grads_target_log_prob=None,<EOL>name=None):
if not tf.executing_eagerly():<EOL><INDENT>raise NotImplementedError("<STR_LIT>")<EOL><DEDENT>with tf.compat.v1.name_scope(<EOL>name,<EOL>default_name="<STR_LIT>",<EOL>values=[<EOL>current_state, step_size, seed, current_target_log_prob,<EOL>current_grads_target_log_prob<EOL>]):<EOL><INDENT>with tf.compat.v1.name_scope("<STR_LIT>"):<EOL><INDENT>current_state = [tf.convert_to_tensor(value=s) for s in current_state]<EOL>step_size = [tf.convert_to_tensor(value=s) for s in step_size]<EOL>value_and_gradients_fn = lambda *args: tfp.math.value_and_gradient(  <EOL>target_log_prob_fn, args)<EOL>value_and_gradients_fn = _embed_no_none_gradient_check(<EOL>value_and_gradients_fn)<EOL>if (current_target_log_prob is None or<EOL>current_grads_target_log_prob is None):<EOL><INDENT>(current_target_log_prob,<EOL>current_grads_target_log_prob) = value_and_gradients_fn(*current_state)<EOL><DEDENT>seed_stream = tfd.SeedStream(seed, "<STR_LIT>")<EOL>current_momentum = []<EOL>for state_tensor in current_state:<EOL><INDENT>momentum_tensor = tf.random.normal(<EOL>shape=tf.shape(input=state_tensor),<EOL>dtype=state_tensor.dtype,<EOL>seed=seed_stream())<EOL>current_momentum.append(momentum_tensor)<EOL><DEDENT>log_slice_sample = tf.math.log(tf.random.uniform([], seed=seed_stream()))<EOL>log_slice_sample += _log_joint(current_target_log_prob,<EOL>current_momentum)<EOL>reverse_state = current_state<EOL>reverse_target_log_prob = current_target_log_prob<EOL>reverse_grads_target_log_prob = current_grads_target_log_prob<EOL>reverse_momentum = current_momentum<EOL>forward_state = current_state<EOL>forward_target_log_prob = current_target_log_prob<EOL>forward_grads_target_log_prob = current_grads_target_log_prob<EOL>forward_momentum = current_momentum<EOL>next_state = current_state<EOL>next_target_log_prob = current_target_log_prob<EOL>next_grads_target_log_prob = current_grads_target_log_prob<EOL>depth = <NUM_LIT:0><EOL>num_states = <NUM_LIT:1><EOL>continue_trajectory = True<EOL><DEDENT>while continue_trajectory:<EOL><INDENT>direction = tfp.math.random_rademacher([], seed=seed_stream())<EOL>if direction < <NUM_LIT:0>:<EOL><INDENT>[<EOL>reverse_state,<EOL>reverse_target_log_prob,<EOL>reverse_grads_target_log_prob,<EOL>reverse_momentum,<EOL>_,<EOL>_,<EOL>_,<EOL>_,<EOL>next_state_in_subtree,<EOL>next_target_log_prob_in_subtree,<EOL>next_grads_target_log_prob_in_subtree,<EOL>num_states_in_subtree,<EOL>continue_trajectory,<EOL>] = _build_tree(<EOL>value_and_gradients_fn=value_and_gradients_fn,<EOL>current_state=reverse_state,<EOL>current_target_log_prob=reverse_target_log_prob,<EOL>current_grads_target_log_prob=reverse_grads_target_log_prob,<EOL>current_momentum=reverse_momentum,<EOL>direction=direction,<EOL>depth=depth,<EOL>step_size=step_size,<EOL>log_slice_sample=log_slice_sample,<EOL>seed=seed_stream())<EOL><DEDENT>else:<EOL><INDENT>[<EOL>_,<EOL>_,<EOL>_,<EOL>_,<EOL>forward_state,<EOL>forward_target_log_prob,<EOL>forward_grads_target_log_prob,<EOL>forward_momentum,<EOL>next_state_in_subtree,<EOL>next_target_log_prob_in_subtree,<EOL>next_grads_target_log_prob_in_subtree,<EOL>num_states_in_subtree,<EOL>continue_trajectory,<EOL>] = _build_tree(<EOL>value_and_gradients_fn=value_and_gradients_fn,<EOL>current_state=forward_state,<EOL>current_target_log_prob=forward_target_log_prob,<EOL>current_grads_target_log_prob=forward_grads_target_log_prob,<EOL>current_momentum=forward_momentum,<EOL>direction=direction,<EOL>depth=depth,<EOL>step_size=step_size,<EOL>log_slice_sample=log_slice_sample,<EOL>seed=seed_stream())<EOL><DEDENT>if continue_trajectory:<EOL><INDENT>accept_state_in_subtree = _random_bernoulli(<EOL>[],<EOL>probs=tf.minimum(<NUM_LIT:1.>, num_states_in_subtree / num_states),<EOL>dtype=tf.bool,<EOL>seed=seed_stream())<EOL>if accept_state_in_subtree:<EOL><INDENT>next_state = next_state_in_subtree<EOL>next_target_log_prob = next_target_log_prob_in_subtree<EOL>next_grads_target_log_prob = next_grads_target_log_prob_in_subtree<EOL><DEDENT><DEDENT>has_no_u_turn = tf.logical_and(<EOL>_has_no_u_turn(forward_state, reverse_state, forward_momentum),<EOL>_has_no_u_turn(forward_state, reverse_state, reverse_momentum))<EOL>continue_trajectory = continue_trajectory and has_no_u_turn<EOL>num_states += num_states_in_subtree<EOL>depth += <NUM_LIT:1><EOL><DEDENT>return next_state, next_target_log_prob, next_grads_target_log_prob<EOL><DEDENT>
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.