Tests crash/failures on `aarch64-linux` with latest jax(lib) (0.4.28)

[GaetanLepage]

When trying to run the test suite for jaxopt 0.8.3 with the latest jax/jaxlib installed (0.4.28), I face the following issues:

• The tests completely crash (core dumped style) if I don't disable the test_dtype_consistency.
• Also, after disabling the former test, there is still one failure (without a crash this time):
  FAILED tests/lbfgsb_test.py::LbfgsbTest::test_multiclass_logreg6 - AssertionError: Array(0.01411963, dtype=float32) not less than or equal to 0.01