Incompatible with jax v0.5.*

[rg936672]

What's the issue?

Updating to jax==0.5.* causes a number of test failures. Mostly these seem to be precision issues (arrays being equal to only four-ish decimal places, which doesn't pass the more demanding tests), but also some of the kernel gradients seem to be being flipped.

Remove the !=0.5.* specifier from the jax dependency in pyproject.toml once done.

[rg936672]

Partially resolved by #936 - but we might still want to investigate why the gradient was previously evaluating as 0.

[rg936672]

Some of the test failures may simply be caused by random seed dependence - jax made some breaking changes to its random number generation with the release of 0.5.0, which has caused many of our tests to run on new inputs.

[bk958178]

waiting until #979 is approved and merged