# JAX

> Source: https://aiwiki.ai/wiki/jax
> Updated: 2026-07-29
> Fact-checked: 2026-07-29
> Categories: Deep Learning, Machine Learning
> License: CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/) - attribute to "AI Wiki (aiwiki.ai)"
> Cite as: AI Wiki. "JAX." aiwiki.ai, 29 Jul 2026. https://aiwiki.ai/wiki/jax
> From AI Wiki (https://aiwiki.ai), the free encyclopedia of artificial intelligence. Reuse freely with attribution.

JAX is an open-source [Python](https://aiwiki.ai/wiki/python) library for accelerator-oriented array computation and program transformation. Its core interface combines a [NumPy](https://aiwiki.ai/wiki/numpy)-inspired array API with transformations for compilation, [automatic differentiation](https://aiwiki.ai/wiki/automatic_differentiation), vectorization, and parallel execution. JAX programs can target CPUs, GPUs, and Google TPUs through the OpenXLA compiler and PJRT runtime stack.[1][5][12]

JAX is a numerical computing library rather than a complete [deep learning](https://aiwiki.ai/wiki/deep_learning) framework. The core package supplies arrays, mathematical primitives, transformations, sharding, and extension interfaces. Model modules, optimizers, data loading, and checkpoint management normally come from separate libraries such as Flax, Optax, Grain, and Orbax. Those projects may work closely with JAX contributors, but the JAX project states that ecosystem libraries are developed independently and are not governed by the JAX core team.[12][27][28][29][30]

The project is released under the Apache License 2.0.[2] At the research cutoff of July 28, 2026, the latest released version recorded in the official change log was JAX 0.11.0, released on July 16, 2026.[14] JAX remained in the "zero version" phase of its effort-based versioning policy, so changes to the second version component can include major breaking changes and changes to the third component can include smaller breaking changes.[15]

## Scope and purpose

JAX is designed around array programs whose numerically significant work can be represented by JAX operations. A user writes ordinary Python functions with `jax.numpy` and related APIs, then applies transformations such as `jax.jit`, `jax.grad`, or `jax.vmap`. During transformation, JAX traces supported operations and constructs a functional intermediate representation. The resulting program can be differentiated, batched, compiled, or partitioned without requiring a separate model-definition language.[1][3][5][6]

This design serves more than neural-network training. JAX can be used for optimization, probabilistic computation, differentiable simulation, numerical linear algebra, and other scientific workloads that benefit from derivatives or accelerator execution. The library does not make every Python program compilable. Data-dependent Python control flow, side effects, dynamic output shapes, foreign code, and unsupported operations require restructuring, a JAX control-flow primitive, a callback, a custom operation, or execution outside a transformed region.[5][8][22][36]

JAX also does not promise a universal performance advantage over NumPy, [PyTorch](https://aiwiki.ai/wiki/pytorch), [TensorFlow](https://aiwiki.ai/wiki/tensorflow), or hand-written kernels. Performance depends on the operation mix, shapes and dtypes, target accelerator, compiler version, transfer costs, compilation cost, and whether a benchmark waits for asynchronous execution to finish. JAX's own documentation requires synchronization with `block_until_ready()` or a host transfer when measuring device execution, and it separates first-call compilation latency from subsequent execution.[5][11]

## Origins and development

### Autograd

An important predecessor was Autograd, introduced at the 2015 ICML AutoML Workshop by Dougal Maclaurin, David Duvenaud, and Ryan P. Adams. Autograd differentiated ordinary Python and NumPy code, including programs with loops, branches, closures, classes, and nested derivative calculations. It supported higher-order differentiation without forcing users to express a model in a separate symbolic language.[4]

JAX retained the idea of transforming Python numerical functions while adding compilation and accelerator execution. The JAX documentation still uses the comparison to explain what the newer system adds: JAX can combine differentiation with compilation and can use `vmap` to batch derivative calculations that Autograd had to evaluate one vector at a time.[7]

### The 2018 system

Roy Frostig, Matthew James Johnson, and Chris Leary described an early JAX system at SysML 2018. Their paper presented JAX as a tracing just-in-time compiler for pure Python and NumPy [machine learning](https://aiwiki.ai/wiki/machine_learning) programs. It used [XLA](https://aiwiki.ai/wiki/xla) to compile selected numerical subroutines for CPUs, GPUs, and TPUs while leaving ordinary Python to coordinate the compiled work.[3]

The paper's use of "high-level tracing" had two meanings. The tracing machinery was written at user level in Python rather than built into the Python interpreter, and its primitives were array operations such as matrix multiplication, convolution, reduction, and multidimensional indexing rather than virtual-machine instructions. The system was compatible with Autograd and supported forward-mode and reverse-mode differentiation to arbitrary order.[3]

The current project has a broader scope than the 2018 prototype. It has a unified `jax.Array` type, multi-device arrays, explicit and automatic sharding, export through StableHLO, custom-kernel support through Pallas, and several extension paths. Development occurs publicly through the `jax-ml/jax` repository, issues, pull requests, discussions, and JAX Enhancement Proposals. The JAX core team leads the project, with contributions from Google DeepMind, other Alphabet teams, hardware companies, and independent contributors.[1][12]

### Selected release milestones

| Date | Documented change |
|---|---|
| 2015 | Autograd paper described differentiation of ordinary Python and NumPy programs.[4] |
| 2018 | SysML paper described JAX's high-level tracing compiler and its use of XLA.[3] |
| December 2020 | DeepMind reported broad internal use of JAX and described Haiku, Optax, RLax, Chex, and Jraph as separate libraries built around it.[24] |
| April 2022 | The PaLM paper reported training a 540-billion-parameter model with a JAX-based system on 6,144 TPU v4 chips.[25] |
| October 15, 2025 | JAX 0.8.0 changed the default `pmap` implementation to one based on `jit` and `shard_map` and placed `pmap` in maintenance mode.[14][19] |
| April 16, 2026 | JAX 0.10.0 removed the legacy C++ `pmap` infrastructure.[14] |
| July 16, 2026 | JAX 0.11.0 was released.[14] |

The milestone table records only events directly supported by project records or primary papers. Repository popularity, package-download counts, unnamed company deployments, and claims that a particular unreleased model used JAX are omitted because those figures and attributions change quickly or lack a stable primary source.

## Core programming model

### Arrays and the NumPy-inspired API

`jax.numpy`, conventionally imported as `jnp`, mirrors much of the NumPy interface. It provides array creation, indexing, broadcasting, reductions, linear algebra, Fourier transforms, and mathematical functions using JAX arrays. The similarity lowers the cost of moving numerical expressions into JAX, but it is not a guarantee that arbitrary NumPy code can be copied unchanged.[1][5]

A `jax.Array` can reside on one device, be replicated, or be sharded across several devices and processes. Its `sharding` attribute describes its placement. JAX operations generally follow the placement of their inputs, subject to explicit sharding rules and compiler partitioning. A NumPy array, by contrast, is ordinarily a host-memory object and does not encode distributed device placement.[5][16]

JAX arrays are immutable at the Python level. An expression such as `x[i] = value` is rejected. Indexed updates use the functional form `x.at[i].set(value)` or related methods, which return an updated value. Inside compiled code, the compiler may implement such an update in place when it can prove that the original buffer is not needed, so functional semantics do not necessarily require an extra physical copy.[8]

JAX's dtype defaults also differ from standard NumPy defaults. The `jax_enable_x64` option is false by default, which normally produces 32-bit integers and floating-point values and prevents creation of 64-bit arrays. Users who need 64-bit numerical work can enable the option globally before computation. This choice is significant for scientific calculations because code that silently runs in 32-bit mode may have different accuracy from an equivalent NumPy program.[35]

### API layers

The public numerical stack has several layers:

| Layer | Role |
|---|---|
| `jax.numpy` | Familiar high-level array API modeled on NumPy |
| `jax.scipy` | Selected SciPy-compatible numerical routines |
| `jax.lax` | Lower-level primitives with stricter shape and dtype rules |
| JAX transformations | Functions such as `jit`, `grad`, `vmap`, and `shard_map` that transform other functions |
| OpenXLA and PJRT | Compiler and runtime interfaces used for optimization, device execution, and pluggable backends |

Many `jax.numpy` operations are implemented through `jax.lax` primitives. These primitives have transformation rules that specify behavior under differentiation, batching, lowering, and other transformations. This is why a single operation can participate in several transformations without each application library implementing those mechanisms again.[6][23]

The backend is deliberately modular. JAX calls the OpenXLA compiler and PJRT runtime for compilation, memory management, and device execution. The project describes XLA and PJRT as external OpenXLA projects rather than components governed by the JAX core team. PJRT's plug-in interface allows a hardware vendor to provide a backend without placing the full implementation in JAX itself.[12][31]

## Function transformations

JAX's best-known APIs accept a function and return another function. Transformations can be nested and composed, although valid composition still depends on the operations and types used by the transformed function.[1][5][6]

| Transformation | Main purpose | Important boundary |
|---|---|---|
| `jax.jit` | Trace, lower, compile, and cache a function for a target backend | Dynamic Python behavior and value-dependent shapes may not be representable |
| `jax.grad` | Return the gradient of a scalar-output function | Differentiated arguments normally need inexact dtypes, and every primitive on the path needs a derivative rule |
| `jax.value_and_grad` | Return a function value and its gradient together | Has the same differentiation constraints as `grad` |
| `jax.jvp` | Compute a Jacobian-vector product for forward-mode differentiation | Cost is favorable when the input dimension is small relative to the output dimension |
| `jax.vjp` | Compute a vector-Jacobian product for reverse-mode differentiation | Reverse mode retains or reconstructs information needed by the pullback |
| `jax.jacfwd` | Construct a Jacobian through forward-mode differentiation | Often preferable for tall Jacobians |
| `jax.jacrev` | Construct a Jacobian through reverse-mode differentiation | Often preferable for wide Jacobians |
| `jax.vmap` | Add a mapped array axis by transforming primitive operations | It is vectorization, not a request to launch one Python call per item |
| `jax.shard_map` | Express a per-shard program over a named device mesh | The user specifies shard placement and explicit collectives |

### Automatic differentiation

JAX implements forward-mode and reverse-mode automatic differentiation. Forward mode propagates Jacobian-vector products alongside primal values. Reverse mode first evaluates the primal program and then propagates cotangents backward through vector-Jacobian products. `jax.grad` is a convenience interface for reverse-mode differentiation of scalar-output functions.[7]

The distinction matters computationally. For a function from an input space of dimension \(n\) to an output space of dimension \(m\), forward mode can obtain a full Jacobian by applying the JVP to \(n\) basis vectors, while reverse mode can obtain it by applying the VJP to \(m\) output basis vectors. `jacfwd` and `jacrev` expose these choices, and the two modes can be composed to calculate higher derivatives such as Hessians or Hessian-vector products.[7]

Differentiation is a program transformation, not symbolic algebra over source text and not numerical finite differencing. JAX traces the operations actually selected for the supplied abstract inputs, then applies registered linearization and transpose rules. Users can define custom derivative behavior for JAX-transformable Python functions with `custom_jvp` and `custom_vjp`. Foreign functions and new primitives require explicit rules if they are to participate in differentiation.[7][22]

Nondifferentiability does not disappear. A mathematical function can lack a unique derivative at a point, and a primitive may define a conventional value there. Integer and Boolean operations generally do not provide the same differentiable structure as floating-point or complex operations. Control-flow primitives have their own differentiation properties, and reverse-mode differentiation through a loop can require bounds or intermediate storage that make some loop forms unsupported or expensive.[7][36]

### Compilation with `jit`

`jax.jit` traces a function for abstract input values, lowers the traced program, and asks the compiler to build an executable for a target platform. Calls with a compatible cache key can reuse that executable. Calls with materially different shapes, dtypes, static arguments, function identities, device topologies, or relevant compiler settings may trigger another trace or compilation.[5][6][33]

Compilation can fuse operations and choose hardware-specific implementations, but the first call includes tracing and compilation work. Small computations, functions called only once, and workloads dominated by transfers or unsupported host code may not benefit. A fair comparison therefore reports compilation separately or states whether compilation was amortized over repeated calls.[5][11]

JAX provides an optional persistent compilation cache. Its key includes the computation, `jaxlib` version, relevant compiler flags, device configuration, compression method, and a custom-hook value. The cache can avoid recompiling compatible programs across process runs. It is treated as trusted executable material, so the project warns against using a cache directory writable by untrusted users.[33]

### Vectorization with `vmap`

`vmap` transforms a function written for an individual item into one that maps across a chosen array axis. Instead of running a Python loop, it applies batching rules to the function's primitives. For example, a vectorized matrix-vector operation may lower to a matrix-matrix operation. `in_axes` and `out_axes` specify which axes are mapped and how mapped results are placed.[1][5]

Vectorization composes with differentiation and compilation. A common use is `vmap(grad(loss))` for per-example gradients, followed by `jit` to compile the combined computation. This does not guarantee that every vectorized form is faster or fits in memory. Batching can increase intermediate sizes, and a primitive without an appropriate batching rule can fail or require a different formulation.[1][7]

## Tracing, jaxprs, and lowering

### Tracers and abstract values

When JAX transforms a function, it commonly executes the Python function with tracer objects instead of concrete runtime arrays. A tracer carries abstract information such as shape and dtype while recording operations. Python executes during tracing, while the recorded numerical operations execute later in the compiled program. This distinction explains several common surprises: a normal `print` can run at trace time, a branch on an unknown tracer value cannot be resolved by Python, and side effects can happen once during tracing rather than once for every device execution.[5][6][8]

A Python value that controls program structure can sometimes be marked static, which makes it part of the compilation cache key. Static arguments should be hashable and stable. Marking a mutable object static can produce stale results because JAX may not see internal mutations as a reason to retrace. Conversely, continually creating new function objects or changing static values can cause avoidable cache misses and repeated compilation.[8][33]

### Jaxpr

JAX represents traced functional programs in an intermediate language called jaxpr, short for JAX expression. A jaxpr contains variables, literals, primitive equations, and nested jaxprs for structured operations. `jax.make_jaxpr` exposes this representation for inspection. Jaxprs are typed by abstract values and omit arbitrary Python object behavior that cannot be expressed in the JAX language.[6]

Jaxpr is central to composability. Differentiation can turn one jaxpr into another, batching can add mapped dimensions, and partial evaluation can separate known from unknown values. A lowering path then translates the transformed program into compiler IR. Users normally interact with public transformations rather than jaxpr internals, but jaxpr helps explain what JAX can and cannot stage.[6][23]

### StableHLO, XLA, and PJRT

Compiled JAX computations are lowered through MLIR-based representations that include StableHLO. StableHLO defines a portable set of high-level operations for exchange between frameworks and compilers. XLA accepts high-level operation graphs, performs optimizations and partitioning, and produces target-specific executable code. PJRT provides the runtime and plug-in boundary used to manage devices, buffers, compilation, and execution.[12][20][31][32]

These layers have different ownership and compatibility contracts. JAX's Python API policy does not automatically apply to compiler internals, `jax.experimental`, or `jax.extend`. StableHLO's general portability goals do not mean every JAX custom call is portable. An exported program can contain custom calls whose C++ targets must also exist in the consumer runtime.[15][20][23]

The `jax.export` API packages StableHLO together with metadata needed to call the exported function. Official documentation gives a compatibility window for the JAX export format: a consumer can be up to six months newer than the exporting JAX version and up to three weeks older, subject to the documented restrictions. Bypassing `jax.export` and taking raw compiler IR does not receive that guarantee. Deserialized artifacts must be trusted because they may invoke registered custom calls.[20]

## Purity, state, randomness, and pytrees

### Functional purity

Transformations are designed for functionally pure functions: outputs should depend on explicit inputs, and program-observable state should flow through return values. Reading a mutable global, consuming an iterator, writing to a list, or mutating an object inside a transformed function can produce tracing-time behavior or incorrect assumptions. JAX supplies callback and effect mechanisms for specific uses, but ordinary side effects are not silently converted into device-side effects.[8]

Stateful algorithms are usually expressed by passing state into a function and returning the new state. A training step might accept model parameters, optimizer state, a random key, and a batch, then return updated parameters and state. Higher-level libraries can present object-oriented interfaces, but they ultimately cross the JAX transformation boundary through explicit values or registered structures.[8][27][28]

### Random-number keys

JAX does not use NumPy's implicit global random state for transformed random operations. Random functions take an explicit key. A program creates a key, splits it to obtain independent subkeys, and passes each subkey to the operation that consumes it. Reusing a key repeats the same pseudorandom computation rather than advancing hidden state.[9]

JEP 263 specifies the original design goals: reproducibility across compilation boundaries, vectorization, parallel execution, and the absence of unnecessary sequencing between random calls. Its core design used a counter-based Threefry generator with functional key splitting. Later typed-key APIs and pluggable implementations refined the representation, but explicit data flow remains the central user model.[9]

JAX does not promise bit-for-bit random values across every release. The compatibility policy states that a pseudorandom function's distribution is covered, while the exact sample for a fixed key may change across releases. Programs that need repeatability should pin relevant versions, record configuration, and test the expected statistical and numerical behavior.[15]

### Pytrees

A pytree is a nested container structure whose leaves are arrays or other values. Built-in container nodes include lists, tuples, and dictionaries. User-defined types can be registered with flattening and unflattening rules. Transformations flatten pytrees at their API boundary, operate over the leaves, and reconstruct the original structure for outputs.[10]

Pytrees let parameters, gradients, optimizer state, and batches retain meaningful Python structure. Axis specifications for `vmap`, differentiation argument selections, and sharding specifications can themselves be tree prefixes that align with an argument pytree. A custom class is treated as a leaf unless registered, so registration is necessary when JAX should see and transform its internal dynamic fields.[10]

## Control flow and shape constraints

Python control flow runs while a function is traced. A Python `if` whose condition depends on a tracer cannot choose a branch because the concrete value is unavailable. JAX provides staged alternatives such as `lax.cond`, `lax.switch`, `lax.scan`, `lax.fori_loop`, and `lax.while_loop`. These primitives place control flow in the traced program rather than deciding it in Python.[36]

Static shapes are another practical boundary. A compiled operation normally needs output ranks and dimension sizes that can be determined from abstract inputs. Filtering an array by a data-dependent Boolean mask, for example, produces a result whose length depends on runtime values and cannot be represented as an ordinary statically shaped output. A fixed-size formulation can instead return a mask, a padded buffer, a count, or a reduction.[5][8]

Shape polymorphism and export support relax some constraints by representing selected dimensions symbolically, but they do not turn arbitrary value-dependent allocation into a static program. The supported symbolic relations, operations, and consumer compatibility must still be checked for the specific export path.[20]

Loop choice also affects differentiation and compilation. `lax.scan` represents a fixed-length iteration compactly and can avoid unrolling a long Python loop into a large jaxpr. `fori_loop` and `while_loop` express other structured cases. Their derivative support differs because reverse mode may need to recover intermediate values, and a loop with an unknown number of iterations cannot always be reversed with bounded storage.[36]

## Parallelism and distributed arrays

### Device meshes and shardings

JAX models device topology with a `Mesh`, a multidimensional grid of devices whose axes have names. A `PartitionSpec`, available through the alias `jax.P`, maps array dimensions to mesh axes. A `NamedSharding` combines a concrete mesh with a partition specification and describes where an array's data is stored.[16]

Sharding is a property of `jax.Array`, not just an annotation on a model. An array can be replicated on a mesh axis, partitioned along one or more axes, or represented with other documented distributed semantics. The order and shape of a mesh can matter because they determine which devices communicate along each logical axis and whether that mapping matches the physical network.[16]

JAX documentation distinguishes three parallel-programming styles:

| Style | User view | Who specifies data sharding? | Who writes collectives? |
|---|---|---|---|
| Compiler-based automatic sharding | Global array | Compiler, guided by constraints | Compiler |
| Explicit sharding with automatic partitioning | Global array | User-visible JAX types | Compiler |
| Manual per-device programming | Local shard | User | User through collectives |

These styles can be combined, but they expose different tradeoffs. Automatic partitioning reduces per-device code, explicit sharding exposes placement earlier in the program, and manual programming gives direct control over collectives and shard-local shapes.[16][17]

### `jit` with sharding

`jax.jit` can compile a global-view function whose inputs and outputs have sharding specifications. The compiler partitions operations and inserts communication needed to preserve the function's semantics. Constraints can guide intermediate layouts. This supports common [data parallelism](https://aiwiki.ai/wiki/data_parallelism), [model parallelism](https://aiwiki.ai/wiki/model_parallelism), and [tensor parallelism](https://aiwiki.ai/wiki/tensor_parallelism) patterns without requiring every operation to be written as device-local code.[16]

Automatic partitioning is not cost-free. A change in layout can cause resharding, and a poorly chosen mesh can create expensive cross-device communication. The correct strategy depends on array dimensions, model structure, collective bandwidth, memory capacity, and topology. JAX provides tools to inspect array shardings, but it does not infer an application-independent optimal cluster layout.[16]

### `shard_map`

`jax.shard_map` expresses an SPMD function over the shards of arrays. The mapped function receives shard-local values according to `in_specs`, produces values assembled according to `out_specs`, and can use named collective operations. Unlike `vmap` and legacy `pmap` behavior, current `shard_map` programming is described as rank-preserving in the migration documentation.[17][19]

The manual form is useful when the communication pattern is part of the algorithm or when the compiler needs explicit guidance. It also requires the author to maintain correct collective semantics and consistent shard specifications. JAX documentation presents `shard_map` as complementary to automatic partitioning under `jit`, not as a replacement for every sharded `jit` program.[17]

### The `pmap` transition

`jax.pmap` historically replicated a function across local devices and exposed an implicit mapped axis. JAX 0.8.0 changed its implementation to use `jit` and `shard_map`, placed the API in maintenance mode, and recommended `shard_map` for new code. JAX 0.10.0 removed the old C++ implementation and associated public types such as `PmapSharding`.[14][19]

Existing `pmap` code is not necessarily invalid, but migration can change rank behavior, sharding objects, data-placement calls, and collective setup. The official guide recommends `jax.jit(jax.shard_map(...))` for explicit migrations and documents replacements for removed replication and sharding helpers.[19]

### Multi-process execution

For a multi-host GPU cluster or Cloud TPU deployment, multiple Python processes run the same program and initialize JAX's distributed system. `jax.distributed.initialize()` must run before device access. It lets processes discover the global topology, performs health checks, and participates in distributed checkpointing.[18]

A global `jax.Array` can span processes, while each process can directly address only its local shards. Every process must execute participating collective computations in a consistent order. If one process skips a collective or follows a different sequence, the distributed program can hang or fail. JAX does not launch remote processes by itself; a scheduler or other orchestration system starts them.[18]

## Hardware and installation

The installation consists of a pure-Python `jax` package and compiled runtime components associated with `jaxlib` and platform plug-ins. The correct installation command depends on the operating system, architecture, and accelerator. The official support matrix, checked for this article at the July 28, 2026 cutoff, marked the following combinations as supported or experimental.[13]

| Platform path | Status in official installation documentation at cutoff |
|---|---|
| CPU on Linux x86-64 and AArch64 | Supported |
| CPU on macOS Apple ARM | Supported |
| CPU on Windows x86-64 | Marked as supported in the summary matrix; the CPU-wheel instructions still label Windows x86-64 experimental |
| NVIDIA GPU on Linux x86-64 and AArch64 | Supported through CUDA plug-in packages |
| NVIDIA GPU on Windows Subsystem for Linux | Experimental |
| Google Cloud TPU on Linux x86-64 | Supported |
| AMD GPU on Linux x86-64 | Supported through an AMD-maintained ROCm plug-in |
| AMD GPU on Windows Subsystem for Linux | Experimental |
| Apple GPU on macOS Apple ARM | Experimental plug-in path |
| Intel GPU on Linux x86-64 | Experimental third-party plug-in path |

This table describes software support, not equal feature coverage or performance. A plug-in may lag JAX releases, support a subset of operations, or require specific driver and toolkit versions. The official installation page should be checked for the exact release rather than copying a command from an older article.[13]

NVIDIA GPU installations have separate CUDA 12 and CUDA 13 paths. The July 2026 documentation recommended the CUDA 13 packages for new installations and listed different minimum driver and compute-capability requirements for the two lines. [CUDA](https://aiwiki.ai/wiki/cuda), cuDNN, and NCCL compatibility therefore belongs to the installed package version, not to JAX as an timeless property.[13]

AMD support is supplied through a [ROCm](https://aiwiki.ai/wiki/rocm) JAX plug-in maintained by AMD, and Apple and Intel accelerator support use plug-in approaches with experimental status in the JAX matrix. The project's backend architecture deliberately places much hardware enablement in OpenXLA and PJRT. This lets vendors add support while also making it necessary to distinguish core-JAX support from a vendor plug-in's own release and compatibility promises.[12][13]

## Extension and interoperability

### Custom derivatives and primitives

The simplest extension is a Python function built entirely from existing JAX operations. It inherits the transformation behavior of those operations. When the mathematically appropriate derivative differs from the automatically generated one, `custom_jvp` or `custom_vjp` can supply a rule while keeping the function in the ordinary JAX language.[7]

A genuinely new primitive needs abstract-evaluation, lowering, differentiation, batching, and other rules for each transformation it should support. `jax.extend` exposes selected internal machinery for extension authors. It is a semi-public interface without the compatibility guarantee of the main API, and the project recommends testing downstream libraries against nightly releases to discover changes early.[15][23]

### Foreign function interface

`jax.ffi` calls external compiled code through XLA's foreign function interface. It is useful when an optimized CPU or GPU routine already exists or when a needed operation cannot reasonably be expressed in JAX. The Python side registers a target and defines an `ffi_call` with abstract result information.[22]

JAX cannot infer how to differentiate or vectorize arbitrary foreign code. Extension authors must provide the relevant rules. The official FFI tutorial describes the interface as a last-resort path because built-in primitives or Pallas usually carry lower development and maintenance costs. FFI targets can also limit export portability if the consumer runtime lacks the registered target.[20][22]

### Pallas

[Pallas](https://aiwiki.ai/wiki/jax_pallas) is JAX's experimental custom-kernel system for GPUs and TPUs. It exposes block specifications, grid execution, references to accelerator memory, and lower-level control while reusing JAX tracing and array expressions. Pallas targets Mosaic GPU on supported GPUs and Mosaic on TPUs in current documentation.[21]

Pallas is explicitly experimental and changes frequently. A valid kernel must respect backend-specific memory layouts, supported operations, synchronization, and hardware restrictions. A benchmark for one Pallas kernel on one accelerator does not establish a general JAX speedup, and a kernel written for a particular tile shape may need revision for another device generation.[21]

## The surrounding ecosystem

JAX core intentionally leaves several common [neural network](https://aiwiki.ai/wiki/neural_network) facilities to external packages. Treating every JAX-compatible library as part of JAX obscures differences in maintainership, versioning, API stability, and scope. The project instead describes a decentralized ecosystem built on a focused core.[12]

| Project | Function | Relationship to JAX core |
|---|---|---|
| Flax | Neural-network module and model APIs, including NNX and Linen | Separate library built on JAX; current Flax documentation encourages new users to use NNX while stating that Linen is not scheduled for near-term deprecation.[27] |
| Optax | Composable gradient transformations and optimization algorithms | Separate JAX library maintained by its own contributors.[28] |
| Orbax | Checkpointing, distributed array persistence, and checkpoint lifecycle management | Separate JAX-native persistence project.[29] |
| Grain | Data loading for training and evaluating JAX models | Separate data-pipeline library.[30] |
| Haiku, RLax, Chex, and Jraph | Neural-network, reinforcement-learning, testing, and graph utilities described by DeepMind in 2020 | Separate projects in the DeepMind JAX ecosystem, not JAX core modules.[24] |

Flax NNX permits Python reference semantics and mutability in its user-facing object model. At transformation boundaries, its functional APIs split, merge, or otherwise expose JAX-compatible state. This is an example of an ecosystem library adapting an object-oriented model to JAX rather than evidence that core JAX arrays have become mutable.[27]

Optax represents an optimizer as transformations with initialization and update functions. Parameters and optimizer states are commonly pytrees. Optax returns updates, which application code applies to parameters, rather than owning a hidden parameter store. This matches the explicit-state pattern used by transformed JAX functions.[28]

Orbax handles checkpoint storage and restoration for pytrees and distributed arrays. Grain supplies samplers, data sources, and transformations for input pipelines. Neither capability is provided by `jax.grad` or `jax.jit`, and neither package's release is automatically synchronized with a JAX release. Compatibility must be checked across the chosen versions.[29][30]

## Documented uses

### Machine-learning research

Google DeepMind reported in December 2020 that JAX had been widely adopted within its research community during the preceding year. The same account described a collection of separately released libraries, including Haiku, Optax, RLax, Chex, and Jraph. This is first-party evidence for institutional research use at that time, not a claim that every later DeepMind model used JAX.[24]

The PaLM paper provides a bounded large-scale training example. The authors reported training the 540-billion-parameter [PaLM](https://aiwiki.ai/wiki/palm) model on 6,144 TPU v4 chips across two pods using Pathways and a JAX-based training system. They reported 46.2 percent model FLOPs utilization and 57.8 percent hardware FLOPs utilization for that configuration.[25]

Those PaLM figures should not be read as general JAX benchmarks. They refer to a named model, a particular parallelization and software stack, 6,144 TPU v4 chips, and the utilization definitions used in the paper. Changing model shape, batch size, accelerator generation, compiler build, input pipeline, or utilization denominator changes the result.[25]

JAX is also used to implement higher-level training systems, but an article about JAX should not infer model provenance from an ecosystem dependency or a job advertisement. A reliable attribution should come from a model paper, official technical report, or code release that explicitly identifies JAX. Claims about Claude, Grok, Gemini, or other model families are therefore excluded here unless the specific model and training system have a primary citation.

### Scientific computing

JAX MD illustrates use outside conventional neural-network training. Its NeurIPS 2020 paper described differentiable molecular-dynamics simulations, statistical-physics environments, interaction potentials, and neural networks implemented with JAX primitives. The authors reported differentiating complete simulation trajectories for meta-optimization and scaling selected spatial primitives to hundreds of thousands of particles on a single GPU.[26]

That claim is bounded to the JAX MD implementation and experiments in the paper. It does not mean every molecular simulation scales to that particle count or that a JAX implementation is automatically faster than a specialized simulator. Interaction range, neighbor-list behavior, precision, integration method, and hardware determine practical cost.[26]

The same transformation model supports other scientific tasks: derivatives of objective functions, sensitivity calculations, batched simulations, and compiled numerical kernels. Suitability depends on whether the required operations have correct JAX implementations and whether static-shape and accelerator constraints match the problem.[5][7][8]

## Performance evaluation

### Compilation and execution

A JAX timing should distinguish at least four costs:

1. Python tracing and jaxpr construction.
2. Compiler lowering, optimization, and executable construction.
3. Transfer of inputs and outputs between host and device.
4. Execution of the compiled computation.

The first invocation of a new compiled signature may include all four. A later invocation can reuse cached tracing and compilation work. Reporting only the later time is valid for a steady-state loop, while reporting only the first call describes startup latency. Neither should be mislabeled as the other.[5][11][33]

Device dispatch is asynchronous. A call can return a `jax.Array` before the accelerator has finished computing its value. Timing only the Python call measures dispatch overhead. The official guidance uses `block_until_ready()` or conversion to a host NumPy value to wait for completion. Host conversion includes transfer overhead, while `block_until_ready()` can measure device completion without necessarily copying the result to the host.[11]

### Benchmark scope

A reproducible benchmark records the JAX and `jaxlib` versions, backend and plug-in, compiler flags, hardware model and count, driver and toolkit versions, input shapes, dtypes, sharding, warmup policy, number of repetitions, synchronization method, and whether compilation is included. Distributed measurements also need mesh topology and communication context.[11][13][16]

Shape changes can cause recompilation. A benchmark that cycles through many shapes may spend substantial time compiling, while one fixed-shape kernel may amortize compilation almost completely. The persistent cache can change startup results across process runs, so its state and location should also be recorded.[33]

Numerical equality across platforms is a separate question from elapsed time. JAX's compatibility policy does not promise exact numerical values across accelerator platforms, inside versus outside `jit`, or across releases. Floating-point reduction order, library kernels, and compiler transformations can change low-order bits. Tests should use tolerances appropriate to the algorithm and dtype rather than assume bitwise equality unless a specific operation guarantees it.[15]

## Limitations and engineering tradeoffs

### Staging boundaries

JAX transforms the supported numerical portion of a program, not the full Python runtime. Dynamic object creation, reflection, arbitrary I/O, and value-dependent Python branching do not become accelerator code merely because a function is decorated with `jit`. Applications commonly keep data orchestration, logging, checkpoint decisions, and irregular preprocessing in Python while staging dense numerical steps.[5][8][36]

Static arguments can bridge some Python choices into a compiled function, but every distinct static value can create another compiled variant. Excessive specialization increases compilation time and memory use. Turning every argument static defeats the purpose of compiling a reusable array program.[8][33]

### Memory and transfers

Accelerator arrays consume device memory, and temporary values created by a computation can exceed the size of visible inputs and outputs. Automatic differentiation may retain intermediates for the backward pass. Vectorization can materialize a larger batch, and sharding can replicate values on some mesh axes. Checkpointing or rematerialization can trade recomputation for memory, but it does not eliminate the need to reason about peak live buffers.[7][16]

Transfers between host and device can dominate small computations. Accidentally converting an array to NumPy or inspecting values inside a loop can synchronize the program and copy data. Conversely, asynchronous dispatch can let Python run ahead and enqueue work. Performance analysis should therefore include placement and synchronization, not only arithmetic operation counts.[11]

### Debugging

Ordinary Python `print` reports tracers inside compiled code because it runs during tracing. JAX provides `jax.debug.print`, `jax.debug.breakpoint`, and other debugging facilities that operate with staged computations. The `jax_disable_jit` option can expose ordinary eager behavior during diagnosis, and `jax_debug_nans` can help locate invalid values, although disabling JIT can change execution and should not be treated as an identical performance environment.[5][34]

Errors involving `TracerBoolConversionError`, `TracerIntegerConversionError`, or non-concrete Boolean indexing usually indicate that Python requested a concrete value or dynamic shape during tracing. The fix is problem-specific: mark a truly compile-time value static, use a JAX control-flow primitive, keep an array fixed-size, or move the operation outside the compiled region.[5][34][36]

### API evolution

JAX evolves quickly. The main public API normally receives a three-month deprecation period for breaking changes, but `jax.experimental` may change with less warning and `jax.extend` has no compatibility guarantee between releases. Numerical and random outputs also have narrower guarantees than API names.[15][23]

The `pmap` transition demonstrates the practical effect. Code that relied on `PmapSharding`, legacy rank reduction, or removed device-placement helpers needed migration even though `jax.pmap` itself remained available. Production systems should pin versions, read the change log, test nightlies where feasible, and verify dependent ecosystem libraries before an upgrade.[14][19]

### Core versus ecosystem support

An installation problem may belong to JAX core, `jaxlib`, an OpenXLA backend, a PJRT plug-in, a driver, or an ecosystem library. The support boundary matters when filing an issue and when evaluating maturity. The official installation documentation routes AMD, Apple, and Intel accelerator issues to their respective plug-in projects where appropriate.[12][13]

Likewise, a Flax module error, an Optax optimizer behavior, an Orbax checkpoint format, or a Grain input-pipeline policy is not automatically a JAX core behavior. These libraries share conventions and often interoperate through arrays and pytrees, but they maintain separate APIs and release histories.[12][27][28][29][30]

## Relationship to other numerical systems

JAX's `jax.numpy` interface resembles NumPy, but device placement, immutability, dtype defaults, asynchronous dispatch, and tracing make the execution model different. NumPy remains appropriate for host-oriented numerical work and for Python programs that need unrestricted dynamic behavior. JAX is most useful when a substantial numerical region can benefit from transformation or accelerator execution.[5][8][11][35]

Compared with an end-to-end deep-learning framework, JAX core exposes lower-level transformations and leaves module, optimizer, input, and persistence abstractions to separate projects. This permits several model-library designs, but it also requires users to choose and version more components. A direct feature table against PyTorch or TensorFlow becomes stale quickly because all three systems have eager execution, compilation, vectorization, distributed APIs, and accelerator backends that change independently.[12][15]

JAX's compiler path is also not unique ownership of XLA. XLA and StableHLO are OpenXLA projects used by multiple front ends. JAX contributes a Python array language, transformations, lowering rules, and runtime integration around that stack. It should therefore be described as using OpenXLA, not as being identical to XLA or as the only framework able to target it.[12][31][32]

## Current status at the research cutoff

JAX 0.11.0, released July 16, 2026, was the latest entry in the official change log before the July 28 cutoff. Its release notes added `jax.custom_remat`, changes to checkpoint-policy organization, and an inlining-policy enum for `jax.jit`. It also removed previously deprecated APIs and dropped some older Python, NumPy, and SciPy versions under the support policy.[14]

Version 0.10.2 had been released on June 17, 2026, and 0.10.1 on May 20, 2026. The earlier 0.10.0 release on April 16 removed the legacy C++ `pmap` implementation. These dates supersede the prior article's statement that 0.10.1 was current in mid-2026.[14]

The project continued to label Pallas experimental, to place `pmap` in maintenance mode, and to recommend modern sharding APIs for new parallel code.[14][19][21] Platform support still varied by operating system and vendor plug-in, so "runs on GPUs" should not be read as uniform support for every GPU and operating-system combination.[13]

## See also

- [Automatic differentiation](https://aiwiki.ai/wiki/automatic_differentiation)
- [NumPy](https://aiwiki.ai/wiki/numpy)
- [XLA](https://aiwiki.ai/wiki/xla)
- [Tensor Processing Unit](https://aiwiki.ai/wiki/tpu)
- [GPU](https://aiwiki.ai/wiki/gpu)
- [Distributed training](https://aiwiki.ai/wiki/distributed_training)
- [Data parallelism](https://aiwiki.ai/wiki/data_parallelism)
- [Model parallelism](https://aiwiki.ai/wiki/model_parallelism)
- [Pallas (JAX kernel language)](https://aiwiki.ai/wiki/jax_pallas)
- [PaLM](https://aiwiki.ai/wiki/palm)

## References

1. JAX authors. "JAX: Python library for accelerator-oriented array computation and program transformation." jax-ml/jax repository README. https://github.com/jax-ml/jax Accessed 2026-07-28.
2. JAX authors. "Apache License, Version 2.0." jax-ml/jax repository. https://github.com/jax-ml/jax/blob/main/LICENSE Accessed 2026-07-28.
3. Frostig, Roy, Matthew James Johnson, and Chris Leary. "Compiling machine learning programs via high-level tracing." SysML 2018. https://cs.stanford.edu/~rfrostig/pubs/jax-mlsys2018.pdf
4. Maclaurin, Dougal, David Duvenaud, and Ryan P. Adams. "Autograd: Effortless gradients in NumPy." ICML 2015 AutoML Workshop. https://indico.ijclab.in2p3.fr/event/2914/contributions/6483/subcontributions/180/attachments/6060/7185/automl-short.pdf
5. JAX authors. "Quickstart: How to think in JAX." JAX documentation. https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html Accessed 2026-07-28.
6. JAX authors. "Key concepts." JAX documentation. https://docs.jax.dev/en/latest/key-concepts.html Accessed 2026-07-28.
7. JAX authors. "The Autodiff Cookbook." JAX documentation. https://docs.jax.dev/en/latest/notebooks/autodiff_cookbook.html Accessed 2026-07-28.
8. JAX authors. "JAX: The Sharp Bits." JAX documentation. https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html Accessed 2026-07-28.
9. JAX authors. "JAX PRNG Design." JAX Enhancement Proposal 263. https://docs.jax.dev/en/latest/jep/263-prng.html Accessed 2026-07-28.
10. JAX authors. "Pytrees." JAX documentation. https://docs.jax.dev/en/latest/pytrees.html Accessed 2026-07-28.
11. JAX authors. "Asynchronous dispatch." JAX documentation. https://docs.jax.dev/en/latest/async_dispatch.html Accessed 2026-07-28.
12. JAX authors. "About the project." JAX documentation. https://docs.jax.dev/en/latest/about.html Accessed 2026-07-28.
13. JAX authors. "Installation." JAX documentation. https://docs.jax.dev/en/latest/installation.html Accessed 2026-07-28.
14. JAX authors. "Change log." JAX documentation. https://docs.jax.dev/en/latest/changelog.html Accessed 2026-07-28.
15. JAX authors. "API compatibility." JAX documentation. https://docs.jax.dev/en/latest/api_compatibility.html Accessed 2026-07-28.
16. JAX authors. "Distributed arrays and automatic parallelization." JAX documentation. https://docs.jax.dev/en/latest/parallel.html Accessed 2026-07-28.
17. JAX authors. "Manual parallelism with shard_map." JAX documentation. https://docs.jax.dev/en/latest/notebooks/shard_map.html Accessed 2026-07-28.
18. JAX authors. "Introduction to multi-controller JAX." JAX documentation. https://docs.jax.dev/en/latest/multi_process.html Accessed 2026-07-28.
19. JAX authors. "Migrating to the new jax.pmap." JAX documentation. https://docs.jax.dev/en/latest/migrate_pmap.html Accessed 2026-07-28.
20. JAX authors. "Exporting and serializing staged-out computations." JAX documentation. https://docs.jax.dev/en/latest/export/export.html Accessed 2026-07-28.
21. JAX authors. "Pallas: a JAX kernel language." JAX documentation. https://docs.jax.dev/en/latest/pallas/ Accessed 2026-07-28.
22. JAX authors. "Foreign function interface (FFI)." JAX documentation. https://docs.jax.dev/en/latest/ffi.html Accessed 2026-07-28.
23. JAX authors. "jax.extend module." JAX documentation. https://docs.jax.dev/en/latest/jax.extend.html Accessed 2026-07-28.
24. Budden, David, and Matteo Hessel. "Using JAX to accelerate our research." Google DeepMind, 2020-12-04. https://deepmind.google/blog/using-jax-to-accelerate-our-research/
25. Chowdhery, Aakanksha, et al. "PaLM: Scaling Language Modeling with Pathways." Journal of Machine Learning Research 24, no. 240 (2023): 1-113. https://www.jmlr.org/papers/volume24/22-1144/22-1144.pdf
26. Schoenholz, Samuel S., and Ekin Dogus Cubuk. "JAX MD: A Framework for Differentiable Physics." Advances in Neural Information Processing Systems 33 (2020). https://papers.neurips.cc/paper_files/paper/2020/file/83d3d4b6c9579515e1679aca8cbc8033-Paper.pdf
27. Flax contributors. "Flax: Neural Networks for JAX." Flax documentation. https://flax.readthedocs.io/en/stable/ Accessed 2026-07-28.
28. Optax contributors. "Optax." Optax documentation. https://optax.readthedocs.io/en/latest/ Accessed 2026-07-28.
29. Orbax contributors. "Orbax." Orbax documentation. https://orbax.readthedocs.io/en/latest/ Accessed 2026-07-28.
30. Grain contributors. "Grain: Feeding JAX Models." Grain documentation. https://google-grain.readthedocs.io/ Accessed 2026-07-28.
31. OpenXLA Project. "XLA." https://openxla.org/xla Accessed 2026-07-28.
32. OpenXLA Project. "StableHLO." https://openxla.org/stablehlo Accessed 2026-07-28.
33. JAX authors. "Persistent compilation cache." JAX documentation. https://docs.jax.dev/en/latest/persistent_compilation_cache.html Accessed 2026-07-28.
34. JAX authors. "Introduction to debugging." JAX documentation. https://docs.jax.dev/en/latest/debugging.html Accessed 2026-07-28.
35. JAX authors. "Default dtypes and the X64 flag." JAX documentation. https://docs.jax.dev/en/latest/default_dtypes.html Accessed 2026-07-28.
36. JAX authors. "Control flow and logical operators with JIT." JAX documentation. https://docs.jax.dev/en/latest/control-flow.html Accessed 2026-07-28.