code: def get_random_number_generator_and_set_seed(seed=None):
    
    random.seed(seed)
    # While python random accepts None as seed and uses time/os seed then,
    # some other functions expect integers so we create one here.
    if seed is None:
      seed = random.randint(0, 2**31 - 1)
    tf.set_random_seed(seed)
    numpy.random.seed(seed)
    return jax_random.get_prng(seed)
docs: Get a JAX random number generator and set random seed everywhere.