[charleswilmot]

Fixes #792 .

What. AbstractRungeKutta.step evaluates its stages in an eqxi.while_loop, the stage loop. On a GPU,
every check of a while loop's predicate is a host round trip: XLA copies the predicate to the host, and the
host waits for it. A Tsit5 step makes 7 of them, a Dopri8 step 14. This PR adds an opt-in for explicit tableaus:

solver = diffrax.Tsit5(scan_kind="unrolled")

Its step evaluates the stages one after another, with static stage indices, and no stage loop:

• Each stage is kept as its own array, and the stages are stacked once, at the end, for the error estimate and
  the dense output. Writing them into the [stages, ...] buffer instead, as the stage loop does, made the
  unrolled step 1.6x slower on the GPU (a lax.scan of step, 6 states: 38.4 ms against 23.7 ms). With
  static indices, XLA rebuilds the whole buffer at every write.
• An FSAL method's first stage, needed only on the first step and after a jump, is a lax.cond. So no step
  evaluates the vector field once more, and the profiler trace shows no host round trip for it.
• The default, scan_kind=None, keeps the stage loop. Nothing changes for anyone who does not ask for it.
• An implicit tableau with scan_kind="unrolled" raises, in __check_init__.
• Other adjoints. DirectAdjoint and ForwardMode replace a None scan kind, but not this one, and
  forward-mode autodiff goes straight through.

Results agree with the stage loop's to rounding, 4e-12 relative: the increments are summed in a different order.
An adaptive solve of a sensitive problem can therefore accept or reject a step differently and take a different
step sequence. Both results are then within the tolerance: in a batch of 1000 below, one solve in six did so, and
the largest error against a reference at rtol 1e-12 was the same for both, 2.2e-4.

Measured: batches of small systems. A small system belongs on a GPU only in a batch. Here it is the issue's
chain of 3 masses (6 states): Tsit5, jax.vmap of diffeqsolve over B initial states, 1000 constant steps.
Times are in ms, for the whole batch:

B	GPU: stage loop	GPU: unrolled	CPU: stage loop	CPU: unrolled
1	140.7	36.3	0.26	0.18
100	118.6	52.6	33.0	11.8
1000	122.6	59.4	332.6	95.4
10 000	398.3	224.8	2 644	619.3
100 000	3 385	1 962	109 669	15 214

B = 1 is a jax.vmap over a batch of one. An unbatched diffeqsolve takes 118.0 → 35.1 ms on the GPU, and
0.31 → 0.18 ms on the CPU.

• A single small solve belongs on the CPU: 0.18 ms, against 36 ms on the GPU.
• The GPU overtakes the CPU between B = 100 and 1000, with the stage loop on both devices and with unrolled
  stages on both. From B = 1000 on, unrolled is 1.7–2.1x faster on the GPU. Today's GPU stage loop is slower than
  the unrolled CPU until B = 10 000.
• On the CPU, batches gain far more than single solves: 2.8x at B = 100, 3.5x at 1000, 4.3x at 10 000 and
  7.2x at 100 000.
• With PIDController(rtol=atol=1e-6) the picture is the same. The GPU beats the CPU from B = 1000 in both
  variants, and there unrolled is 1.6–1.8x faster on the GPU. On the CPU, unrolled is 2.8–3.1x faster for
  B = 100–10 000.

Measured: host round trips per step, counted in a jax.profiler trace:

	diffeqsolve	lax.scan of step
Tsit5	8 → 1	7 → 0
Dopri8	15 → 1	–

The one left in diffeqsolve is its own loop over the steps.

Measured: single solves of various vector fields. Tsit5 and Dopri8, with constant steps and with
PIDController; solves, and jax.grad through the default RecursiveCheckpointAdjoint. The speed-up ranges over
solver and stepping:

vector field	GPU: solve	GPU: jax.grad	CPU: solve	CPU: jax.grad
the chain, 6 states	2.6–4.5x	2.7–3.8x	1.1–1.6x	1.3–1.8x
the chain, 600 states	2.0–3.0x	2.0–2.7x	0.89–0.97x	0.90–1.11x
a neural ODE, MLP 64 → 128 → 128 → 128 → 64	1.6–2.0x	1.9–2.1x	1.00–1.04x	1.14–1.25x
the same, width 1024	1.2x	1.3–1.5x	0.96–1.07x	1.16–1.28x
the chain, 4000 states (Tsit5, constant steps)	1.12x	0.95x	–	–

The cost is compile time. The vector field is traced once per stage instead of once. Across the cases above,
compile time grows 0.8–3.2x. benchmarks/compile_times.py (Dopri8, with a vector field built to be expensive to
trace) gives 0.8–1.9x with eqxi.noinline, and 4.3–13.5x without it. That cost, and single CPU solves up to
17 % slower at mid sizes (the chain, 60 states: 1.51 → 1.77 ms), are why this is opt-in.

Tests (test/test_solver.py, 22 new):

• scan_kind="unrolled" against the stage loop, for every explicit Runge-Kutta solver (Heun, Midpoint, Ralston,
  Bosh3, Tsit5, Dopri5, Dopri8). Each with constant steps and with adaptive steps across a jump, so that FSAL
  re-evaluates its first stage partway through. Checked: the saved values, through dense output, and the gradients.
• ForwardMode, DirectAdjoint, BacksolveAdjoint; vmap; an SDE (Heun on a MultiTerm).
• A step has no while in its jaxpr; an implicit solver given "unrolled" raises.

diffrax's whole suite passes, and passes again with every explicit Runge-Kutta step forced to
scan_kind="unrolled". The forcing is done by a local pytest plugin, which is not part of this PR.

jax / jaxlib 0.12.0.1 (jax[cuda12]), equinox 0.13.8, Python 3.14. NVIDIA RTX 4060 Laptop GPU, driver 595.99;
CPU: Intel i5-12500H, pinned to its four performance cores. Medians of 11 to 14 calls, compile excluded.