jax

8abb1a73 - Merge pull request #25490 from jax-ml:more-linearize-fixes