`AssertionError` is several tests

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 = (.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])
```