def _choose_random_direction(current_state_parts, batch_rank, seed=None):
  """Chooses a random direction in the event space."""
  seed_gen = distributions.SeedStream(seed, salt='_choose_random_direction')
  # Chooses the random directions across each of the input components.
  rnd_direction_parts = [
      tf.random.normal(
          tf.shape(input=current_state_part), dtype=tf.float32, seed=seed_gen())
      for current_state_part in current_state_parts
  ]

  # Sum squares over all of the input components. Note this takes all
  # components into account.
  sum_squares = sum(
      tf.reduce_sum(
          input_tensor=rnd_direction**2.,
          axis=tf.range(batch_rank, tf.rank(rnd_direction)),
          keepdims=True) for rnd_direction in rnd_direction_parts)

  # Normalizes the random direction fragments.
  rnd_direction_parts = [rnd_direction / tf.sqrt(sum_squares)
                         for rnd_direction in rnd_direction_parts]

  return rnd_direction_parts
Chooses a random direction in the event space.