`AssertionError` is several tests

[GaetanLepage]

The following tests fail (distrax==0.1.5):

=========================== short test summary info ============================
FAILED distrax/_src/bijectors/tanh_test.py::TanhTest::test_stability__with_device - AssertionError: 
FAILED distrax/_src/bijectors/tanh_test.py::TanhTest::test_stability__with_jit - AssertionError: 
FAILED distrax/_src/bijectors/tanh_test.py::TanhTest::test_stability__without_device - AssertionError: 
FAILED distrax/_src/bijectors/tanh_test.py::TanhTest::test_stability__without_jit - AssertionError: 
FAILED distrax/_src/distributions/categorical_test.py::CategoricalTest::test_method_with_input_unnormalized_probs__with_device - AssertionError: 
FAILED distrax/_src/distributions/categorical_test.py::CategoricalTest::test_method_with_input_unnormalized_probs__with_jit - AssertionError: 
FAILED distrax/_src/distributions/categorical_test.py::CategoricalTest::test_method_with_input_unnormalized_probs__without_device - AssertionError: 
FAILED distrax/_src/distributions/categorical_test.py::CategoricalTest::test_method_with_input_unnormalized_probs__without_jit - AssertionError: 
FAILED distrax/_src/distributions/gamma_test.py::GammaTest::test_method_with_value_1d dist, 1d value, edge case__with_device - AssertionError: 
FAILED distrax/_src/distributions/gamma_test.py::GammaTest::test_method_with_value_1d dist, 1d value, edge case__with_jit - AssertionError: 
FAILED distrax/_src/distributions/gamma_test.py::GammaTest::test_method_with_value_1d dist, 1d value, edge case__without_device - AssertionError: 
FAILED distrax/_src/distributions/gamma_test.py::GammaTest::test_method_with_value_1d dist, 1d value, edge case__without_jit - AssertionError: 
FAILED distrax/_src/distributions/von_mises_test.py::VonMisesTest::test_von_mises_sample_gradient - AssertionError: 
FAILED distrax/_src/distributions/von_mises_test.py::VonMisesTest::test_von_mises_sample_moments - AssertionError: 
=== 14 failed, 8067 passed, 1444 skipped, 3049 warnings in 941.66s (0:15:41) ===

All of the errors are caused by some numerical results being too far from the expected values:

distrax/_src/utils/equivalence.py:91: in f
    np.testing.assert_allclose(x, y, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

args = (<function assert_allclose.<locals>.compare at 0x7ffa23d6b380>, array([-0.18232156, -0.69314718,        -inf]), array([ -0.18232156,  -0.69314718, -36.73680057]))
kwds = {'equal_nan': True, 'err_msg': '', 'header': 'Not equal to tolerance rtol=1e-07, atol=1e-05', 'verbose': True}

    @wraps(func)
    def inner(*args, **kwds):
        with self._recreate_cm():
>           return func(*args, **kwds)
E           AssertionError: 
E           Not equal to tolerance rtol=1e-07, atol=1e-05
E           
E           x and y -inf location mismatch:
E            x: array([-0.182322, -0.693147,      -inf])
E            y: array([ -0.182322,  -0.693147, -36.736801])