Tests fail with: jax.errors.UnexpectedTracerError: Encountered an unexpected tracer

[GaetanLepage]

Version of Awkward Array

2.7.2

Description and code to reproduce

When running the test suite, I get the following failures:

FAILED tests/test_1447_jax_autodiff_slices_ufuncs.py::test_regular_array_2 - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1447_jax_autodiff_slices_ufuncs.py::test_regular_array_4 - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1447_jax_autodiff_slices_ufuncs.py::test_regular_array_5 - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1447_jax_autodiff_slices_ufuncs.py::test_recordarray_4 - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1447_jax_autodiff_slices_ufuncs.py::test_recordarray_5 - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1447_jax_autodiff_slices_ufuncs.py::test_recordarray_6 - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1490_jax_reducers_combinations.py::test_reducer[min-0] - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1490_jax_reducers_combinations.py::test_reducer[min-1] - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1490_jax_reducers_combinations.py::test_reducer[max-0] - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
FAILED tests/test_1490_jax_reducers_combinations.py::test_reducer[max-1] - jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A funct...
================= 10 failed, 2957 passed, 36 skipped in 44.86s =================

Here is an example of the error message I get each time:

    def check_eval_args(args):
      for arg in args:
        if isinstance(arg, Tracer):
>         raise escaped_tracer_error(arg)
E         jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with>
E         JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
E         To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
E         See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.UnexpectedTracerError

FYI, I'm running the latest jax version (0.4.38).