The Wire · Showcase
JAX ADDS HARDWARE-NATIVE INTEGER MULTIPLY, TIGHTENS jax2tf SCOPE
By RepoJournal · Filed · About Google
JAX shipped a new mulhi primitive that maps directly to the hardware instruction, while cleaning up jax2tf to focus on what it actually supports.
JAX landed the mulhi primitive [1], a new operation that lowers directly to the chlo.mulhi op, giving you access to native high-bit multiplication on supported hardware without the overhead of decomposition. This is a precision win for numerical work where you need both halves of the multiply result. In parallel, the jax2tf toolchain got a documentation and examples refresh [2] that strips out TensorFlow.js references and removes outdated serialization patterns, making it clearer what the bridge actually supports and reducing confusion for new users. The update doesn't break existing code but does expect you to drop any TF.js assumptions if you had them. On the testing front [3], JAX added hardening for remat transpose when primal inputs are reduced, ensuring cotangent handling stays correct through that optimization path.
Action items
- → Review jax2tf examples if you're using the bridge, especially TF.js paths jax-ml/jax [plan]
- → Test mulhi on your hardware target if you're doing integer-heavy numerics jax-ml/jax [monitor]
References
FAQ
- What changed in Google on May 24, 2026?
- JAX shipped a new mulhi primitive that maps directly to the hardware instruction, while cleaning up jax2tf to focus on what it actually supports.
- What should Google teams do about it?
- Review jax2tf examples if you're using the bridge, especially TF.js paths • Test mulhi on your hardware target if you're doing integer-heavy numerics
- Which Google repositories shipped on May 24, 2026?
- google/jax