jax

7dd13d70 - Rollback https://github.com/jax-ml/jax/pull/29410 due to downstream pytype failures.