The Wire · Showcase
JAX FIXES REMAT3 CONSTANT HANDLING CRASH, GOOGLE CLOUD PYTHON CENTRALIZES REST TRANSCODING
By RepoJournal · Filed · About Google
Matthew Johnson patched a critical remat3 regression where closed-over constants were silently dropped, causing immediate crashes in production rematerialization pipelines [ref:1].
The remat3 flag flip exposed a subtle bug in _remat_jaxpr: to_jaxpr returns both jaxpr and constants, but the consts were discarded, leaving Jaxpr objects with constvars but no attached values [1] [2]. Any rematted computation capturing arrays (remat-of-jit patterns) would crash in pjit_staging_rule with mismatched input counts. Johnson's fix attaches constants via with_consts, mirroring how jaxpr constants are handled elsewhere. On the Google Cloud Python side, Heba Alazzeh shipped a centralized REST transcoding helper under google.api_core.gapic_v1.rest_transcoding [3] [4], collapsing redundant code from REST transport constructors and reducing serialization errors. Simultaneously, Mosaic GPU gained a layout inference rule for MemRefReshapeOp [5] and refactored GPU Interpret Mode to use GpuClockBundle for cleaner, less error-prone clock management [6]. Cloud Python also absorbed this week's API definition updates [7] and added agent identity credentials v1beta [8]. Across both repos, 27 commits and 23 PRs landed in the last 24 hours.
One email a day. Unsubscribe in one click.
Keep up with Google in about 3 minutes a day: what actually shipped — the commits, pull requests, releases, and security advisories that matter.
One email a day. Unsubscribe in one click. Read a past issue →
Action items
- → Test any remat3 pipelines that capture external arrays immediately google/jax [immediate]
- → Review REST transport code for integration with the new centralized transcode() helper googleapis/google-cloud-python [plan]
- → Monitor Mosaic GPU reshape handling for tiled layout changes google/jax [monitor]
References
- [1] [remat3] fix failures / regressions, work with mutable arrays google/jax
- [2] [remat3] fix failures / regressions, work with mutable arrays ↗ google/jax
- [3] feat(api-core): centralize rest transcoding helpers (#17765) googleapis/google-cloud-python
- [4] feat(api-core): centralize rest transcoding helpers ↗ googleapis/google-cloud-python
- [5] [Mosaic GPU] Add layout inference rule for `MemRefReshapeOp`. ↗ google/jax
- [6] [Mosaic GPU] Refactor GPU Interpret Mode to use `GpuClockBundle`s ↗ google/jax
- [7] feat: update googleapis and regenerate ↗ googleapis/google-cloud-python
- [8] feat(google/cloud/agentidentitycredentials/v1beta): add google-cloud-agentidentitycredentials v1beta ↗ googleapis/google-cloud-python