JAX

Computes gradients of arbitrary Python functions through reverse-mode automatic differentiation (grad) that decomposes composite functions into primitive operations with known derivatives. JAX traces function execution to build a computational graph, then applies the chain rule in reverse to compute gradients with respect to any input. Supports higher-order derivatives (hessian, jacobian) by composing grad with itself, enabling second-order optimization and sensitivity analysis without manual derivative specification.

Traces a Python function once to extract its computational graph, then compiles it to XLA (Accelerated Linear Algebra) bytecode for execution on CPUs, GPUs, or TPUs with near-native performance. The jit decorator intercepts all primitive operations during a symbolic trace, builds a static computation graph, and uses XLA's compiler to fuse operations, eliminate dead code, and generate device-specific machine code. Subsequent calls with the same input shapes execute the compiled code directly, bypassing Python interpretation.

Provides jax.lax.scan for efficient sequential computation (e.g., RNNs, sequential processing) that applies a function repeatedly over a sequence while maintaining state. scan traces the function once and generates XLA code that loops over the sequence, updating state at each step. Unlike explicit loops, scan enables compiler optimizations (kernel fusion, memory layout optimization) and works correctly within jit, vmap, and pmap, making it the preferred way to implement sequential algorithms in JAX.

Automatically places computations on available devices (CPU, GPU, TPU) without explicit device specification, enabling code portability across different hardware. JAX detects available devices at runtime and places arrays on the default device, with automatic data transfer between devices when needed. Users can control placement via jax.device_put and specify device constraints, but the default behavior is transparent device management that enables the same code to run on different hardware.

Enables JIT compilation of functions that work with variable input shapes by using abstract shapes during tracing, generating code that handles multiple concrete shapes without recompilation. When a function is JIT-compiled with abstract_eval, JAX traces it with symbolic shapes (e.g., (batch, 128) where batch is a variable dimension) and generates XLA code that works for any concrete batch size. This avoids recompilation when batch size changes, a common scenario in ML training and inference.

Provides utilities for saving and loading JAX arrays and PyTree structures via jax.experimental.io_callback and third-party libraries (e.g., flax.serialization, orbax), enabling model checkpointing and state persistence. JAX itself does not provide built-in serialization (by design, to keep the core library minimal), but the ecosystem provides robust solutions for saving model parameters, optimizer states, and training metadata. Checkpointing is essential for long-running training and enables resuming from interruptions.

Provides eager execution mode (jax.config.update('jax_disable_jit', True)) that disables JIT compilation and executes functions immediately, enabling step-by-step debugging and error inspection. In eager mode, Python control flow works normally, print statements execute, and errors occur at the line that causes them, making debugging much easier than in compiled mode. Eager execution is essential for development and debugging, though it sacrifices performance.

Automatically transforms a function written for a single input into a batched function via vmap (vectorized map), which adds a batch dimension and applies the function element-wise across that dimension. vmap traces the function once with a single example, then generates code that processes an entire batch in parallel using SIMD instructions or vectorized operations. Unlike explicit loops, vmap enables the compiler to fuse batch operations and optimize memory access patterns, often achieving near-peak hardware throughput.

Distributes a function across multiple devices (GPUs, TPUs) via pmap (parallel map), which replicates the function across devices and automatically handles data distribution and communication. pmap traces the function once per device, generates device-specific code, and inserts collective communication operations (all-reduce, all-gather) at specified points. Unlike vmap which vectorizes within a single device, pmap partitions computation across devices and manages synchronization, enabling near-linear scaling for data-parallel training.

Provides a NumPy-compatible API (jax.numpy) that implements standard array operations (matmul, sum, reshape, etc.) as JAX primitives that can be traced, differentiated, and compiled. jax.numpy wraps NumPy functions and maps them to XLA operations, ensuring that code written for NumPy can be executed on JAX with minimal changes. All operations are functional (no in-place mutations) and return immutable arrays, enabling safe composition with transformations like grad and jit.

Enables users to define custom operations (primitives) with specified gradient rules, XLA lowering rules, and batching rules, integrating them seamlessly into JAX's transformation system. Users register a primitive by defining its forward computation, then provide implementations for grad (reverse-mode derivative), jvp (forward-mode derivative), vmap (batching behavior), and xla_lowering (XLA compilation). Once registered, the custom operation can be used within any JAX transformation without modification, enabling extension of JAX with domain-specific operations.

Provides a functional random number API (jax.random) that generates deterministic, reproducible random numbers via explicit PRNG keys, avoiding global state and enabling safe parallelization. Instead of NumPy's stateful RNG, JAX uses a key-based system where each random operation consumes a key and returns a new key, ensuring reproducibility and enabling safe use in transformations. The API supports multiple RNG algorithms (threefry, philox) and enables splitting keys for parallel random number generation across devices.

Provides a PyTree system that treats nested structures (dicts, lists, tuples, custom classes) as first-class data structures in JAX transformations. PyTrees enable automatic handling of complex data structures in grad, jit, vmap, and pmap without manual flattening/unflattening. Users can register custom classes as PyTree nodes, and JAX automatically applies transformations to all leaves (arrays) while preserving structure, enabling clean code for handling model parameters, optimizer states, and other nested data.

Computes directional derivatives (Jacobian-vector products) via forward-mode AD (jvp), which traces a function while propagating derivative information forward through the computation graph. Unlike reverse-mode grad which is efficient for scalar outputs, forward-mode is efficient for functions with few inputs and many outputs. JAX implements jvp as a tracer that intercepts primitives and applies their forward-mode derivative rules, enabling efficient computation of Jacobians and sensitivity analysis.

Provides functional control flow primitives (jax.lax.cond, jax.lax.while_loop, jax.lax.fori_loop) that enable data-dependent branching and looping within JIT-compiled functions. Unlike Python's if/while which cannot depend on traced values, these primitives accept boolean conditions or loop bounds as inputs and generate XLA code that handles both branches/iterations. This enables dynamic computation (e.g., early stopping, adaptive algorithms) within compiled functions without breaking the compilation boundary.