RepoJournal
Google

Google

JAX, the GenAI SDK, and the Cloud libs — Google's open source layer

Pick a date

  1. Wed 20 may
  2. Thu 21 may
  3. Sat 23 may
  4. Sun 24 may
  5. Mon 25 may
  6. Tue 26 may
  7. Thu 28 may
  8. Fri 29 may
  9. Sat 30 may
  10. Sun 31 may
  11. Mon 1 jun
  12. Mon 8 jun
  13. Tue 9 jun
  14. Wed 10 jun
  15. Thu 11 jun
  16. Fri 12 jun
  17. Sat 13 jun
  18. Sun 14 jun
  19. Mon 15 jun
  20. Tue 16 jun
  21. Wed 17 jun
  22. Thu 18 jun
  23. Fri 19 jun
  24. Sat 27 jun
  25. Sun 28 jun
  26. Mon 29 jun
  27. Tue 30 jun
  28. Wed 1 jul
  29. Thu 2 jul
  30. Today 3 jul

The Wire · Showcase

JAX TIGHTENS GRADIENT CHECKPOINTING, FIXES SCIPY PARITY BUGS

By RepoJournal · Filed · About Google

JAX's custom_vjp3 primitive now supports checkpoint_name in forward passes, unlocking memory-efficient gradient computation for complex AD workflows.

The big move is giving custom_vjp3 a remat rule that applies JAX's remat_transform to primal functions [1]. This means you can now checkpoint intermediate values in custom VJP forward passes without breaking the AD chain, which matters for training large models where memory is the bottleneck. Alongside this, JAX made primal_left_tangent_right a hijax primitive so it lowers away in lojax, letting the compiler dead-code-eliminate unused gradients [1].

The team also shipped fixes to scipy.special parity bugs that have been lingering [2] [3]. lax.special.gammainc and lax.special.gammaincc now return 0.0 and 1.0 respectively when a=infinity instead of NaN, matching scipy's behavior. Same fix went into igamma and igammac for a=inf with finite x. These are edge cases but they matter if you're doing numerical computation or porting scipy code to JAX.

Test coverage for remat3 got a refresh to close gaps [4], so expect fewer surprises in production rematerialization rules. Custom AD primitives and scipy compatibility are both stability plays, not headline features, but both are the kind of fixes that prevent silent correctness bugs.

Action items

References

  1. [1] support checkpoint_name in custom_vjp fwd if custom_vjp3=1 ↗ google/jax
  2. [2] Fix `lax.special.gammainc` and `lax.special.gammaincc` for a = infinity ↗ google/jax
  3. [3] Fix `lax.special.igamma` and `lax.special.igammac` for `a = inf` & `x` finite. google/jax
  4. [4] [remat3] update tests to finish coverage of remat3 ↗ google/jax

FAQ

What changed in Google on June 28, 2026?
JAX's custom_vjp3 primitive now supports checkpoint_name in forward passes, unlocking memory-efficient gradient computation for complex AD workflows.
What should Google teams do about it?
Review custom_vjp forward passes for checkpoint_name opportunities if you're memory-constrained on large models • Test scipy.special edge cases (gammainc, gammaincc, igamma, igammac) if you rely on infinity bounds
Which Google repositories shipped on June 28, 2026?
google/jax

Related across the cluster

For your repos

The showcase is a teaser.
Your wire is the product.

Same engine. Different stack. Below: what changes when the wire is yours.

Showcase wire

  • •14 famous open source orgs
  • •One wire per day
  • •Public, generic
  • •Read on the web, when you remember

Your wire

  • →Up to 1,500 of your repos - orgs, deps, vendors
  • →Morning and evening briefs
  • →Action items routed to your team
  • →Slack delivery, email, breaking-news CVE alerts

Want a hands-on demo first? Ask a current user for an invite link.