t(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. 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. The first inference run takes more time because it needs to compile the code, b