Need one numerical function to run faster, differentiate automatically, process batches, or execute across accelerators? A composable function transformation takes a function and produces a transformed version. JAX applies this functional style to Python and NumPy code, so transformations can be combined instead of implemented separately. section 1.2.2 section 1.4
JAX’s core transformations cover four complementary jobs:
- `grad()` applies automatic differentiation, producing gradients without requiring derivatives to be written by hand. section 1.2.2
- `jit()` compiles a function with XLA to produce efficient code for supported CPUs, GPUs, and TPUs. section 1.2.2
- `vmap()` auto-vectorizes a function, converting single-item computation into batch computation. section 1.2.2
- `pmap()` compiles and replicates a function for parallel execution across XLA devices using a single-program, multiple-data pattern. section 1.2.2
Compose transformations
A representative composition is `jit(vmap(grad(some_function)))`: differentiate a function, apply it across a batch, then compile the resulting computation. This style works best when functions are pure—defined only by their inputs, with no hidden state or side effects. Parameters, gradients, optimizer state, and random-number state therefore remain explicit arguments. Ignoring purity can produce surprising behavior because JIT compiles and caches computations. section 2.9 section 1.1.2
Figure 1.3 JAX (right) is much more than NumPy (left). It provides a set of composable function transformations for Python + NumPy code for compilation, vectorization, parallelization, and automatic differentiation.
Use the comparison figure to see how JAX extends a NumPy-style API with compilation, vectorization, parallelization, and automatic differentiation. section 1.2.2
