Fix tests with jax 0.5.1

[stephen-huan]

Checklist

• I've formatted the new code by running hatch run dev:format before committing.
• I've added tests for new code.
• I've added docstrings for the new code.

Description

Fix some failing tests with jax 0.5.1. It seems from 0.5.0 -> 0.5.1 the numerics got slightly worse (especially for inverse).

For the get_batch test, as I understand since it's sampling with replacement it's possible to get unlucky and have two rows coincide. This is most likely when the batch size is large (the tests that fail are when the batch size is the same as the dataset size). We'd expect n_dim * batch_size / n_data elements to collide (each row has a 1 / n_data chance of colliding, each collision makes n_dim elements the same, and there are batch_size rows).

Related to #490.

[stephen-huan]

CI failures seem unrelated.

[thomaspinder]

Yes, unclear what is causing the CI failures right now. I am unable to reproduce locally.