code: def _choose_random_direction(
      current_state_parts,
      batch_rank,
      seed=None,
      experimental_shard_axis_names=None,
  ):
      """Chooses a random direction in the event space."""
      seeds = list(
          samplers.split_seed(seed, n=len(current_state_parts))
      )
      seeds = distribute_lib.fold_in_axis_index(
          seeds, experimental_shard_axis_names
      )
  
      # Sample random directions across each of the input components.
      def _sample_direction_part(state_part, part_seed):
          state_part_shape = ps.shape(state_part)
          batch_shape = state_part_shape[:batch_rank]
          dimension = ps.reduce_prod(state_part_shape[batch_rank:])
          return ps.reshape(
              random_ops.spherical_uniform(
                  shape=batch_shape,
                  dimension=dimension,
                  dtype=state_part.dtype,
                  seed=part_seed,
              ),
              state_part_shape,
          )
  
      return [
          _sample_direction_part(state_part, seed)
          for state_part, seed in zip(current_state_parts, seeds)
      ]
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import slice_sampler_utils
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow_probability.python.random import random_ops
label_desc: Use image node_id 1 for calling a global function with example usage: _choose_random_direction(current_state_parts, batch_rank, seed, experimental_shard_axis_names) and returns: unknown, ps