jax

2616b4d3 - Loosen test tolerance for `tests/linalg_sharding_test.py::LinalgShardingTest::test_batch_axis_sharding_jvp`.