safe_zip() error when passing arrays into args

Hi @patrick-kidger thank you for all your work. I have an ODE in JAX using scan that I am trying to port over to Diffrax to learn the library.

I have a vector field with 1 variable and defining my field / equinox call with ([self], t, y, args) and my return as new_y

x1 = Arr
x2 = Arr
assert len(x1) == len(x2)

t0 = 0,
t1 = len(x1) - 1

y0: float
args: tuple[Arr, Arr, float, float]

I am getting this error in runge_kutta.py

ks = ty_map(lambda x, xs: xs.at[stage_index].set(x), ki, ks)
....
  File "/Users/alessandrolatif/watz/venv/lib/python3.11/site-packages/jax/_src/numpy/util.py", line 398, in _broadcast_to
    for arr_d, shape_d in safe_zip(arr_shape, shape_tail))
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: safe_zip() argument 2 is shorter than argument 1

I am unsure what I am not understanding to get this to work and would be grateful for your guidance.

Start with the reported safe_zip() failure in runge_kutta.py and reproduce it using the JAX scan setup, the vector field call, and the array arguments described in the issue. Determine whether the failure is caused by Diffrax or by the supplied inputs, then document a minimal reproduction and verify the resulting behavior.