The Wire · Showcase
JAX OVERHAULS REF SYSTEM AS GENAI SDK LOCKS IN JSON SCHEMA SUPPORT
By RepoJournal · Filed · About Google
JAX's reference layer just got a major surgery: pin control arguments are in, version guards are out, and Pallas GPU kernels gained conditional execution.
The JAX team shipped a significant refactoring of its reference system [1], adding pin=True/False arguments to new_ref and deleting unnecessary config overhead. This cleanup is paired with removal of redundant ifrt_version guards [2], signaling a maturing platform moving away from version-specific workarounds. On the GPU side, Pallas for Mosaic added predicate arguments to copy_gmem_to_smem [3], enabling conditional kernel execution and updating the expect_tx dialect signature from I32 attr to I32. The compiler desk cleaned up redundant constraint canonicalization [4] and hardened the leak-checker in lower_jaxpr [5], stabilizing the entire compilation pipeline. Over at the GenAI SDK, JSON schema support landed in FunctionDeclaration parsing [6], while LiveServerContent gained interaction_status tracking [ref:7, ref:9] to better monitor real-time agent behavior. On the cloud Python side, BigQuery storage validation tests now cover pandas-gbq delegation [9] with real protobuf decoding across multiple data types, and AsyncAppendableObjectWriter finally has retry support for close() and finalize() operations [11], plugging a gap in transient error recovery. Generator compatibility imports [10] added smart fallbacks for google-api-core helpers, ensuring older versions don't break. Across all three repos: 41 commits, 34 PRs.
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
- → Review JAX ref system changes before pinning your jax version in production configs google/jax [plan]
- → Test AsyncAppendableObjectWriter retry behavior in staging if you use gRPC append streams googleapis/google-cloud-python [plan]
- → Integrate JSON schema support into FunctionDeclaration calls for stricter type validation googleapis/python-genai [monitor]
References
- [1] [refs] add pin=True/False arg to new_ref, delete config ↗ google/jax
- [2] Remove the ifrt_version guard that's unnecessary google/jax
- [3] [Pallas][Mosaic GPU] Added predicate arg to copy_gmem_to_smem ↗ google/jax
- [4] [NFC] Remove redundant `AlwaysTrue` canonicalization of `OneOf` constraints. google/jax
- [5] make leak-checker more robust in lower_jaxpr ↗ google/jax
- [6] feat: enable json schema in FunctionDeclaration parser ↗ googleapis/python-genai
- [7] feat: Add interaction_status to LiveServerContent googleapis/python-genai
- [8] feat: Add interaction_status to LiveServerContent ↗ googleapis/python-genai
- [9] test: add dual-environment validation tests for pandas-gbq delegation ↗ googleapis/google-cloud-python
- [10] feat(generator): add fallback compatibility imports for google-api-core helpers ↗ googleapis/google-cloud-python
- [11] fix(storage): add retry support for finalize and close in AsyncAppendableObjectWriter ↗ googleapis/google-cloud-python