JAX 0.11.0 Makes AI Training Memory Policy Programmable

A training team can now tell JAX which activations to save, move to host memory, or recompute at a function boundary. JAX 0.11.0 turns that choice into a programmable policy for AI workloads. It does not turn memory pressure into a solved problem.

The JAX 0.11.0 release introduces custom_remat, named policy classes for saving, offloading and recomputation, and an option to inline jitted functions into their callers. The same release raises the supported runtime floor and changes the behavior of empty array creation. Those details make this more than a feature announcement. It is a migration decision.

The useful change is not that rematerialization exists. JAX already let developers trade computation for memory. The change is that a function can carry its own forward, rematerialization and backward behavior while still seeing the checkpoint policy around it. That gives framework and model authors a more precise place to express memory intent. It also puts more responsibility on them to prove that the policy helps on the graph, compiler, accelerator and interconnect they actually run.

What changed inside JAX rematerialization

Automatic differentiation needs intermediate values from the forward pass to calculate gradients. Keeping every activation is fast to retrieve but expensive in accelerator memory. Discarding selected activations and rebuilding them during the backward pass lowers storage pressure at the cost of extra computation. JAX calls that trade rematerialization, and its gradient-checkpointing guide explains how policies decide which values remain available.

JAX 0.11.0 adds custom_remat for cases where a plain wrapper around a function is too blunt. Its forward rule can produce the normal output plus residual state. Its rematerialization rule can define what gets rebuilt. Its backward rule can consume the residuals and cotangents. The initial remat3 implementation also shows why policy visibility matters: the custom function can inspect the ambient checkpoint policy instead of operating as an isolated black box.

That is a better fit for modern AI systems than a single global switch. Large training graphs mix attention, expert routing, normalization, communication and custom differentiation. Their tensors have different sizes, reuse patterns and reconstruction costs. A policy that is sensible for a large attention activation can be wasteful for a small value that is expensive to reproduce.

TECHi’s recent analysis of Qwen3.5 on Google Ironwood showed the same systems principle from the inference side: model architecture, kernels, memory layout and communication have to be tuned together. That article mentions JAX only as the language behind Pallas kernels. The 0.11.0 release addresses a different question—how training code expresses what should survive the forward pass.

Save, offload and recompute are different bets

The new named checkpoint-policy classes make three choices easier to reason about.

Saving keeps an activation close to the accelerator and avoids rebuilding it, but consumes scarce high-bandwidth memory. Recomputing frees that storage, then spends accelerator cycles to reproduce the value later. Offloading moves a saved value to another supported memory kind, such as host memory, and pays transfer latency plus bandwidth to bring it back.

Those choices are not interchangeable. An offload policy can help when a large tensor is cheap to transfer relative to the memory it releases and when copies overlap with useful work. It can hurt when the tensor is small, the interconnect is congested, the backend does not support the requested memory kind, or the backward pass needs the value before the transfer completes. JAX’s host-offloading notebook makes the constraint explicit: device placement and memory-kind support are part of the policy.

Recomputation has a similar boundary. A cheap elementwise result may be a good candidate to rebuild. A collective communication step or a costly custom kernel may not be. The right objective is not the smallest possible peak-memory number. It is the best capacity and step-time result for the workload’s real training target.

That distinction matters for models whose published weights already imply a steep hardware floor. Inkling’s open weights still need a data center because only a fraction of a mixture-of-experts model may activate per token while the full parameter set still has to live somewhere. Activation policy does not erase weight memory. It can, however, decide whether intermediate state prevents a model, sequence length or batch from fitting within the remaining budget.

The new path is deliberately opt-in

customremat depends on the experimental remat3 implementation. The jaxremat3 configuration flag is false by default in the 0.11.0 source, so upgrading the package does not silently move existing workloads onto the new machinery. Teams must enable it intentionally.

That is the right posture for a transformation that touches forward and backward execution. JAX development leading into 0.11.0 included fixes involving scan, custom VJP and JVP behavior, higher-order automatic differentiation and symbolic-zero cotangents. A training stack that combines those transforms should not infer safety from a simple example. It needs tests shaped like its own graph.

