jax-ml/jax · error · ValueError

Scan with {num_extensive_inputs} extensive argument(s) is no

Error message

Scan with {num_extensive_inputs} extensive argument(s) is not supported. Found {num_consts} consts and {num_carry} carry arguments.

What it means

When pattern-matching a lax.scan to a fori_loop inside Pallas, the scanned jaxpr may only take consts and carry arguments. Extra (extensive) inputs — len(invars) - num_consts - num_carry > 0 — cannot be lowered to the Triton loop.

Source

Thrown at jax/_src/pallas/utils.py:122

def next_power_of_2(x: int) -> int:
  """Returns the next power of two greater than or equal to `x`."""
  if x < 0:
    raise ValueError("`next_power_of_2` requires a non-negative integer.")
  return 1 if x == 0 else 2 ** (x - 1).bit_length()


def pattern_match_scan_to_fori_loop(
    jaxpr: jax_core.Jaxpr, num_consts: int, num_carry: int
) -> tuple[jax_core.Jaxpr, bool]:
  num_extensive_inputs = len(jaxpr.invars) - num_consts - num_carry
  num_extensive_outputs = len(jaxpr.outvars) - num_carry
  if num_extensive_outputs:
    raise ValueError(
        f"Scan with {num_extensive_outputs} extensive output(s) is not"
        " supported."
    )
  if num_extensive_inputs:
    raise ValueError(
        f"Scan with {num_extensive_inputs} extensive argument(s) is not"
        f" supported. Found {num_consts} consts and {num_carry} carry"
        " arguments."
    )
  if num_carry > 0:
    # Pattern match onto fori_loop:
    # We expect the first carry argument to the jaxpr to be the loop index and
    # for the loop index + 1 to be returned as the first value out of the loop.
    in_index_var = jaxpr.invars[num_consts]
    out_index_var = jaxpr.outvars[0]
    assert isinstance(in_index_var.aval, jax_core.ShapedArray)
    # Check that the loop index argument is an int32 scalar
    if (in_index_var.aval.shape or
        in_index_var.aval.dtype not in (jnp.int32, jnp.int64)):
      # The loop index is not an int32 scalar so we assume that the loop index
      # has been DCEd and the body does *not* expect a loop index as an
      # argument.
      return jaxpr, False

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move per-iteration data into the carry tuple and update it each iteration
2. Hoist the data out as a const (compute outside the kernel or before scan)
3. Convert the scan to an explicit Python/fori loop indexing into a Ref

Example fix

# before
_, out = lax.scan(lambda c, xs: f(c, xs), c0, stacked)  # xs is extensive
# after
c, _ = lax.scan(lambda c, _: f(c, stacked_const), c0, None, length=n)

# all scan inputs must be consts or carry; fold per-iter data into carry
x, _ = lax.scan(lambda c, _: step(c, const_data), c0, None, length=n)

When it happens

Trigger: A pallas kernel body calls lax.scan on a closure that captures or takes additional per-iteration array inputs beyond consts and carry, e.g. scan(f)(x_stacked) where the stacked input is sliced per iteration.

Common situations: Using scan over a leading axis of an input array instead of folding that data into the carry or consts; closures capturing block arrays inside scan.

Related errors

• Scan with {num_extensive_outputs} extensive output(s) is not
• Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value
• Only positional arguments are supported by debug_print on Pa
• unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}
• end must be greater than start, but got: {end} <= {start}

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