The Wire · Showcase
JAX FIXES CUSTOM_VJP3 AUTODIFF BUGS; PYTHON-GENAI BREAKS BACKWARD COMPATIBILITY
By RepoJournal · Filed · About Google
JAX shipped critical fixes to custom_vjp3 that were producing wrong-shaped outputs under vmap-of-jvp, while python-genai landed a breaking change removing three parameters and reshaping safety settings.
The JAX team patched custom_vjp3 to correctly handle symbolic-zero tangent avals in VmapOf.jvp [1], a bug that produced misshapen outputs when constants flowed through jvp-of-vmap combinations. The same PR also restored proper error messaging for jvp of symbolic_zeros=True custom_vjp, replacing bare NotImplementedError with the expected TypeError [2]. Separately, JAX's Pallas GPU support expanded with cp.async backing for copy_gmem_to_smem [3], controlled by a new impl parameter and supporting barriers for async waits. Over in python-genai, a major refactor [4] removes cached_content, presence_penalty, and frequency_penalty entirely while exposing safety_settings and labels without private SDK gates. This is a breaking change that will require client updates. The genai SDK also dropped the default api-revision header [5] to fix upstream issues. BigFrames 2.45.0 shipped with ai.classify, ai.score, and ai.if_ accessors for BigQuery DataFrames [6], plus local UDF execution support. Authentication landed new mTLS helpers [7] for custom connection pools to load default client certificates and resolve the GOOGLE_API_USE_MTLS_ENDPOINT environment variable.
Action items
- → Update python-genai clients to remove cached_content, presence_penalty, frequency_penalty usage before next release googleapis/python-genai [immediate]
- → Upgrade JAX if using vmap-of-jvp with custom_vjp3; verify symbolic-zero handling in your gradients google/jax [plan]
- → Review BigFrames 2.45.0 for new ai.classify and local UDF execution in your pipelines googleapis/google-cloud-python [monitor]
- → Test mTLS helpers in google.auth.transport.mtls if building custom connection pools googleapis/google-cloud-python [plan]
References
- [1] [hijax] custom_vjp3 fixes: VmapOf sym-zero jvp, errors, pretty-printing google/jax
- [2] [hijax] custom_vjp3 fixes: VmapOf sym-zero jvp, errors, pretty-printing ↗ google/jax
- [3] [pallas:mgpu] `plgpu.copy_gmem_to_smem` now supports `cp.async` ↗ google/jax
- [4] refactor(interactions)!: remove cached_content, presence_penalty, and frequency_penalty; expose safety_settings and labels ↗ googleapis/python-genai
- [5] chore: Remove default api-revision header googleapis/python-genai
- [6] bigframes: v2.45.0 ↗ googleapis/google-cloud-python
- [7] feat(auth): Implement python mtls helpers ↗ googleapis/google-cloud-python
FAQ
- What changed in Google on July 9, 2026?
- JAX shipped critical fixes to custom_vjp3 that were producing wrong-shaped outputs under vmap-of-jvp, while python-genai landed a breaking change removing three parameters and reshaping safety settings.
- What should Google teams do about it?
- Update python-genai clients to remove cached_content, presence_penalty, frequency_penalty usage before next release • Upgrade JAX if using vmap-of-jvp with custom_vjp3; verify symbolic-zero handling in your gradients • Review BigFrames 2.45.0 for new ai.classify and local UDF execution in your pipelines
- Which Google repositories shipped on July 9, 2026?
- google/jax, googleapis/python-genai, googleapis/google-cloud-python