jax

cd1b9520 - Merge pull request #29937 from jax-ml:avoid-strong-refs