Reverse-mode `jax.grad` fails through `recomb_model` — `lax.while_loop` blocks AD

[MinhMPA]

Summary

Forward-mode AD (jax.jvp) works through recomb_model.__call__. Reverse-mode (jax.grad) fails:

ValueError: Reverse-mode differentiation does not work for lax.while_loop or                                                                                                 
lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using                                                                                                   
fori_loop with static start/stop.                                                                                                                                            

The error originates at three lax.while_loop sites:

• hyrex/helium.py:243 (Saha 4He, stop = xHeIII < threshold)
• hyrex/helium.py:394 (post-Saha 2He, stop = |xe_saha − xe| > threshold)
• hyrex/hydrogen.py:311 (post-Saha H + two-photon, stop = |xe_Saha − xe| > threshold)

Why it matters

Reverse-mode AD is required for HMC/NUTS over CMB likelihoods — forward-mode scales O(d) in parameter count. Without it, HyRex is forward-only / Fisher-only.

Reproducer

import jax, jax.numpy as jnp                                                                                                                                                 
jax.config.update("jax_enable_x64", True)                                                                                                                                  
from hyrex.hyrex import recomb_model
                                            
recomb = recomb_model()                 

def xe_at_zrec(omega_b):                                                                                                                                                     
    out_xe, out_lna, _, _ = recomb(
        jnp.float64(0.6736), omega_b, jnp.float64(0.12),                                                                                                                     
        jnp.float64(3.044), jnp.float64(0.2453),                                                                                                                           
    )                                   
    target_loga = -jnp.log(1 + 1090)
    valid_lna = jnp.where(jnp.isfinite(out_lna.arr), out_lna.arr, -1e10)                                                                                                     
    valid_xe  = jnp.where(jnp.isfinite(out_lna.arr), out_xe.arr,  0.0)
    return valid_xe[jnp.argmin(jnp.abs(valid_lna - target_loga))]                                                                                                            
                                                                                                                                                                           
jax.jvp(xe_at_zrec, (jnp.float64(0.02237),), (jnp.float64(1.0),))  # works                                                                                                   
jax.grad(xe_at_zrec)(jnp.float64(0.02237))                         # raises ValueError                                                                                       

Proposed fix

Each loop iterates over a static-size grid with a data-dependent early exit. Convert to lax.scan over jnp.arange(axis_size) with a stop-mask:

def scan_step(state, _):                                                                                                                                                     
    xe_output, lna_output, xe, iz, stop = state
    new_xe_output, new_lna_output, new_xe, new_iz, new_stop = compute_xe(state)                                                                                              
    return (                                                                                                                                                               
        jnp.where(stop, xe_output,  new_xe_output),
        jnp.where(stop, lna_output, new_lna_output),                                                                                                                         
        jnp.where(stop, xe,         new_xe),
        jnp.where(stop, iz,         new_iz),                                                                                                                                 
        stop | new_stop,                                                                                                                                                     
    ), None                             
                                                                                                                                                                             
final_state, _ = jax.lax.scan(scan_step, initial_state, jnp.arange(lna_axis.size))                                                                                           

Outputs are identical (masked branch is a no-op). Forward pass is ~2× more expensive since the body now runs to the static bound regardless of when the natural stop would
have triggered. Reverse-mode AD works after the change.

Happy to contribute a PR. Was there a specific reason while_loop was preferred — compile time, memory, something else I'm missing?