Jax release v0.3.1

• Changes:
  • jax.test_util.JaxTestCase and jax.test_util.JaxTestLoader are now deprecated.
    The suggested replacement is to use parametrized.TestCase directly. For tests that
    rely on custom asserts such as JaxTestCase.assertAllClose(), the suggested replacement
    is to use standard numpy testing utilities such as numpy.testing.assert_allclose(),
    which work directly with JAX arrays (#9620 ).
  • jax.test_util.JaxTestCase now sets jax_numpy_rank_promotion='raise' by default
    (#9562 ). To recover the previous behavior, use the new
    jax.test_util.with_config decorator:

    @jtu.with_config(jax_numpy_rank_promotion='allow')
    class MyTestCase(jtu.JaxTestCase):
      ...

  • Added jax.scipy.linalg.schur, jax.scipy.linalg.sqrtm,
    jax.scipy.signal.csd, jax.scipy.signal.stft,
    jax.scipy.signal.welch.