[BUG] rng runs with each microbatch

[ZYHowell]

Each microbatch runs the traced Jaxpr once. However, some parts of the Jaxpr is not related to microbatch. This results in redundant computations and wrong behavior. For example:

def f(x, key):
    random_num = jax.random.uniform(key, (x.shape[-1], 2))
    return x @ random_num

Here the random_num is not related to batch args(x), but each microbatch generates its own random_num. If there are two microbatches x = [x1, x2], the expected result is [x1@rn, x2@rn], but we actually get [x1@rn1, x2@rn2]. We should separate the rng and reuse the result in all microbatches