JAX

JAX is an open-source Python library for accelerator-oriented array computation and program transformation. Its core interface combines a NumPy-inspired array API with transformations for compilation, 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 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, 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 programs. It used 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

Scroll sideways for more →

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:

Scroll sideways for more →

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]

Scroll sideways for more →

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:

Scroll sideways for more →

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, model parallelism, and 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]

Scroll sideways for more →

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, 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 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 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 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]

Scroll sideways for more →

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 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
• NumPy
• XLA
• Tensor Processing Unit
• GPU
• Distributed training
• Data parallelism
• Model parallelism
• Pallas (JAX kernel language)
• PaLM

References

1. ^JAX authors. "JAX: Python library for accelerator-oriented array computation and program transformation." jax-ml/jax repository README. github.com/...jax Accessed 2026-07-28.
2. ^JAX authors. "Apache License, Version 2.0." jax-ml/jax repository. github.com/...LICENSE Accessed 2026-07-28.
3. ^Frostig, Roy, Matthew James Johnson, and Chris Leary. "Compiling machine learning programs via high-level tracing." SysML 2018. cs.stanford.edu/...jax-mlsys2018.pdf
4. ^Maclaurin, Dougal, David Duvenaud, and Ryan P. Adams. "Autograd: Effortless gradients in NumPy." ICML 2015 AutoML Workshop. indico.ijclab.in2p3.fr/...automl-short.pdf
5. ^JAX authors. "Quickstart: How to think in JAX." JAX documentation. docs.jax.dev/...thinking_in_jax Accessed 2026-07-28.
6. ^JAX authors. "Key concepts." JAX documentation. docs.jax.dev/...key-concepts Accessed 2026-07-28.
7. ^JAX authors. "The Autodiff Cookbook." JAX documentation. docs.jax.dev/...autodiff_cookbook Accessed 2026-07-28.
8. ^JAX authors. "JAX: The Sharp Bits." JAX documentation. docs.jax.dev/...Common_Gotchas_in_JAX Accessed 2026-07-28.
9. ^JAX authors. "JAX PRNG Design." JAX Enhancement Proposal 263. docs.jax.dev/...263-prng Accessed 2026-07-28.
10. ^JAX authors. "Pytrees." JAX documentation. docs.jax.dev/...pytrees Accessed 2026-07-28.
11. ^JAX authors. "Asynchronous dispatch." JAX documentation. docs.jax.dev/...async_dispatch Accessed 2026-07-28.
12. ^JAX authors. "About the project." JAX documentation. docs.jax.dev/...about Accessed 2026-07-28.
13. ^JAX authors. "Installation." JAX documentation. docs.jax.dev/...installation Accessed 2026-07-28.
14. ^JAX authors. "Change log." JAX documentation. docs.jax.dev/...changelog Accessed 2026-07-28.
15. ^JAX authors. "API compatibility." JAX documentation. docs.jax.dev/...api_compatibility Accessed 2026-07-28.
16. ^JAX authors. "Distributed arrays and automatic parallelization." JAX documentation. docs.jax.dev/...parallel Accessed 2026-07-28.
17. ^JAX authors. "Manual parallelism with shard_map." JAX documentation. docs.jax.dev/...shard_map Accessed 2026-07-28.
18. ^JAX authors. "Introduction to multi-controller JAX." JAX documentation. docs.jax.dev/...multi_process Accessed 2026-07-28.
19. ^JAX authors. "Migrating to the new jax.pmap." JAX documentation. docs.jax.dev/...migrate_pmap Accessed 2026-07-28.
20. ^JAX authors. "Exporting and serializing staged-out computations." JAX documentation. docs.jax.dev/...export Accessed 2026-07-28.
21. ^JAX authors. "Pallas: a JAX kernel language." JAX documentation. docs.jax.dev/...pallas Accessed 2026-07-28.
22. ^JAX authors. "Foreign function interface (FFI)." JAX documentation. docs.jax.dev/...ffi Accessed 2026-07-28.
23. ^JAX authors. "jax.extend module." JAX documentation. docs.jax.dev/...jax.extend Accessed 2026-07-28.
24. ^Budden, David, and Matteo Hessel. "Using JAX to accelerate our research." Google DeepMind, 2020-12-04. deepmind.google/...-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. jmlr.org/...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). papers.neurips.cc/...515e1679aca8cbc8033-Paper.pdf
27. ^Flax contributors. "Flax: Neural Networks for JAX." Flax documentation. flax.readthedocs.io/...stable Accessed 2026-07-28.
28. ^Optax contributors. "Optax." Optax documentation. optax.readthedocs.io/...latest Accessed 2026-07-28.
29. ^Orbax contributors. "Orbax." Orbax documentation. orbax.readthedocs.io/...latest Accessed 2026-07-28.
30. ^Grain contributors. "Grain: Feeding JAX Models." Grain documentation. google-grain.readthedocs.io Accessed 2026-07-28.
31. ^OpenXLA Project. "XLA." openxla.org/xla Accessed 2026-07-28.
32. ^OpenXLA Project. "StableHLO." openxla.org/stablehlo Accessed 2026-07-28.
33. ^JAX authors. "Persistent compilation cache." JAX documentation. docs.jax.dev/...persistent_compilation_cache Accessed 2026-07-28.
34. ^JAX authors. "Introduction to debugging." JAX documentation. docs.jax.dev/...debugging Accessed 2026-07-28.
35. ^JAX authors. "Default dtypes and the X64 flag." JAX documentation. docs.jax.dev/...default_dtypes Accessed 2026-07-28.
36. ^JAX authors. "Control flow and logical operators with JIT." JAX documentation. docs.jax.dev/...control-flow Accessed 2026-07-28.