The Wire · Showcase
JAX FIXES FP8 ATTENTION, LANDS TPU CUMULATIVE REDUCTIONS
By RepoJournal · Filed · About Google
JAX shipped critical fixes for fp8 batching rules and native TPU support for cumulative operations, while the Gemini SDK crossed v2.16.0 with new robotics models and expanded grounding capabilities.
The JAX team merged a follow-up fix for fp8 attention batching that applies the same normalization logic as the earlier non-fp8 correction [1], addressing a root cause that affected multiple code paths. Simultaneously, cumulative reductions (cumsum, cumprod, cummax, cummin) now lower to native chlo.ScanOp on TPU [2], mirroring GPU behavior while preserving backward compatibility through the TPU compiler's rewrite pipeline. On the SDK side, python-genai v2.16.0 landed with environment resource support, native GoogleMaps Tool grounding types, and the Gemini Robotics ER 2 Preview model [3], expanding multimodal capabilities for production workloads. The google-cloud-python ecosystem saw Protobuf 6.33.5 become a hard requirement [6] to unblock client library modernization, while SQLAlchemy-Spanner gained native UUID data type support with opt-in backward compatibility [5]. One caveat: google-cloud-bigquery was temporarily dropped from bulk release due to a socket leak in the dbapi connection tests [4], requiring isolated investigation before reintegration.
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
- → Pull JAX fp8 and TPU cumulative fixes immediately if running fp8 attention or cumulative reductions on TPU google/jax [immediate]
- → Upgrade python-genai to v2.16.0 and integrate Gemini Robotics ER 2 if deploying multimodal applications googleapis/python-genai [plan]
- → Enforce Protobuf 6.33.5+ in google-cloud-python dependencies before next release cycle googleapis/google-cloud-python [plan]
- → Hold google-cloud-bigquery deployments until socket leak fix lands; monitor PR for resolution googleapis/google-cloud-python [monitor]
References
- [1] Fix cudnn fp8 attention batching rules ↗ google/jax
- [2] [JAX] Lower cumulative reductions to chlo.ScanOp on TPU. ↗ google/jax
- [3] v2.16.0 ↗ googleapis/python-genai
- [4] chore: remove google-cloud-bigquery from bulk release due to failing system test (#17944) googleapis/google-cloud-python
- [5] feat(sqlalchemy-spanner): native UUID data type support in dialect ↗ googleapis/google-cloud-python
- [6] fix: require Protobuf 6.33.5+ ↗ googleapis/google-cloud-python