# JAX 0.12.0 Upgrade Guide: CPU Requirements, NumPy and MXFP4

JAX 0.12.0 changes x86 CPU and NumPy requirements while adding opt-in MXFP4 and remote execution, with practical checks for existing workloads.

By Paras Madan | 2026-10-19


JAX 0.12.0 requires AVX2 and FMA for its x86-64 binary packages and NumPy 2.2 or newer. Check your deployment machines and dependency pins before upgrading, then test the new MXFP4 and remote-runtime options separately from the existing workload.

The [stable package](https://github.com/jax-ml/jax/releases/tag/jax-v0.12.0) was published on October 19 and is available now. Its practical significance depends on where JAX runs: an older CPU can block the upgrade, a changed array contract can break application code, and low-precision operations need measurements on the intended accelerator.

## Check the machines behind the environment

The new x86 baseline is comparable to Intel Haswell-era hardware, with AVX2 and FMA instruction support. Older AVX-only processors need a source build. Check those capabilities on the actual host or virtual machine; a processor family name alone is a poor deployment check.

This matters beyond the main training server. Include CPU-only CI runners, preprocessing jobs and development machines in the inventory. A successful installation on a recent workstation gives you little evidence about an older worker that loads the same environment. Run a small compiled computation on each machine class before starting a long job.

The [0.12.0 package metadata](https://github.com/jax-ml/jax/blob/jax-v0.12.0/setup.py) sets NumPy's minimum to 2.2. Python still requires version 3.12 or newer, as it already did in 0.11.2. SciPy's minimum is 1.15 and `ml_dtypes` requires 0.5.0 or newer.

Rebuild the dependency lock in a clean test environment and check the resolved versions. Updating JAX while retaining a NumPy 2.1 constraint is an avoidable conflict. Keep the previous lockfile and deployment image available so that rollback restores the whole environment.

## Test array contracts before tuning performance

Two small API changes deserve targeted tests. Multi-input calls to `jax.numpy.atleast_1d`, `atleast_2d` and `atleast_3d` now return tuples rather than lists. A single input still returns an array. Check callers that append to the result or explicitly expect a list; ordinary unpacking can remain unchanged.

For `jax.random.multinomial`, `shape` now describes the batch dimensions. JAX appends the outcome dimension from the probability array. If existing code supplies the full desired output shape, remove its final dimension from that argument.

For example, ten samples over four outcomes should use a batch shape of `(10,)` and produce a `(10, 4)` result. Add an explicit shape assertion near sampling code, particularly where downstream logic can accept an extra dimension without immediately raising an error.

These checks belong in the baseline upgrade test. Keep the same data, model settings and precision while comparing outputs and gradients against the previous environment. That makes a later low-precision experiment easier to interpret.

## MXFP4 requires an explicit numerical choice

The new `jax.nn.get_scaled_dot_general_config('mxfp4')` option supplies four-bit E2M1 data with an E8M0 scale for each block of 32 elements. The [tagged implementation](https://github.com/jax-ml/jax/blob/jax-v0.12.0/jax/_src/nn/functions.py) rejects a `global_scale` argument for MXFP4, so an existing NVFP4 configuration should not be copied over unchanged.

Use the configuration with `jax.nn.scaled_dot_general`. Its configuration list covers the two operands and the gradient, and the documented example combines MXFP4 operands with an MXFP8 gradient. When no configurations are supplied, the function uses ordinary `lax.dot_general`.

Hardware support limits the performance story. The release specifies native quantization on AMD gfx950 GPUs and a dequantize-and-dot fallback elsewhere. Package support alone therefore gives no basis for predicting a speedup on a particular GPU.

Start with one representative operation and compare numerical error, forward and backward results, memory use and execution time. Measure compilation separately and wait for device work to finish before stopping the execution timer. Then test the surrounding model: an acceptable matrix-product error does not establish that training convergence or application quality will stay within your tolerances.

## Remote execution needs an existing runtime

JAX also registers an IFRT proxy backend. Set `JAX_PLATFORMS=proxy` and point `JAX_BACKEND_TARGET` at the server address to use devices managed by a remote IFRT runtime, such as Pathways on Cloud.

The [proxy implementation](https://github.com/jax-ml/jax/blob/jax-v0.12.0/jax/_src/xla_bridge.py) gives this backend a negative priority, keeping it from replacing a local backend automatically. A missing server address causes an error before an RPC is attempted. You need a separately available server and access to its compute resources.

Before moving a workload, confirm the selected backend and visible devices, then run a small computation through the intended connection. Include connection failures in the test so operational problems are recognizable. For teams using only local accelerators, this configuration can remain outside the initial upgrade.

Roll out the verified environment first. Introduce MXFP4 or proxy execution in separate changes, with a recorded baseline and a clear acceptance threshold. That keeps each result attributable to the dependency upgrade, numerical choice or execution location that produced it.

I'm Paras. I work at the intersection of AI, engineering and distribution.

I help technical teams turn infrastructure changes and migration guides into clear stories that reach the developers who need them.

[Book a 30-minute call](https://cal.com/paras-madan/30min).