jax-ml/jax · error · ValueError

Reverse-mode differentiation does not work for lax.while_loo

Error message

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.

What it means

Reverse-mode differentiation (grad/vjp) of a while_loop/fori_loop requires knowing the number of iterations ahead of time so JAX can build the linearized transpose. When the loop bounds (or the while condition's trip count) are dynamic (traced) values, JAX cannot reverse the loop and raises this error, suggesting alternatives that have known trip counts.

Source

Thrown at jax/_src/lax/control_flow/loops.py:2076

                                      params_known['cond_nconsts'])
  eqn_known = pe.new_jaxpr_eqn(ins_known, out_binders_known, while_p,
                               params_known, effects_known, eqn.source_info,
                               eqn.ctx)
  # Typecheck known eqn.
  _while_loop_abstract_eval(
      *[v.aval for v in eqn_known.invars], cond_jaxpr=cond_jaxpr_known,
      body_jaxpr=body_jaxpr_known, body_nconsts=params_known['body_nconsts'],
      cond_nconsts=params_known['cond_nconsts'])

  # Staged eqn is same as input eqn.
  eqn_staged = eqn

  unks_out = carry_uk
  inst_out = [True] * len(unks_out)
  return eqn_known, eqn_staged, unks_out, inst_out, new_inst

def _while_transpose_error(*_, **kwargs):
  raise 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.")

# For a while loop with ordered effects in the cond, we need a special
# lowering. Fundamentally, we'd like to rewrite a while loop that looks like
# this:
# ```
# while cond(x):
#   x = body(x)
# ```
# into something that looks like this:
# ```
# while True:
#   token, pred = cond(token, x)
#   if not pred:
#     break
#   token, x = body(token, x)
# ```

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use lax.fori_loop with static (Python int or concrete) start/stop so JAX specializes the trip count
2. Replace the dynamic loop with lax.scan if you can express the fixed maximum iterations with masking (e.g. run to max_len and mask inactive steps)
3. Recompute the bound outside jit and pass it as a static argument (functools.partial or static_argnums)
4. Wrap the whole loop in a custom_vjp defining explicit forward/backward rules if differentiation through a truly dynamic loop is required

Example fix

// before
def f(x, n):
  return lax.fori_loop(0, n, lambda i, a: a + x[i], 0.0)
jax.grad(lambda x: f(x, n_tapped))(x)
// after
def f(x, n):
  n = int(n)  # pass as static_argnums instead if inside jit
  return lax.fori_loop(0, n, lambda i, a: a + x[i], 0.0)

import jax
def bounds_are_static(lower, upper):
    return jax.core.is_concrete(lower) and jax.core.is_concrete(upper)
# or: type(lower) is int and type(upper) is int outside jit

def differentiable_loop_safe(lower, upper):
    return isinstance(lower, (int,)) and isinstance(upper, (int,)) or \
           (jax.core.is_concrete(lower) and jax.core.is_concrete(upper))

When it happens

Trigger: Calling jax.grad (or vjp/linearize) on a function using lax.while_loop with a dynamic condition, or lax.fori_loop whose lower/upper are traced arrays (e.g. computed from inputs or inside jit with dynamic values), rather than Python ints / static values.

Common situations: Migrating a Python for loop to fori_loop inside jit where bounds come from data (e.g. sequence lengths, number of steps from a tensor); using while(cond_fn) loops in optimization/solver code and then trying to differentiate through them.

Related errors

• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• primal and tangent arguments to jax.jvp do not match; dtypes
• jvp called with different primal and tangent shapes;Got prim
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/356c1f5d1d00a22f. Report an issue: GitHub.