`jax.jacfwd` and `jax.jacrev(jax.grad(...))` return NaN through `experimental.run_loop_jit`, while `jax.grad` works fine

## Summary

`jax.grad` works correctly through `torax.experimental.run_loop_jit` — this
is the pattern documented in
`torax/examples/iter_hybrid_rampup_grad_and_vmap.ipynb`, and it's worked
reliably for us across many production runs. But when we tried to go one
step further and compute a Hessian (for a local optimum
identifiability/sensitivity check) or a first-order Jacobian of a residual
vector via forward-mode, both failed:

- `jax.jacfwd` (first-order, forward-mode) returns **NaN**.
- `jax.jacrev(jax.grad(...))` (second-order, reverse-over-reverse — the
usual way to get a Hessian without needing forward-mode at all) also
returns **NaN**.

Both fail on a minimal, fully public config attached below: circular
geometry (no external file), the transport/source defaults from your own
`examples/basic_config.py`, and a standard fixed-dt linear solver
(predictor-corrector + Pereverzev-Corrigan stabilization — not an exotic
setup). No custom transport models, no MHD/sawtooth model, no
`jax.checkpoint`. **Only 2 real timesteps are needed to reproduce both
failures.**

## How to reproduce

```
python3 repro.py
```

The necessary script file is attached.

## Expected behavior

`jax.jacfwd` and `jax.jacrev(jax.grad(...))` should return the same kind of
finite result `jax.grad` does (or, if this combination is genuinely
unsupported, we'd expect an explicit error rather than a silent NaN).

## Actual behavior

```
Simulation took final_i=2 steps (of MAX_STEPS=3).

=== jax.grad -- works correctly ===
loss=2130.704021017755
grad=-177.3106891539839
finite: loss=True, grad=True

=== jax.jacfwd -- expected: NaN ===
jacfwd result = nan
finite = False

=== jax.jacrev(jax.grad(...)) -- expected: NaN ===
hessian = nan
finite = False
```

## Investigation / what we ruled out

We hit this originally on our full production pipeline (custom EQDSK
geometry, a `CombinedTransportModel` patch overridden via list-indexed
`update_provider`, `jax.checkpoint`-wrapped `run_loop_jit`, an MHD
sawtooth-crash trigger, ~250-400 timesteps) and narrowed it down to the
minimal case above. Along the way we tested and ruled out each of the
following as *necessary* to reproduce (removing each individually still
reproduces both failures):

- `jax.checkpoint` / gradient checkpointing.
- The MHD sawtooth-crash trigger's discrete branching (we originally
suspected this, since it's the one deliberately-discrete physics event in
our scenario — but the failure reproduces with no `mhd` config at all).
- `CombinedTransportModel` + list-indexed `update_provider` (reproduces with
a single top-level `ConstantTransportModel`, no list indexing).
- EQDSK/CHEASE geometry or any custom scenario-building code (reproduces
with plain analytic circular geometry).
- A large step count (reproduces with only 2 real steps).

The one thing that *is* required: genuine multi-step time evolution under
`time_step_calculator.calculator_type='fixed'`. With the default `'chi'`
adaptive calculator, this particular toy scenario collapses to a single
step, and in that single-step case `jax.jacrev(jax.grad(...))` does *not*
fail (though `jax.jacfwd` still does, even at a single step) — so the two
failures may not share a root cause.

[repro.py](https://github.com/user-attachments/files/30683025/repro.py)

[environment.txt](https://github.com/user-attachments/files/30683029/environment.txt)

Start by running the attached repro.py with the stated environment and compare its jax.grad, jax.jacfwd, and jax.jacrev(jax.grad(...)) results. Read torax.experimental.run_loop_jit and the documented usage in torax/examples/iter_hybrid_rampup_grad_and_vmap.ipynb, focusing on fixed multi-step evolution. Done means finite Jacobian and Hessian results, or an explicit error if the combination is unsupported.

Written by the indexing model from the issue text.