JAX bug in field line integrate when using JAX version `0.6.0`

[unalmis]

Looks like #1839 caused a regression. Seems to have not been caught since the PR was merged without running those JAX tests.

216 durations < 0.005s hidden.  Use -vv to show these durations.)
=========================== short test summary info ============================
FAILED tests/test_magnetic_fields.py::TestMagneticFields::test_field_line_integrate_jax_transforms - TypeError: Argument 'Zero(ShapedArray(float0[1]))' of type '<class 'jax._src.ad_util.Zero'>' is not a valid JAX type
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[f0uriest]

which version of jax/diffrax etc was this using?

[unalmis]

https://github.com/PlasmaControl/DESC/actions/runs/18185201751

[YigitElma]

patrick-kidger/diffrax#622 (comment)

Let me try to change the test to prevent it.