The JAX configuration reference is therefore operational documentation, not a footnote. A rollout should record the flag state with the package version and accelerator environment. Otherwise, two jobs described as running JAX 0.11.0 can execute through different rematerialization implementations.

The dependency floor changes at the same time

The PyPI package metadata requires Python 3.12 or newer. The release also drops support for NumPy 2.0 and SciPy 1.14, establishing NumPy 2.1 and SciPy 1.15 as the new minimums. Python 3.13t support is removed as well.

That package floor can be a larger migration than custom_remat itself. Notebook images, training containers, cluster base layers and compiled extensions may all move together. An environment resolver can produce a valid installation while application code or a binary dependency still assumes the older stack.

TECHi’s examination of Dependabot’s install-layer gap is relevant here. A controlled update proposal is not the same thing as a controlled runtime change. JAX teams need a pinned environment, a reproducible lockfile or image digest, and a clean rollback target. The version number alone does not describe the deployed system.

JAX 0.11.0 also changes jax.numpy.empty and empty_like so they return genuinely uninitialized arrays instead of zero-initialized storage. That aligns the name with NumPy’s documented empty semantics, but it can expose code that accidentally relied on the old zeros. The risk is not necessarily a loud failure. A buffer that is partially written before being read can now carry arbitrary prior memory values into a result.

This is a correctness audit, not just a performance benchmark. Search for empty and empty_like allocations, prove that every element is written before use, and add a test that would fail if stale values survive. A speedup or lower memory peak does not compensate for silent numerical contamination.

One release-note claim still needs a clean runtime check

The release notes say jax.checkpoint_policies is now a submodule and that importing names from jax.checkpoint_policies should work. TECHi inspected the tagged source and the uploaded PyPI wheel before publication. The archive contains no jax/checkpoint_policies.py module. In the tagged ad_checkpoint source, checkpoint_policies is still built as a types.SimpleNamespace.

Attribute access through jax.checkpoint_policies is present in the inspected code. The new submodule import path is not something we would treat as verified without a clean-wheel runtime test or a follow-up package. This may be packaging behavior that is clarified quickly, but it is exactly the kind of edge that an upgrade can expose in framework code.

The practical response is narrow: keep imports on the documented attribute path that already works in the environment, add an import test to continuous integration, and watch the JAX release record for clarification. Do not make the entire migration depend on a newly described import form until the shipped artifact proves it.

Measure memory and step time together

A useful rematerialization trial starts with the workload that is currently constrained. Record peak accelerator memory, compiled step time, steady-state step time, host-memory use, transfer volume and any change in batch size or sequence length. Warm-up and compilation should be separated from steady-state execution.

Then change one policy boundary at a time. A global reduction in peak memory can conceal a local regression if a large activation is repeatedly transferred or an expensive operation is rebuilt. Per-function custom_remat is valuable precisely because it gives teams a smaller unit to test.

The result also has to survive the transformations used in production. Test scan loops, custom differentiation, distributed sharding and higher-order gradients where they exist. Compare outputs and gradients with the prior implementation at tolerances appropriate for the model. A graph that compiles and completes one step has not yet proved training stability.

What would invalidate the case for adopting it

The feature earns its place if it unlocks a material capacity gain, keeps numerical behavior stable and produces an acceptable end-to-end time or cost trade. The case weakens when offload traffic stalls the step, recomputation consumes the saved capacity benefit, backend support narrows portability, or the experimental path behaves differently under the transforms a model depends on.

The release is still important even for teams that wait. It exposes memory policy at a more useful software boundary and gives framework authors a way to encode model-specific choices. That is a real advance in the control surface around AI training.

But JAX 0.11.0 is not a free-memory button. It is a sharper instrument arriving alongside a higher platform floor and a changed allocation contract. Teams should adopt it the way they would adopt a compiler or runtime change: with artifact-level verification, workload-shaped tests and a rollback that has already been rehearsed.