JAX 101: expressing computations — JAX documentation
JAX is a Python library for high-performance numerical computing and machine learning. Its interface, centered on jax.numpy, will look familiar if you’ve used NumPy. What sets JAX apart is what it can do with the functions you write: transform them, to compute gradients or to vectorize over batches; and compile them, to run fast on CPU, GPU, and TPU, at any scale. These pages cover the first half of that story: how to express computations in JAX. They’re meant to be read in order: Arrays and jax.numpy — JAX’s array type and the jax.numpy API: what’s the same as NumPy, what’s different, and why. Transformations: grad and vmap — computing gradients with jax.grad() and vectorizing with jax.vmap(), plus the tracing model that underlies every JAX transformation. Pytrees — how JAX handles structured data, like nested dictionaries and lists of arrays. Pseudorandom numbers — pseudorandom numbers with explicit PRNG keys: pure functions of key values, with no hidden generator state. Stateful com