{"record":{"id":"356c1f5d1d00a22f","repo":"jax-ml/jax","slug":"reverse-mode-differentiation-does-not-work-for-lax","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2076,"sourceCode":"                                      params_known['cond_nconsts'])\n  eqn_known = pe.new_jaxpr_eqn(ins_known, out_binders_known, while_p,\n                               params_known, effects_known, eqn.source_info,\n                               eqn.ctx)\n  # Typecheck known eqn.\n  _while_loop_abstract_eval(\n      *[v.aval for v in eqn_known.invars], cond_jaxpr=cond_jaxpr_known,\n      body_jaxpr=body_jaxpr_known, body_nconsts=params_known['body_nconsts'],\n      cond_nconsts=params_known['cond_nconsts'])\n\n  # Staged eqn is same as input eqn.\n  eqn_staged = eqn\n\n  unks_out = carry_uk\n  inst_out = [True] * len(unks_out)\n  return eqn_known, eqn_staged, unks_out, inst_out, new_inst\n\ndef _while_transpose_error(*_, **kwargs):\n  raise ValueError(\"Reverse-mode differentiation does not work for \"\n                   \"lax.while_loop or lax.fori_loop with dynamic start/stop values. \"\n                   \"Try using lax.scan, or using fori_loop with static start/stop.\")\n\n# For a while loop with ordered effects in the cond, we need a special\n# lowering. Fundamentally, we'd like to rewrite a while loop that looks like\n# this:\n# ```\n# while cond(x):\n#   x = body(x)\n# ```\n# into something that looks like this:\n# ```\n# while True:\n#   token, pred = cond(token, x)\n#   if not pred:\n#     break\n#   token, x = body(token, x)\n# ```","sourceCodeStart":2058,"sourceCodeEnd":2094,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2058-L2094","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Use lax.fori_loop with static (Python int or concrete) start/stop so JAX specializes the trip count","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)","Recompute the bound outside jit and pass it as a static argument (functools.partial or static_argnums)","Wrap the whole loop in a custom_vjp defining explicit forward/backward rules if differentiation through a truly dynamic loop is required"],"exampleFix":"// before\ndef f(x, n):\n  return lax.fori_loop(0, n, lambda i, a: a + x[i], 0.0)\njax.grad(lambda x: f(x, n_tapped))(x)\n// after\ndef f(x, n):\n  n = int(n)  # pass as static_argnums instead if inside jit\n  return lax.fori_loop(0, n, lambda i, a: a + x[i], 0.0)","handlingStrategy":"validation","validationCode":"import jax\ndef bounds_are_static(lower, upper):\n    return jax.core.is_concrete(lower) and jax.core.is_concrete(upper)\n# or: type(lower) is int and type(upper) is int outside jit","typeGuard":"def differentiable_loop_safe(lower, upper):\n    return isinstance(lower, (int,)) and isinstance(upper, (int,)) or \\\n           (jax.core.is_concrete(lower) and jax.core.is_concrete(upper))","tryCatchPattern":null,"preventionTips":["Pass loop bounds as static_argnums when they determine trip counts","Default to lax.scan with a static length + masking for dynamic-length data","Wrap dynamic loops needing gradients in jax.custom_vjp with manual rules"],"tags":["jax","autodiff","grad","while-loop","fori-loop","dynamic-bounds"],"backgroundTag":"jax-grad-dynamic-loop-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}