Tutorial: JAX 101# This is a tutorial developed by engineers and researchers at DeepMind. Tutorials JAX As Accelerated NumPy Just In Time Compilation with JAX Automatic Vectorization in JAX Advanced Automatic Differentiation in JAX Pseudo Random Numbers in JAX Working with Pytrees Parallel Evaluation in JAX Stateful Computations in JAX