jax-ml/jax · error · ValueError

`unroll` must be a `bool` or a non-negative `int`.

Error message

`unroll` must be a `bool` or a non-negative `int`.

What it means

scan_nocarry validates that the unroll argument is a non-negative int or bool. A negative unroll (e.g. -1 meaning 'auto' in other libraries) is meaningless for scan and raises ValueError.

Source

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

  jaxpr, y_avals = pe.trace_to_jaxpr(f, args_avals, dbg_body)
  jaxpr, consts = pe.separate_consts(jaxpr)

  if config.mutable_array_checks.value:
    _check_no_aliased_closed_over_refs(dbg_body, consts, list(xs_flat))

  disallowed_effects = effects.control_flow_allowed_effects.filter_not_in(jaxpr.effects)
  if disallowed_effects:
    raise NotImplementedError(
        f'Effects not supported in `scan`: {disallowed_effects}')

  unroll = core.concrete_or_error(
      None, unroll,
      "The `unroll` argument to `scan` expects a concrete `int` or `bool` "
      "value.")
  if isinstance(unroll, bool):
    unroll = max(length, 1) if unroll else 1
  if unroll < 0:
    raise ValueError("`unroll` must be a `bool` or a non-negative `int`.")

  args = list(consts) + list(xs_flat)
  # TODO(dougalm): handle traceable-level forwarding
  out = Scan3(
      extensives = [False] * len(consts) + [True] * len(xs_flat),
      length=length, jaxpr=jaxpr, reverse=reverse, unroll=unroll)(args)

  return y_avals.update(out).unflatten()

@partial(api_boundary, repro_api_name="jax.lax.scan")
def scan3[Carry, X, Y](
    f: Callable[[Carry, X], tuple[Carry, Y]],
    init: Carry,
    xs: X | None = None,
    length: int | None = None,
    reverse: bool = False,
    unroll: int | bool = 1,
    _split_transpose: bool = False) -> tuple[Carry, Y]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass a positive int (e.g. 1, 2, 4) or a bool
  2. If you wanted automatic unrolling, pick an explicit divisor of the scan length
  3. Validate unroll >= 0 before calling scan

Example fix

// before
lax.scan(f, init, xs, unroll=-1)
// after
lax.scan(f, init, xs, unroll=4)
Defensive patterns

Strategy: validation

Validate before calling

assert isinstance(unroll, (bool, int)) and (isinstance(unroll, bool) or unroll >= 0), 'unroll must be bool or non-negative int'

Type guard

def valid_unroll(u) -> bool:
    return isinstance(u, bool) or (isinstance(u, int) and u >= 0)

Try / catch

null

Prevention

When it happens

Trigger: Calling lax.scan(f, init, xs, unroll=-1) or scan_nocarry with a negative integer unroll.

Common situations: Porting code from libraries where negative values mean 'auto unroll'; passing a computed unroll that underflows to negative.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors


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