```python
rng = create_key(0)
rng = jax.random.split(rng, jax.device_count())
```  
To take advantage of JAX's optimized speed on a TPU, pass `jit=True` to the pipeline to compile the JAX code into an efficient representation and to ensure the model runs in parallel across the 8 devices.  
<Tip warning={true}>  
You need to ensure all your inputs have the same shape in subsequent calls, otherwise JAX will need to recompile the code which is slower.  
</Tip>