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

[aaronkho]

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

environment.txt

[jcitrin]

@sbodenstein , @Nush395

[Nush395]

Hi @aaronkho, thanks for raising the issue and the detailed investigation. The forward mode NaN is something we discovered just recently too and didn't have time to look into yet. The jacrev failing is interesting and unknown to me so really good to know about. It's not intended behaviour so I'd prefer not to raise any errors until we've taken a look at what the issue might be. Will prioritise doing an investigation into this later today. Our suspicion is it's in one of our custom gradient implementations.

@sbodenstein also has a new while loop implementation landing soon that may fix these problems (though I believe it is orthogonal). Will keep you posted!

[Nush395]

Hi @aaronkho, I tested your repro with #2336 and that hopefully fixes things! We're just discussing internally what the correct behaviour is but a fix should be merged soon.