xwhile

t3toolbox.backend.common.xwhile(cond, body, init_state, use_jit=False)

def xwhile(
        cond,                  # state -> 0-d bool (eager) / traced bool (jit): loop while True
        body,                  # state -> state (a pytree of the SAME structure and leaf shapes)
        init_state,            # the loop-carried state (pytree)
        use_jit: bool = False, # True + jax state -> jax.lax.while_loop; else a Python while loop
):

Data-dependent while with the numpy / eager-jax / jit dispatch – the xscan precedent for a while. With use_jit and a jax state it compiles via jax.lax.while_loop; otherwise (numpy, eager jax, or jax not installed) it runs while bool(cond(state)): state = body(state), so it works on every backend and silently falls back to eager when jit is unavailable. Write cond/body backend-agnostically (cond returns a 0-d boolean; body uses xnp.where, not Python branches) so the SAME pair drives both paths.

Parameters:

use_jit (bool)