code: def _start_trajectory_batched(
      current_state, current_target_log_prob, seed_stream
  ):
      """Computations needed to start a trajectory."""
      with tf.name_scope("start_trajectory_batched"):
          batch_size = tf.shape(input=current_state[0])[0]
          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.log1p(
              -tf.random.uniform(
                  shape=[batch_size],
                  dtype=current_target_log_prob.dtype,
                  seed=seed_stream(),
              )
          )
          log_slice_sample += _log_joint(
              Point(
                  None, current_target_log_prob, None, current_momentum
              )
          )
  
          return current_momentum, log_slice_sample
  
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 3 for calling a global function with example usage: _start_trajectory_batched(current_state, current_target_log_prob, seed_stream) and returns: current_momentum,  log_slice_sample