jax

7e235e3a - jax.test_util: improve type annotations