One function can become a gradient, a batch, and compiled code
JAX starts with a NumPy-like array API, then changes what can be done with a numerical function. grad differentiates it, jit compiles it through XLA, and vmap maps it across array axes. These transformations compose. A researcher can write one loss function and derive compiled per-example gradients instead of maintaining separate scalar, batched, and differentiated versions. That is the reason to choose JAX over an ordinary array library.
The bargain has terms. JAX traces functions, so transformed code must be pure, and arrays do not accept NumPy-style in-place assignment. Python control flow that depends on traced values also needs JAX-aware forms. This is a programming model, not a drop-in speed switch for an existing scientific script. If mutation and data-dependent Python branches dominate your code, migration work can outweigh the compiler benefit.
The CPU install is simple, while hardware support is uneven
The README gives CPU users a one-command pip install. Accelerator paths split into CUDA 13, TPU, local ROCm 7, and Intel GPU packages or instructions. Its support table is unusually direct: native Windows supports CPU execution but not NVIDIA GPUs, WSL2 NVIDIA support is experimental, and Apple GPU support is experimental. A team should pick the target hardware before it treats installation as solved.
commit 656227d was a sizeable source checkout: 1,993 files, about 678,577 lines of source, and 54.8 MB. It contained 36 CI workflow files and a tests directory, but no Dockerfile. The missing container recipe does not block pip, though contributors who want a reproducible source environment must build that layer themselves. The public package path and the repository development path are different jobs.
What happened when we ran it
Our sandbox installed commit 656227d in 36 seconds, pulled 41 packages, and occupied 551 MB on disk. The build succeeded in 11 seconds. We used a fresh Debian container with Python 3.12, 3 CPUs, 8 GB of RAM, no secrets, and no elevated privileges. Pip-audit reported 0 known vulnerabilities in that installed dependency set.
Pytest failed with exit code 1 after 15 seconds. It recorded 0 passed, 0 failed, and 200 collection or setup errors out of 200 before stopping. The final lines name mgpu_matmul_test.py, mgpu_ragged_dot_test.py, mgpu_torch_test.py, mosaic_gpu_test.py, and several other Pallas modules. Those lines do not include the underlying exceptions.
That distinction matters. We did not watch 200 assertions fail, because pytest never reached a completed test. We also cannot call the checkout healthy under the full command we ran. The log does not prove that a GPU was required, that a system library was absent, or that JAX itself was broken. It proves a narrower and still useful point: the generic CPU container installed and built the repository, but it could not collect the full suite.
JIT changes Python semantics before it changes speed
The sharp-bits guide says JAX transformations are designed for functionally pure Python functions. Side effects may occur only while tracing, globals can be captured in surprising ways, and JAX arrays are immutable. Indexed changes use .at[...] operations rather than assignment. Dynamic shapes and boolean indexing also meet compiler constraints that ordinary NumPy code does not have.
Compilation brings its own measurement trap. Device execution can be asynchronous, and the first call includes tracing and compilation work that later calls may reuse. We did not benchmark either path, so the 11-second repository build says nothing about model training speed or kernel latency. Teams should measure their actual workload after forcing completion, on the same accelerator and software versions they intend to deploy.
A current CPU deadlock report deserves a targeted regression test
Open issue 41265 reports that JAX 0.11.1 and 0.11.2 can deadlock when several FFTs run concurrently inside a compiled while_loop or fori_loop on CPU. The report includes reproductions on 2-CPU and 32-thread systems and describes workarounds, with performance costs for one of them. It is a specific concurrency case, not evidence that ordinary FFT calls hang.
Release 0.11.2 shipped on September 17 with new primitives, Python 3.15 work, distributed mutual TLS settings, numerical fixes, and a move to Bazel 8.7.0 with Bzlmod for source builds. It also removed a deprecated API. That mix is normal for a fast-moving compiler-facing library: fixes arrive, and callers still need pinned versions plus migration review.
October 5 activity shows maintenance, not a small support surface
GitHub showed 36,377 stars, a push on October 5, 2026, and 2,640 open issues and pull requests combined. Several issues and pull requests were updated on October 4 and 5. The combined count is not a bug total. It reflects a broad surface that includes Python APIs, compilers, CPUs, several accelerator families, distributed execution, and experimental kernels.
JAX earns its complexity when function transformations are the architecture rather than a convenience. Our 36-second install and successful build make the CPU entry path credible. The 200 collection errors keep us from recommending the source checkout as a frictionless contributor setup. Adopt it with backend-specific tests, pinned jax and jaxlib versions, and a clear reason to prefer its transformation model.