JAX 101: expressing computations

JAX 101: expressing computations#

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:

  1. Arrays and jax.numpy — JAX’s array type and the jax.numpy API: what’s the same as NumPy, what’s different, and why.

  2. Transformations: grad and vmap — computing gradients with jax.grad() and vectorizing with jax.vmap(), plus the tracing model that underlies every JAX transformation.

  3. Pytrees — how JAX handles structured data, like nested dictionaries and lists of arrays.

  4. Pseudorandom numbers — pseudorandom numbers with explicit PRNG keys: pure functions of key values, with no hidden generator state.

  5. Stateful computations — stateful computations: threading state through pure functions, and in-place mutation with refs, JAX’s mutable array type.

Reference pages are there for when you need them, in any order:

  1. Errors — common JAX errors, explained: most arise from expressing something in a way that’s incompatible with tracing, so the tracing model from Transformations: grad and vmap is the key to fixing them.

  2. Convolutions — batched N-dimensional convolutions with jax.lax.conv_general_dilated(): dimension numbers, strides, dilation, and transposed convolutions.

  3. Default dtypes and the X64 flag — how JAX chooses default dtypes, and the X64 flag.

  4. Type promotion semantics — the full type promotion semantics: the result of combining every pair of types.

  5. Rank promotion warning — configuring jax.numpy to warn or error on implicit rank promotion.

The performance and scaling docs, JAX 201: performance and scaling, cover making these computations fast: compilation with jax.jit(), sharded arrays and parallelism, and profiling.