def _init_random_number_generators(seed=None):
  """Initializes random generators for Python, NumPy, TensorFlow, and JAX."""
  # Seed Python random (None as seed is okay), then use it to seed the others.
  random.seed(seed)
  if seed is None:
    seed = random.randint(0, 2**31 - 1)
  logging.info('using seed %d', seed)
  np.random.seed(seed)
  tf.random.set_seed(seed)
  return jax_random.get_prng(seed)
Initializes random generators for Python, NumPy, TensorFlow, and JAX.
trax/supervised/training.py