JAX ships 0.11.1 and reworks custom_vjp, Cloud Python hardens Spanner

JAX spent August reworking its VJP system, culminating in 0.11.1. The core change adds `defvjp_with_logs` so `custom_vjp` rules can emit logs, a new primitive that keeps call sites compatible. The team also caught itself: a plan to make `fun` in `custom_gradient` positional-only was reverted because "that's a breaking change" [1], blocking a silent API break for existing callers. The cumulative effect is a more instrumentable autodiff core without forcing migration on anyone.

Beyond the headline work, the month showed a pattern of restrained reversals and cleanups. A TPU accumulator API was rolled back, a `shard_alike` feature was deleted, and legacy call machinery was stripped. Each of these reduces the surface area JAX has to support, steering the design toward a smaller, more intentional core. The 0.11.1 release also shipped "export guardrails" and a tracing speedup, though the exact guardrails were left unspecified in the daily briefings.

Google Cloud Python's standout story is the Spanner `JsonObject` fix, which implements the dict protocol and nested unwrapping. This resolves a class of bugs where JSON values weren't treated like dictionaries, which previously forced users to write manual extraction code [2]. Spanner also fixed a data-loss bug in autocommit mode, and Bigtable saw a data-loss patch, both of which are safety-critical for production code paths.

On the build and tooling front, JAX switched to bzlmod and dropped Python 3.11 and legacy Bazel in one morning. This is a deliberate modernization push, but it means downstream users still on Python 3.11 will need to plan an upgrade before adopting newer JAX versions. The month closed with the XLA dependency updated at least twice, a routine sync that nonetheless keeps JAX aligned with upstream.

Finally, the hardware-specific work deserves notice: Intel GPU support landed in CI, native TPU convolutions shipped, and Mosaic GPU gained efficient transposed 8-bit stores using `ld/stmatrix` for Hopper+ and `cp.async` on older architectures [3]. These are the kind of low-level kernel decisions that determine whether the high-level API can actually deliver performance on next-gen hardware.