Three composable transformations are JAX's reason to exist
JAX begins with NumPy-like arrays, then lets you transform a numerical function. grad differentiates it, jit compiles it through XLA, and vmap maps it across array axes without making the author carry batch dimensions through every operation. Those transformations can wrap one another. A researcher can write one loss function and derive compiled per-example gradients without maintaining a separate implementation.
CPU installation is one command, while accelerators split by platform
The documented CPU start is pip install -U jax. Accelerator users choose separate extras for CUDA 13, TPU, or local ROCm 7, while Intel GPU users follow another project's instructions. The support table is frank about limits: native Windows supports CPU use but no NVIDIA GPU path, WSL2 GPU support is experimental, and Apple GPU support is also experimental.
Our checkout contained 1,968 files, about 651,242 source lines, and 54.3 MB before installation. It also had 37 CI workflow files and a tests directory, but no Dockerfile. That is substantial compiler-facing infrastructure rather than a small Python utility. The missing Dockerfile does not prevent installation, though teams wanting a reproducible source-development image must supply that layer themselves.
What happened when we ran it
Our sandbox installed the checked-out commit in 26 seconds. It added 39 Python packages and occupied 546 MB on disk. The build completed successfully in another 31 seconds. Pip-audit reported 0 known vulnerabilities in the installed dependency set, which is a clean result for that specific environment and commit rather than a guarantee about later releases.
The test command failed with exit code 1 after 73 seconds. Pytest reported 0 passed, 0 failed, and 200 collection or setup errors out of 200 before its stop threshold ended the run. The final log lines list Pallas modules, including TPU all-gather, fusible matmul, paged attention, and asynchronous tests. Those lines do not state the underlying exception, so blaming a missing package, unsupported CPU, or test bug would be guesswork.
Two hundred collection errors block a clean source-checkout verdict
No test body completed in our run because pytest failed during collection or setup. That distinction matters: we did not observe 200 behavioral assertions fail, and we also cannot claim the suite passed any exercised functionality. The installation and build results show that the repository can be prepared and compiled in the stated Debian container. They do not establish that this checkout is healthy under its complete test command.
The mismatch between a 31-second successful build and 200 setup-stage errors is useful purchasing information. A developer evaluating only the CPU quick start may never touch Pallas or TPU modules, yet a contributor or team packaging from source needs a documented way to collect the suite in a clean environment. The log tail alone cannot tell us which additional setup, if any, would have changed the result.
JIT requires pure functions and predictable shapes
Compilation changes how ordinary Python behaves. JAX traces a function to build a computation, so data-dependent branches, mutation, iterators, side effects, and changing shapes need care. The first call pays tracing and compilation costs that later calls may reuse. Device execution can also be asynchronous, which means a casual timer can measure dispatch rather than completed work.
Three scaling modes expose how work reaches devices
JAX documents automatic, explicit, and manual parallelism. Automatic mode keeps a global view while the compiler selects much of the partitioning. Explicit sharding records layouts in array types and keeps them inspectable. Manual mode gives the programmer a per-device view and explicit collectives. This progression is useful when a workload outgrows one device and compiler choices need tighter control.
Hardware portability still needs proof on the intended machine. A function that compiles for 1 CPU may hit different lowering, memory, or precision behavior on a GPU or TPU. Pallas goes further by supporting custom accelerator kernels, and several of the modules named in our 200-error log tail sit in that area. Teams using those paths should keep backend-specific tests instead of treating XLA as a uniform abstraction.
An open float16 report makes numerical regression tests necessary
Issue 39899 reports that jnp.linalg.norm can narrow a float16 intermediate before the square root. Its reproducer returns infinity for a representable norm and zero gradients for the example, with no invalid value in the gradient to alert downstream code. The report remained open when checked. It concerns a specific dtype and operation, but silent numerical errors deserve more weight than an ordinary exception.
Release 0.11.1 shows active repair alongside migration cost. It fixed numerical instability in 2x2 and 3x3 determinants, fused-attention batching, split behavior, sharded reshapes, and diagnostics. The same release removed 2 configuration flags, changed several NumPy-style return containers, and ended deserialization support for exports older than the stated compatibility window. Pinning jax with jaxlib and reading release notes is routine maintenance here.
Same-day activity supports JAX, while the queue stays large
GitHub showed a push on August 26, 2026, plus newly updated issue activity that day. The repository had 36,213 stars and 2,487 open issues and pull requests when fetched. That combined count is not a bug total. It reflects a wide surface spanning Python APIs, compilers, device backends, distributed execution, and experimental kernels, and it makes issue triage quality more informative than the raw number alone.
JAX 0.11.1 arrived on August 17, 2026, 9 days before that push. The current activity and detailed release notes support confidence that the project is maintained. Our failed collection run prevents an equally confident claim about building from this source checkout in a generic Debian container. Adopt JAX for its transformation model, then make backend, dtype, export compatibility, and test-environment checks part of the engineering work.