jax-ml/jax · error · ValueError

Can only use `unroll` in `fori_loop` if the loop bounds are

Error message

Can only use `unroll` in `fori_loop` if the loop bounds are statically known.

What it means

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.

Source

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

  if use_scan:
    if unroll is None:
      unroll = False
    length = max(upper_ - lower_, 0)
    if config.disable_jit.value and length == 0:
      # non-jit implementation of scan does not support length=0
      return init_val
    scan_body = _fori_scan_body_fun(body_fun, body_fun_dbg)
    (_, result), _ = scan(
        scan_body,
        (lower_, init_val),
        None,
        length=length,
        unroll=unroll,
    )
    return result
  if unroll is not None and unroll is not False and unroll != 1:
    raise ValueError("Can only use `unroll` in `fori_loop` if the loop bounds "
                     "are statically known.")

  if lower_dtype != dtype:
    lower = lax.convert_element_type(lower, dtype)
  if upper_dtype != dtype:
    upper = lax.convert_element_type(upper, dtype)
  while_body_fun = _fori_body_fun(body_fun, body_fun_dbg)
  _, _, result = while_loop(_fori_cond_fun, while_body_fun,
                            (lower, upper, init_val))
  return result

### map and miscellaneous rules

def _scan_leaf(leaf, batch_elems, num_batches, batch_size):
  def f(l):
    return l[:batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:])

  aval = core.typeof(leaf)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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
2. Drop the unroll argument (use unroll=None/1) when bounds must stay dynamic
3. Cap iterations at a static maximum and mask updates per iteration (common pattern for dynamic-length sequences)
4. Replace with lax.scan(fixed length + mask, unroll=N) to keep unrolling benefits

Example fix

// before
@jax.jit
def f(x, n):
  return lax.fori_loop(0, n, body, x, unroll=4)  # n is traced
// after
@partial(jax.jit, static_argnums=(1,))
def f(x, n):
  return lax.fori_loop(0, n, body, x, unroll=4)

import jax
def can_unroll(lower, upper, unroll):
    if unroll in (None, False, 1):
        return True
    return jax.core.is_concrete(lower) and jax.core.is_concrete(upper)

def bounds_static(lo, hi) -> bool:
    return isinstance(lo, int) and isinstance(hi, int)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• `unroll` must be a `bool` or a non-negative `int`.
• Reverse-mode differentiation does not work for lax.while_loo
• lax.fori_loop: body_fun argument should be callable.
• lower and upper arguments to fori_loop must have equal types
• numpy masked arrays are not supported as direct inputs to JA

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