1 Introduction
JAX is an open-source Python library developed by Google Research for high‑performance numerical computing, with a primary focus on machine learning research. It provides a NumPy‑compatible application programming interface (API) combined with automatic differentiation (autograd), just‑in‑time compilation via XLA (Accelerated Linear Algebra), and a set of composable function transformations (vmap, pmap, pjit) that enable efficient execution on a variety of hardware platforms, including CPUs, GPUs, and TPUs. JAX is designed to allow researchers to write clean, NumPy‑style code while achieving hardware‑accelerated performance and straightforward gradient computation, thereby accelerating the iterative cycle of experimentation and discovery.
1.1 Background and motivation
The motivation for JAX arose from the need for a flexible, high‑performance framework that could combine the ease of use of NumPy with the automatic differentiation and hardware acceleration required for modern machine learning. Existing solutions often forced a trade‑off between expressiveness and performance: TensorFlow’s static graph approach limited dynamic control flow, while PyTorch’s eager execution sometimes incurred overhead. JAX addresses these challenges by leveraging XLA compilation and a functional programming paradigm that transforms user functions into efficient computation graphs without sacrificing Pythonic readability.
1.2 Key design principles
JAX is built on three core design principles that underlie its architecture and usage.
1.2.1 Composability of transformations
Function transformations in JAX—such as jit, grad, vmap, and pmap—can be arbitrarily nested and combined. A user can, for example, compute the gradient of a vectorized, JIT‑compiled function with a single line of code. This composability encourages building complex pipelines from simple, pure functions.
1.2.2 Functional programming paradigm
JAX encourages a functional style: functions must be pure (no side effects), and data is immutable. This purity enables transformations to safely trace and modify function behavior, ensures deterministic results, and simplifies reasoning about parallelism. The library provides controlled state management via explicit PRNG sequences and state‑carrying constructs (e.g., jax.lax.scan).
1.2.3 Asynchronous dispatch
JAX operations are dispatched asynchronously to the accelerator. When a JIT‑compiled function is called, JAX returns a DeviceArray (a future) immediately, while the actual computation and data transfer occur on the device in a separate stream. This allows the host to continue executing Python code in parallel with device computation, improving overall throughput in interactive and batch settings.
2 Core Components
JAX is organized around a set of core components that replace and extend NumPy’s functionality with hardware acceleration and automatic differentiation.
2.1 jax.numpy (jnp)
The jax.numpy module (commonly imported as jnp) provides a NumPy‑compatible interface. Almost all NumPy functions have a corresponding JAX version that operates on JAX arrays and works seamlessly with transformations.
2.1.1 Array creation and manipulation
jnp.array, jnp.zeros, jnp.ones, jnp.arange, and other creation functions produce JAX arrays (jax.Array). Manipulation routines (jnp.reshape, jnp.concatenate, jnp.transpose, etc.) behave identically to NumPy but return arrays that can be traced and compiled. Array indexing, slicing, and broadcasting follow NumPy conventions.
2.1.2 Random number generation (jax.random)
Unlike NumPy, JAX uses an explicit, functional random number generator (RNG) to maintain purity. The jax.random module requires an explicit PRNG key (a jax.Array of shape (2,)). Functions such as jax.random.normal, jax.random.uniform, and jax.random.split consume the key and return both the random values and a new key. This design avoids global state and ensures reproducibility across transformations.
2.1.3 Control flow and lax primitives
JAX provides low‑level linear algebra primitives in jax.lax (e.g., lax.add, lax.conv, lax.dot_general). While jnp functions often map to these primitives, for advanced control flow (conditional loops, scans) users employ lax.cond, lax.while_loop, or lax.scan to maintain traceability and avoid Python‑level dynamic control flow that would break JIT compilation.
2.2 Automatic Differentiation (autograd)
JAX’s autograd system computes derivatives of arbitrary functions with respect to their inputs.
2.2.1 grad, jacfwd, jacrev
The primary transformation jax.grad returns a function that computes the gradient of a scalar‑valued function. For vector‑valued or matrix‑valued outputs, jax.jacfwd and jax.jacrev compute Jacobians using forward‑ and reverse‑mode differentiation, respectively. jax.hessian provides second‑order derivatives by composing these.
2.2.2 Custom gradients (custom_vjp, custom_jvp)
For functions with non‑differentiable operations or where performance can be improved, JAX allows users to define custom forward (JVP) and reverse (VJP) rules via jax.custom_jvp and jax.custom_vjp. These decorators enable overriding the default differentiation behavior, facilitating integration of external algorithms or manual gradient scaling.
2.2.3 Higher-order derivatives
Because grad itself is a transformation, higher‑order derivatives are obtained by repeated application: jax.grad(jax.grad(f)) computes the second derivative. This composability extends to arbitrary orders, limited only by memory and numerical stability.
2.3 Just-in-Time Compilation (jit)
The jax.jit transformation compiles a Python function into an optimized XLA computation graph, which is then executed on the accelerator.
2.3.1 XLA compilation pipeline
When a function decorated with @jax.jit is called, JAX traces its operations (using abstract Trace objects) to build a stable‑hash Computation (HLO) graph. XLA further optimizes this graph (fusing operations, rematerializing buffers) before generating device‑specific code. Subsequent calls with the same input shapes and types reuse the compiled binary.
2.3.2 Static arguments and tracing
By default, jit traces all arguments. Arguments that do not change shape or type between calls can be marked as static using the static_argnums or static_argnames parameters. Static arguments are converted to Python integers at compile time, allowing the traced function to specialize on their values. This is essential for loops or array shapes that depend on constants.
2.3.3 Practical performance considerations
JIT compilation introduces a one‑time overhead on the first call (compilation time). For subsequent calls, performance is typically near native. Users should avoid recompilation by ensuring consistent input shapes and using static arguments where possible. Overly large functions may lead to long compilation times, and some dynamic operations (e.g., Python‑level conditional on array values) can force re‑tracing.
2.4 Vectorization and Batching (vmap)
jax.vmap automatically vectorizes a function by mapping it over leading batch dimensions of its inputs, effectively executing the function in a single batched call.
2.4.1 Automatic batching semantics
vmap takes an axis of size batch_size from each input and passes a slice to the inner function. It handles broadcasting and produces outputs with the batch dimension. The transformation can be applied to any function, including ones that already use jit, enabling efficient inner‑loop vectorization.
2.4.2 Combining vmap with jit and grad
Because transformations compose, jax.vmap(jax.jit(f)) compiles a vectorized version of f, and jax.grad(jax.vmap(f)) computes per‑sample gradients of a batched function. This composition is key to writing concise, efficient training loops.
2.5 Parallelization (pmap, pjit)
JAX provides transformations for distributing computations across multiple devices.
2.5.1 Single-program multiple-data (SPMD) with pmap
jax.pmap executes a function in parallel across multiple devices (e.g., multiple GPUs or TPU cores). It replicates the function on each device and passes a shard of the input data to each replicant. Communication between replicants uses collective operations. pmap automatically handles gradient averaging during distributed training.
2.5.2 shard_map and pjit for model parallelism
For more sophisticated parallelization strategies (e.g., model parallelism, pipeline parallelism), JAX offers jax.shard_map and jax.experimental.pjit. These allow users to specify how arrays are partitioned across devices (using jax.sharding.Sharding objects) and to define axis‑wise reductions and communications. pjit compiles the function specifically for the given mesh of devices and memory layouts.
2.5.3 Collective communication (psum, all_gather)
JAX provides collective operations like jax.lax.psum (parallel sum), jax.lax.pmax, and jax.lax.all_gather for aggregating data across devices. These are used inside pmap or pjit‑compiled functions to synchronize gradients, parameters, or statistics.
3 Ecosystem and Integrations
JAX’s ecosystem includes a variety of third‑party libraries that build on its core transformations for specialized domains.
3.1 Neural network libraries
Several deep‑learning libraries provide high‑level neural network layers, training utilities, and model management for JAX.
3.1.1 Flax
Flax (by Google Research) is a flexible neural‑network library for JAX. It offers a Module abstraction similar to PyTorch’s nn.Module, with flax.linen providing common layers, optimizers via flax.training, and utilities for serialization and checkpointing. Flax is widely used in research and production within the Google ecosystem.
3.1.2 Haiku
Haiku (by DeepMind) provides an object‑oriented interface for defining neural networks. It uses a function‑transform pattern: users define a forward function that creates parameters using hk.nets and hk.Conv2D etc. The library handles parameter management via hk.transform, which separates network definition from parameter storage.
3.1.3 Equinox
Equinox is a lightweight library that treats model parameters as PyTrees (nested dictionaries of arrays). It uses standard Python classes and __call__ methods, making it intuitive for users familiar with PyTorch. Equinox integrates with JAX transformations and provides type‑safe parameter handling.
3.2 Optimization and training
3.2.1 Optax
Optax is a gradient‑processing and optimization library for JAX. It provides a composable collection of optimizers (SGD, Adam, AdamW, LAMB, etc.), learning‑rate schedules, gradient transformations (clipping, normalization, weight decay), and utilities for combining them. Optax’s design uses pure functions that accept parameters and gradients and return new parameters, making it naturally compatible with JAX’s functional style.
3.2.2 Other optimizers and schedulers
Other optimization libraries exist, such as jaxopt (convex and non‑convex solvers) and custom implementations within Flax or Haiku. Most of these follow Optax’s update‑function pattern, and many support second‑order methods (e.g., L‑BFGS) via JAX’s automatic differentiation.
3.3 Data loading and pipelines
3.3.1 TensorFlow Datasets integration
JAX works seamlessly with TensorFlow Datasets (TFDS). Because JAX arrays can be created from TensorFlow tensors (via jax.device_put), TFDS can be used as a data source. The dm‑tree library (DeepMind) further facilitates converting TF data to JAX PyTrees. Many JAX examples use tf.data pipelines with jax.numpy conversion.
3.3.2 Custom data pipelines with jax.experimental
For advanced data loading, JAX’s experimental modules (jax.experimental.map and jax.experimental.io_callback) allow users to plug in custom data sources using pure functions. The jax.experimental.shard_map can also be used to distribute data loading across devices.
3.4 Probabilistic programming
3.4.1 NumPyro
NumPyro is a probabilistic programming library that uses JAX’s automatic differentiation and JIT compilation for Bayesian inference. It supports Hamiltonian Monte Carlo (HMC), NUTS, variational inference (VI), and other MCMC methods. NumPyro leverages JAX for efficient gradient‑based sampling on accelerators.
3.4.2 BlackJAX
BlackJAX is a library of MCMC sampling algorithms implemented in JAX. It provides building blocks for custom samplers, including HMC, NUTS, SGLD, and more. BlackJAX is designed to be compatible with any model that can be expressed as a JAX‑differentiable log‑probability function.
3.5 Reinforcement learning
3.5.1 RLlib and JAX-backed environments
Ray RLlib supports JAX as a backend for training. Several environment libraries (e.g., dm_env, Gym‑style) can be adapted to JAX by wrapping observation and action spaces as JAX arrays. JAX’s vmap and pmap are used to parallelize environment rollouts.
3.5.2 PureJAX RL implementations
Many custom reinforcement learning agents are implemented purely in JAX. Typical examples include Deep Q‑Networks (DQN), Proximal Policy Optimization (PPO), and Soft Actor‑Critic (SAC). These implementations leverage jit for the training step, vmap for batched environment steps, and pmap for multi‑agent or multi‑worker setups.
4 Use Cases and Examples
JAX is applied across a wide range of scientific and machine learning domains.
4.1 Scientific computing
4.1.1 Differentiable physics simulations
JAX’s automatic differentiation enables gradient‑based optimization of physical simulations (e.g., fluid dynamics, rigid‑body mechanics). Libraries like jax‑md (molecular dynamics) and optax are used to either fit simulation parameters to data or to compute control policies. The vmap transformation accelerates batch simulations.
4.1.2 Bayesian inference with MCMC
Probabilistic programming libraries (NumPyro, BlackJAX) use JAX to sample from posterior distributions. The JIT compilation speeds up the repeated likelihood evaluations required by MCMC, while grad provides the gradients needed for HMC/NUTS. This approach is used in fields like epidemiology, astronomy, and econometrics.
4.2 Machine learning research
4.2.1 Training small-scale models (MLPs, CNNs)
Typical small‑scale models (multilayer perceptrons, convolutional networks) are implemented in JAX with Flax or Equinox. Training loops use jit for the forward‑backward pass, grad for gradient computation, and Optax for updates. Data is batched using vmap and loaded via TFDS.
4.2.2 Transformer language models
JAX is used for training and serving transformer models (e.g., BERT, GPT‑style). Libraries like Flax and Haiku provide Transformer layers. pmap and pjit are used for distributed training across many TPUs. The shard_map transformation enables efficient tensor parallelism for large models.
4.2.3 Generative models (diffusion, GANs)
Generative models that require sampling and derivatives (score‑based diffusion, generative adversarial networks) benefit from JAX’s fast compilation and vectorization. For example, the score_sde library uses JAX for neural SDE solvers, and GAN training loops exploit vmap for batch‑wise discriminators.
4.3 Large-scale distributed training
4.3.1 Data parallelism
Data parallelism distributes a model’s training batch across multiple devices. Using pmap, each device holds a full copy of the model parameters and processes a sub‑batch. Gradients are averaged with jax.lax.psum. This is the most common distributed training pattern in JAX.
4.3.2 Model parallelism with pjit
Model parallelism partitions a model’s layers or parameters across devices. pjit allows users to specify sharding constraints via jax.sharding.Mesh and jax.sharding.PartitionSpec. This approach is used for very large models that exceed the memory of a single device.
4.3.3 Mixed-precision training
JAX natively supports bfloat16 and float16 via array dtypes. Training loops can cast parameters or activations to half‑precision while keeping optimizer state in float32. With pmap and a custom loss‑scaling function (or using Optax’s scale_by_loss), mixed‑precision training achieves significant speedups on TPUs and GPUs.
5 Comparison with Other Frameworks
JAX occupies a unique position in the numerical computing landscape, distinct from both TensorFlow and PyTorch.
5.1 JAX vs TensorFlow
5.1.1 Eager execution vs functional transformations
TensorFlow (TF) introduced eager execution as default in TF 2.x, but its functional style (via tf.function) still modifies Python semantics. JAX is purely functional from the start: every function transformation is explicit and composable. This often leads to cleaner, more predictable code for research.
5.1.2 AutoDiff and gradient computation
Both JAX and TF support automatic differentiation, but JAX’s grad is a simple identity on functions, while TF uses tf.GradientTape which records operations in a context manager. JAX’s approach is more flexible for higher‑order gradients and custom VJP rules.
5.1.3 Production readiness and deployment
TensorFlow has a mature serving ecosystem (TF Serving, TF Lite, TensorRT). JAX’s production tooling is less mature, though initiatives like JAX2TF (converting JAX models to TF graphs) and XLA‑based serving are emerging. For research‑focused workflows, JAX often has fewer deployment overheads.
5.2 JAX vs PyTorch
5.2.1 Dynamic vs static computation graphs
PyTorch uses a dynamic graph (define‑by‑run) by default, which is intuitive for programming with control flow. JAX compiles static graphs through jit, which can limit dynamic branching but yields better performance. JAX’s lax.cond and lax.while_loop provide a functional way to handle dynamic shapes efficiently.
5.2.2 JIT compilation and performance
JAX’s JIT compilation via XLA often yields faster execution than PyTorch’s eager mode, especially on TPUs. PyTorch’s torch.jit and torch.compile (TorchDynamo) have been catching up, but JAX’s transformations are more tightly integrated with the compiler.
5.2.3 Ecosystem maturity and community
PyTorch has a larger community, more pre‑trained models, and extensive third‑party libraries. JAX’s ecosystem is smaller but growing rapidly, with strong support from Google and DeepMind for research‑focused libraries (Flax, Haiku, Optax). For cutting‑edge ML research, JAX is often preferred for its flexibility and performance.
5.3 JAX vs NumPy with Numba/Cupy
5.3.1 Automatic differentiation advantage
NumPy alone lacks automatic differentiation; users must manually derive gradients. Tools like Numba or CuPy can accelerate NumPy code but provide no mechanism for gradient computation. JAX integrates both acceleration and autodiff, enabling end‑to‑end differentiable programs without additional libraries.
5.3.2 Portability across hardware backends
NumPy runs only on CPUs. CuPy targets NVIDIA GPUs, and Numba can target CPUs and GPUs via LLVM, but each requires different code paths. JAX’s XLA backend abstracts over CPU, GPU, and TPU, allowing the same user code to execute on any supported device with no modifications.
6 Advanced Topics
JAX offers several advanced mechanisms for customizing and extending its capabilities.
6.1 Custom primitives (lax.custom_lin)
jax.lax.custom_linear_solve and jax.lax.custom_root allow users to define custom linear‑solve and root‑finding operations with user‑specified Jacobian‑vector products. This is valuable for implementing implicit layers (e.g., Neural ODEs, DEQs) while maintaining efficient gradient propagation.
6.2 Metaprogramming and abstract evaluation
JAX provides jax.make_jaxpr to obtain a Jaxpr (JAX expression) representation of a function’s computation graph. Users can inspect, modify, or re‑compile Jaxprs for meta‑programming. Abstract evaluation (via jax.core.abstract_eval) allows reasoning about shapes and dtypes without executing concrete values.
6.3 Serialization and checkpointing
Model parameters and optimizer states are Python trees of arrays. They can be serialized using standard libraries like pickle or safetensors. For large‑scale training, jax.experimental.checkpoint and orbax.checkpoint provide utilities for saving and restoring training state across distributed runs.
6.4 Debugging and profiling
6.4.1 jax.debug and breakpoint()
The jax.debug module offers print and breakpoint functions that can be used inside JIT‑compiled functions. They are implemented as special primitives that trigger side effects during tracing, allowing users to inspect intermediate values without disabling compilation.
6.4.2 TensorBoard integration
JAX can write logs compatible with TensorBoard via the tensorboard package. Users often combine jax.experimental.host_callback with callbacks that write scalar and histogram summaries during training.
6.4.3 XLA dump and visualisation
XLA compilation can produce HLO graphs and op‑level statistics. By setting XLA_FLAGS to dump HLO, users can visualize the computation graph using tools like tensorboard or XLA_HLO_GRAPH. This aids in debugging performance bottlenecks and verifying compilation.
7 Limitations and Future Directions
Despite its strengths, JAX has several known limitations and ongoing development efforts.
7.1 Known limitations
7.1.1 Functional purity constraints
JAX’s requirement that functions be pure (no side effects) can be limiting for tasks that rely on mutable state, such as certain reinforcement‑learning algorithms or interactive simulation loops. While workarounds exist (e.g., lax.scan for accumulation), they may feel unnatural to users accustomed to imperative programming.
7.1.2 Dynamic control flow performance
Python‑level control flow (if‑statements, loops) inside JIT‑compiled functions must be traced, which can either force recompilation (if branching on array values) or require using lax.cond / lax.while_loop. These functional primitives can be less intuitive and may incur overhead compared to native control flow.
7.1.3 Device memory management
JAX does not automatically garbage‑collect device memory as aggressively as some frameworks. Users must be careful to release references to large arrays (e.g., by rebinding variables) or use jax.lax.stop_gradient to avoid retaining unnecessary computation graphs. Memory fragmentation on accelerators can also be an issue.
7.2 Ongoing developments
7.2.1 JAX2TF and multi-framework interop
Work is underway to improve interoperability between JAX and TensorFlow. The JAX2TF tool converts JAX‑compiled functions to TensorFlow graphs, allowing deployment in TF Serving or integration with existing TF pipelines. Similarly, experiments with JAX‑Torch bridging are being explored.
7.2.2 Improved error messages
Early JAX error messages could be opaque, especially when tracing large functions. Recent releases have significantly improved error reporting, including more informative stack traces and warnings about unexpected shape changes or re‑tracing. The community continues to refine these diagnostics.
7.2.3 Community‑driven enhancements
The JAX ecosystem benefits from active contributions. Ongoing efforts include expanding the library of primitives (e.g., custom callbacks, sparse array support), improving documentation, and refining the developer experience for large‑scale distributed training. The rise of projects like jax2torch and eqx exemplifies the community’s drive to make JAX more accessible and interoperable.