ng = create_key(0)
rng = jax.random.split(rng, jax.device_count())
JAX code can be compiled to an efficient representation that runs very fast. However, we need to ensure that all inputs have the same shape in subsequent calls; otherwise, JAX will have to recompile the code, and we wouldn’t be able to take advantage of the optimized speed.
The Flax pipeline can compile the code for us if we pass jit = True as an argument. It will also ensure th