{"record":{"id":"b425358778068907","repo":"jax-ml/jax","slug":"can-only-use-unroll-in-fori-loop-if-the-loop-b","errorCode":null,"errorMessage":"Can only use `unroll` in `fori_loop` if the loop bounds are statically known.","messagePattern":"Can only use `unroll` in `fori_loop` if the loop bounds are statically known\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2660,"sourceCode":"\n  if use_scan:\n    if unroll is None:\n      unroll = False\n    length = max(upper_ - lower_, 0)\n    if config.disable_jit.value and length == 0:\n      # non-jit implementation of scan does not support length=0\n      return init_val\n    scan_body = _fori_scan_body_fun(body_fun, body_fun_dbg)\n    (_, result), _ = scan(\n        scan_body,\n        (lower_, init_val),\n        None,\n        length=length,\n        unroll=unroll,\n    )\n    return result\n  if unroll is not None and unroll is not False and unroll != 1:\n    raise ValueError(\"Can only use `unroll` in `fori_loop` if the loop bounds \"\n                     \"are statically known.\")\n\n  if lower_dtype != dtype:\n    lower = lax.convert_element_type(lower, dtype)\n  if upper_dtype != dtype:\n    upper = lax.convert_element_type(upper, dtype)\n  while_body_fun = _fori_body_fun(body_fun, body_fun_dbg)\n  _, _, result = while_loop(_fori_cond_fun, while_body_fun,\n                            (lower, upper, init_val))\n  return result\n\n### map and miscellaneous rules\n\ndef _scan_leaf(leaf, batch_elems, num_batches, batch_size):\n  def f(l):\n    return l[:batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:])\n\n  aval = core.typeof(leaf)","sourceCodeStart":2642,"sourceCodeEnd":2678,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2642-L2678","documentation":"fori_loop only supports the unroll option when the trip count is statically known, because unrolling requires the compiler to know the loop bounds at trace time. With dynamic (traced) bounds the loop lowers to a while_loop, where unrolling of a data-dependent iteration count is impossible, so JAX raises this ValueError.","triggerScenarios":"Calling lax.fori_loop(lower, upper, body, init, unroll=N) (N>1 or True) where lower/upper are traced arrays or otherwise non-concrete (e.g. computed inside jit from inputs), not Python ints.","commonSituations":"Porting scan or static loops to fori_loop with unroll for performance, but forgetting that bounds coming from data (sequence lengths, batch sizes) inside jit are dynamic; also unroll=N where N was intended for a different API.","solutions":["Make the bounds static: pass lower/upper as Python ints or via static_argnums/partial so JAX specializes the trip count and uses scan-based unrolling","Drop the unroll argument (use unroll=None/1) when bounds must stay dynamic","Cap iterations at a static maximum and mask updates per iteration (common pattern for dynamic-length sequences)","Replace with lax.scan(fixed length + mask, unroll=N) to keep unrolling benefits"],"exampleFix":"// before\n@jax.jit\ndef f(x, n):\n  return lax.fori_loop(0, n, body, x, unroll=4)  # n is traced\n// after\n@partial(jax.jit, static_argnums=(1,))\ndef f(x, n):\n  return lax.fori_loop(0, n, body, x, unroll=4)","handlingStrategy":"validation","validationCode":"import jax\ndef can_unroll(lower, upper, unroll):\n    if unroll in (None, False, 1):\n        return True\n    return jax.core.is_concrete(lower) and jax.core.is_concrete(upper)","typeGuard":"def bounds_static(lo, hi) -> bool:\n    return isinstance(lo, int) and isinstance(hi, int)","tryCatchPattern":null,"preventionTips":["Only pass unroll when bounds are static Python ints","Use static_argnums for bounds that control unrolled loops","For dynamic-length work, use lax.scan with static length plus masking"],"tags":["jax","fori-loop","unroll","static-shape","traced-values"],"backgroundTag":"jax-static-value-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}