code: def _update_trajectory_grad(
      previous_kernel_results,
      previous_state,
      proposed_state,
      proposed_velocity,
      trajectory_jitter,
      accept_prob,
      step_size,
      criterion_fn,
      max_leapfrog_steps,
      experimental_shard_axis_names=None,
      reduce_chain_axis_names=None,
  ):
      """Updates the trajectory length."""
  
      # Compute criterion grads.
      def leapfrog_action(dt):
          # This represents the effect on the criterion value as the state follows the
          # proposed velocity. This implicitly assumes an identity mass matrix.
          def adjust_state(x, v, shard_axes=None):
              broadcasted_dt = distribute_lib.pbroadcast(
                  bu.left_justified_expand_dims_like(dt, v), shard_axes
              )
              return x + broadcasted_dt * v
  
          adjusted_state = _map_structure_up_to_with_axes(
              proposed_state,
              adjust_state,
              proposed_state,
              proposed_velocity,
              experimental_shard_axis_names=experimental_shard_axis_names,
          )
          return criterion_fn(
              previous_state=previous_state,
              proposed_state=adjusted_state,
              accept_prob=accept_prob,
              # We add the step size here because we effectively do `floor(traj +
              # step_size) / step_size` when computing the number of leapfrog steps.
              trajectory_length=(
                  trajectory_jitter
                  * previous_kernel_results.max_trajectory_length
                  + step_size
                  + dt
              ),
          )
  
      criterion, trajectory_grad = gradient.value_and_gradient(
          leapfrog_action, tf.zeros_like(accept_prob)
      )
      trajectory_grad *= trajectory_jitter
  
      # Weight by acceptance probability.
      reduce_chain_axis_names = distribute_lib.canonicalize_named_axis(
          reduce_chain_axis_names
      )
      trajectory_grad = tf.where(
          accept_prob > 1e-4, trajectory_grad, 0.0
      )
      trajectory_grad = tf.where(
          tf.math.is_finite(trajectory_grad), trajectory_grad, 0.0
      )
      trajectory_grad = _reduce_sum_with_axes(
          trajectory_grad * accept_prob, None, reduce_chain_axis_names
      ) / _reduce_sum_with_axes(
          accept_prob + 1e-20, None, reduce_chain_axis_names
      )
  
      # Compute Adam/RMSProp step size.
      dtype = previous_kernel_results.adaptation_rate.dtype
      iteration_f = tf.cast(previous_kernel_results.step, dtype) + 1.0
      averaged_sq_grad_adaptation_rate = (
          previous_kernel_results.averaged_sq_grad_adaptation_rate
      )
      new_averaged_sq_grad = (
          (1 - averaged_sq_grad_adaptation_rate)
          * previous_kernel_results.averaged_sq_grad
          + averaged_sq_grad_adaptation_rate * trajectory_grad**2
      )
      adjusted_averaged_sq_grad = new_averaged_sq_grad / (
          1.0 - (1 - averaged_sq_grad_adaptation_rate) ** iteration_f
      )
      trajectory_step_size = (
          previous_kernel_results.adaptation_rate
          / tf.sqrt(adjusted_averaged_sq_grad + 1e-20)
      )
  
      # Apply the gradient. Clip absolute value to ~log(2)/2.
      log_update = tf.clip_by_value(
          trajectory_step_size * trajectory_grad, -0.35, 0.35
      )
      new_max_trajectory_length = (
          previous_kernel_results.max_trajectory_length
          * tf.exp(log_update)
      )
  
      # Iterate averaging.
      average_weight = iteration_f ** (-0.5)
      new_averaged_max_trajectory_length = tf.exp(
          average_weight * tf.math.log(new_max_trajectory_length)
          + (1 - average_weight)
          * tf.math.log(
              1e-10
              + previous_kernel_results.averaged_max_trajectory_length
          )
      )
  
      # Clip the maximum trajectory length.
      new_max_trajectory_length = _clip_max_trajectory_length(
          new_max_trajectory_length,
          step_size,
          previous_kernel_results.adaptation_rate,
          max_leapfrog_steps,
      )
  
      return previous_kernel_results._replace(
          criterion=criterion,
          max_trajectory_length=new_max_trajectory_length,
          averaged_sq_grad=new_averaged_sq_grad,
          averaged_max_trajectory_length=new_averaged_max_trajectory_length,
      )
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  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 unnest
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 15 for calling a global function with example usage: _update_trajectory_grad(previous_kernel_results, previous_state, proposed_state, proposed_velocity, trajectory_jitter, accept_prob, step_size, criterion_fn, max_leapfrog_steps, experimental_shard_axis_names, reduce_chain_axis_names) and returns: previous_kernel_results, criterion_fn, unknown