Demos not working

Hi,
I can't run the demos presented in the getting started section: https://docs.kidger.site/diffrax/#quick-example
I just copied the code from the site and get an error.
I installed diffrax through conda and checked, that the versions are up to date:
Python - 3.11.8
diffrax - 0.5.1
jax - 0.4.27
equinox - 0.11.4
so I should fulfill the requirements listed.
There is a JAX issue on this topic, so I was wondering, if there is way around that.
What am I missing here?
Thanks.

Traceback (most recent call last):
File "/home/soeren-nagel/Projects/co-evo-social-dynamics/demos/test.py", line 10, in
solution = diffeqsolve(term, solver, t0=0, t1=1, dt0=0.1, y0=y0)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/equinox/_jit.py", line 206, in call
return self._call(False, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/equinox/_module.py", line 1053, in call
return self.func(self.self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/equinox/_jit.py", line 200, in _call
out = self._cached(dynamic_donate, dynamic_nodonate, static)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/diffrax/_integrate.py", line 993, in diffeqsolve
final_state, aux_stats = adjoint.loop(
^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/diffrax/_adjoint.py", line 292, in loop
final_state = self._loop(
^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/diffrax/_integrate.py", line 509, in loop
_, traced_jump, traced_result = eqx.filter_eval_shape(body_fun_aux, init_state)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/diffrax/_integrate.py", line 305, in body_fun_aux
(y, y_error, dense_info, solver_state, solver_result) = solver.step(
^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/diffrax/_solver/runge_kutta.py", line 1149, in step
final_val = eqxi.while_loop(
^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/equinox/internal/_loop/loop.py", line 107, in while_loop
return checkpointed_while_loop(
^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/equinox/internal/loop/checkpointed.py", line 249, in checkpointed_while_loop
final_val = _checkpointed_while_loop(
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/equinox/internal/_loop/checkpointed.py", line 270, in _checkpointed_while_loop
return while_loop(cond_fun, _body_fun, init_val)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/contextlib.py", line 81, in inner
return func(*args, **kwds)
^^^^^^^^^^^^^^^^^^^
File "/home/soeren-nagel/miniconda3/envs/social_space/lib/python3.11/site-packages/equinox/internal/_loop/checkpointed.py", line 268, in
_body_fun = lambda x: body_fun(x) # hashable wrapper; JAX issue #13554
^^^^^^^^^^^
ValueError: safe_map() argument 2 is shorter than argument 1

Start with the quick example in the getting started documentation and reproduce it using Python 3.11.8, diffrax 0.5.1, JAX 0.4.27, and Equinox 0.11.4. Trace the failure at diffeqsolve and the reported safe_map error, then verify that the example runs successfully with the documented requirements.