Who contributed to Google on July 6, 2026?

1 developer shipped this update, including mattjj.

What were the notable Google updates?

[remat3] support offloading via save_and_offload_only_these_names, [remat3] support offloading via offload_dot_with_no_batch_dims, and [custom_vjp3] handle unhashable nondiff_argnums in custom_vjp3.

Google

JAX, the GenAI SDK, and the Cloud libs - Google's open source layer

Pick a date

Topics: AI / ML Python Full archive →

The Wire · Showcase

JAX REMAT3 OVERHAUL CUTS MEMORY PRESSURE WITH SELECTIVE OFFLOADING

By RepoJournal · Filed 06:02 UTC on July 6, 2026 · About Google

1 person shipped this

mattjj @mattjj 2 cited

JAX is shipping a fundamental redesign of its rematerialization system that lets you offload specific activations to slower storage during training, reclaiming GPU memory without the checkpointing tax.

The remat3 refactor hits production with two major pieces: support for save_and_offload_only_these_names [1] and offload_dot_with_no_batch_dims [2], both defunctionalizing previous approaches into dataclass policies. The first lets you name exactly which residuals get device_put to offload storage on the forward pass and restored on the backward pass. The second does the same for matrix multiplications without batch dimensions, offloading their outputs while rematerializing any batched dots as before. This keeps the remat2 functional interface intact [1] so your existing code doesn't break. Meanwhile, custom_vjp3 now handles unhashable nondiff_argnums [3], closing a gap that was blocking some VJP patterns. The wins here are real: you trade slow storage I/O for GPU memory headroom on large training runs without recomputation overhead.

Action items

→ Review remat3 offloading docs and test against your largest models google/jax [plan]
→ Check custom_vjp3 patterns if you're using unhashable arguments google/jax [monitor]

References

[1] [remat3] support offloading via save_and_offload_only_these_names google/jax
[2] [remat3] support offloading via offload_dot_with_no_batch_dims ↗ google/jax
[3] [custom_vjp3] handle unhashable nondiff_argnums in custom_vjp3 ↗ google/jax

Quick answers

What shipped in Google on July 6, 2026?: JAX is shipping a fundamental redesign of its rematerialization system that lets you offload specific activations to slower storage during training, reclaiming GPU memory without the checkpointing tax. In total, 7 commits and 4 pull requests landed.
Who contributed to Google on July 6, 2026?: 1 developer shipped this update, including mattjj.
What were the notable Google updates?: [remat3] support offloading via save_and_offload_only_these_names, [remat3] support offloading via offload_dot_with_no_batch_dims, and [custom_vjp3] handle unhashable nondiff_argnums in custom_vjp3.

JAX SYNCS XLA COMPILER STACK

JAX updated its XLA dependency to pull in the latest compiler improvements and optimizations.

ai 1 shipped 1-min read

Google 2 days ago

JAX OVERHAULS AUTODIFF INTERNALS WITH CUSTOM_VJP3, EXPANDS MESH API

JAX shipped a major redesign of custom differentiation rules that simplifies how you define vector-Jacobian products, while simultaneously opening up mesh operations to the public API.

ai 24 shipped 1-min read

Google 3 days ago

JAX FIXES REMAT3 CONSTANT HANDLING CRASH, GOOGLE CLOUD PYTHON CENTRALIZES REST TRANSCODING

Matthew Johnson patched a critical remat3 regression where closed-over constants were silently dropped, causing immediate crashes in production rematerialization pipelines [ref:1].

ai 50 shipped 1-min read

Google 4 days ago

JAX TIGHTENS XLA OPTIMIZATION CONTROLS, GOOGLE CLOUD PYTHON SETS UP FEATURE GATES

JAX makes XLA optimization levels human-readable while Google Cloud Python client libraries prepare a unified feature gating system for observability.

ai 54 shipped 1-min read

Elsewhere on the wire

AI Agents about 21 hours ago

CLINE SHIPS FIVE CRITICAL AGENT FIXES AFTER END-TO-END TESTING BLITZ

Saoudrizwan just landed a cascade of fixes that prevent the agent loop from running wild, stop file reads from resolving to the wrong directory, and plug the compaction pipeline.

ai-agents 16 shipped 2-min read

Local LLMs about 21 hours ago

LLAMA.CPP SHIPS MINIMAX-M3 SPARSE ATTENTION, VLLM UNLOCKS HETEROGENEOUS PREFILL-DECODE

The inference stack gets a new model architecture and a fundamental shift in how prefill and decode can run on different hardware configurations.

+10

llms 125 shipped 1-min read

@CachyOS about 21 hours ago

GLIBC 2.43 REBUILD SHIPS, INSTALLER JUMPS TO C++26

CachyOS pushed glibc 2.43 through the pipeline while the CLI installer upgraded its C++ standard to C++26, signaling a major refresh across core infrastructure.

infra 97 shipped 1-min read

Elixir & Phoenix about 21 hours ago

ECTO FIXES SUBQUERY PLANNING IN JOINS, ELIXIR PATCHES IEX PARSING

Ecto's query planner now handles subqueries correctly in interpolated join filters, while Elixir fixes a syntax error that broke 'not in' as a continuation operator in IEx.