JAX

Computes gradients of arbitrary Python functions through reverse-mode (grad) and forward-mode automatic differentiation by tracing function execution and building a computational graph. JAX's grad function transforms a scalar-output function into one that returns both the output and gradient vector, supporting higher-order derivatives (hessian, jacobian) through function composition. Differentiates through control flow, loops, and nested function calls without explicit graph definition.

Traces Python functions to XLA intermediate representation and compiles them to optimized native code (CPU/GPU/TPU) via the XLA compiler, eliminating Python interpreter overhead. The jit decorator caches compiled kernels by input shape/dtype, reusing them across calls. Supports control flow through XLA's conditional and while_loop primitives, enabling Python-like syntax that compiles to efficient machine code.

JAX enforces functional programming by requiring explicit state management through carry parameters in loops (lax.scan, lax.while_loop) and transformations. State is passed as function arguments and returned as outputs, eliminating hidden state and making computations pure and composable. This enables deterministic execution, easy parallelization, and automatic differentiation through stateful computations.

JAX's transformations (grad, jit, vmap, pmap) are fully composable — you can nest them arbitrarily (e.g., jit(grad(vmap(f)))) and JAX automatically optimizes the composed computation. Each transformation is implemented as a function that takes a function and returns a transformed function, enabling functional composition. The composition order matters for performance but not correctness.

JAX integrates with Google's XLA (Accelerated Linear Algebra) compiler, which performs whole-program optimization including kernel fusion, memory layout optimization, and dead code elimination. jit compilation targets XLA, which generates optimized code for CPU/GPU/TPU. XLA's optimization is transparent — JAX automatically applies it to all jitted code, enabling significant performance improvements without manual optimization.

JAX enables pure functional neural network training where model parameters are explicit function arguments rather than stored in modules. Training loops are written as pure functions that take parameters and data, return updated parameters and loss. This approach enables automatic differentiation through entire training loops, easy parallelization across devices, and composability with all JAX transformations. Libraries like Flax and Optax provide higher-level abstractions on top of this functional foundation.

The vmap transformation automatically vectorizes functions across a specified axis, generating code that processes batches in parallel without explicit loop unrolling. vmap traces the function once with a single example, then generates vectorized code that applies the same computation to all batch elements. Composes with jit and grad — you can vmap a jitted function or differentiate a vmapped function, enabling batched gradient computation across distributed arrays.

The pmap transformation partitions arrays across multiple devices (GPUs, TPUs) and executes functions in parallel on each partition. pmap traces the function with a single device's slice of data, then replicates the computation across all devices with automatic communication (via collective ops like all_reduce) for cross-device operations. Integrates with jit for per-device compilation and with grad for distributed gradient computation.

jax.numpy provides a NumPy-compatible API for array operations (matmul, reshape, sum, etc.) that works with JAX's transformations. Operations are pure functions returning new arrays rather than mutating in-place, enabling composition with grad/jit/vmap. Supports broadcasting, indexing, and most NumPy functions, with some operations (like in-place updates) requiring functional alternatives (e.g., array.at[idx].set(value)).

The jax.custom_vjp (vector-jacobian product) and jax.custom_vmap decorators allow defining custom gradient rules for functions, enabling implementation of operations with non-standard differentiation (e.g., operations where the gradient differs from the forward pass, or where you want to optimize gradient computation). You define forward and backward passes separately, giving fine-grained control over gradient computation while maintaining composability with other JAX transformations.

JAX's random module uses explicit PRNG keys (jax.random.PRNGKey) instead of global state, enabling deterministic and reproducible randomness that composes with jit/vmap/pmap. Each random operation consumes a key and returns a new key, making randomness functional and parallelizable. Supports multiple PRNG algorithms (threefry, philox) and key splitting for generating independent random streams across devices.

jax.lax provides control flow primitives (cond, while_loop, fori_loop, scan) that compile to efficient XLA code while remaining differentiable. These replace Python's if/while statements inside jitted functions, enabling data-dependent control flow without breaking compilation or differentiation. scan is particularly powerful for sequential operations (RNNs, sequential models) with automatic gradient computation through time.

JAX's pytree system treats nested Python structures (dicts, lists, tuples, custom classes) as first-class objects, enabling transformations to work on entire data structures. grad/vmap/pmap automatically handle pytrees, applying transformations to all leaves (arrays) while preserving structure. Custom pytrees can be registered via jax.tree_util.register_pytree_node, enabling transformations on user-defined data structures.

JAX arrays are device-agnostic — operations automatically run on the default device (CPU/GPU/TPU) without explicit device placement. jax.device_put explicitly moves arrays to devices, and jax.devices() lists available hardware. Operations transparently use available accelerators, enabling code that works identically on CPU, GPU, or TPU without modification.