test_chees_adaptation fails with latest jax 0.7.0 #795
Description
Describe the issue as clearly as possible:
When updating jax to 0.7.0, tests/adaptation/test_adaptation.py::test_chees_adaptation[adaptation_filters0] fails with AssertionError:
np.testing.assert_allclose(harmonic_mean, 0.75, atol=1e-1)
np.testing.assert_allclose(parameters["step_size"], 1.5, rtol=2e-1)
> np.testing.assert_array_less(infos.num_integration_steps.mean(), 15.0)
E AssertionError:
E Arrays are not strictly ordered `x < y`
E
E Mismatched elements: 1 / 1 (100%)
E Max absolute difference among violations: 0.16400051
E Max relative difference among violations: 0.01093337
E x: array(15.164001, dtype=float32)
E y: array(15.)
tests/adaptation/test_adaptation.py:124: AssertionError
Do you think the test is just flaky?
Context: bumping jax in nixpkgs
Steps/code to reproduce the bug:
pytestExpected result:
All tests succeedError message:
Blackjax/JAX/jaxlib/Python version information:
BlackJAX 1.2.5
Python 3.13.5 (main, Jun 11 2025, 15:36:57) [GCC 14.3.0]
Jax 0.7.0
Jaxlib 0.7.0Context for the issue:
No response