jax-ml/jax · error · ValueError

scan got value with no leading axis to scan over: {}.

Error message

scan got value with no leading axis to scan over: {}.

What it means

scan infers the loop length from x.shape[0] of each xs element. If an element has no .shape (a scalar, a Python number, or a non-array object mixed into the xs pytree), inference fails and raises ValueError naming the offending values.

Source

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

    length: Any | None) -> int:

  # TODO(dougalm): put this in some sort of `scannable` typeclass
  from jax._src.hijax import HiType
  is_hi = [isinstance(a, HiType) for a in xs_avals]
  if xs_flat and all(is_hi):
    if length is None:
      raise ValueError(
          "must provide `length` to `scan`, since the leading-axis size of "
          "non-array (hijax) types cannot be inferred")
    return length
  xs_flat = [x for x, h in zip(xs_flat, is_hi) if not h]
  xs_avals = [a for a, h in zip(xs_avals, is_hi) if not h]

  try:
    lengths: list[int] = [x.shape[0] for x in xs_flat]
  except AttributeError as err:
    msg = "scan got value with no leading axis to scan over: {}."
    raise ValueError(
      msg.format(', '.join(str(x) for x in xs_flat
                           if not hasattr(x, 'shape')))) from err

  xs_shaped_avals = lax_utils.ensure_shaped(*xs_avals)
  if not all(a.sharding.spec.partitions[0] is None for a in xs_shaped_avals):
    raise ValueError('0th dimension of all xs should be replicated. Got '
                     f'{", ".join(str(a.sharding.spec) for a in xs_shaped_avals)}')

  if length is not None:
    try:
      length = int(length)
    except core.ConcretizationTypeError:
      msg = ('The `length` argument to `scan` expects a concrete `int` value.'
             ' For scan-like iteration with a dynamic length, use `while_loop`'
             ' or `fori_loop`.')
      raise core.ConcretizationTypeError(length, msg) from None
    else:
      if not all(length == l for l in lengths):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Close over constants in f instead of putting them in xs
2. Convert scalars to at least 1-D arrays with shape (length,) via jnp.broadcast_to or [c]*n wrapped in jnp.asarray
3. Ensure every leaf of xs is an ndarray with a leading axis

Example fix

// before
const = 2.0
lax.scan(lambda c, t: (c, t[0] + const), 0.0, (xs, const))
# const has no shape -> error
// after
def body(c, x):
  return c, x + 2.0
carry, ys = lax.scan(body, 0.0, xs)

flat, _ = jax.tree_util.tree_flatten(xs)
assert all(hasattr(x, 'shape') and x.ndim >= 1 for x in flat), \
    f'leaves without leading axis: {[x for x in flat if not hasattr(x, "shape")]}'

def all_xs_have_leading_axis(xs) -> bool:
    return all(hasattr(x, 'shape') and getattr(x, 'ndim', 0) >= 1
               for x in jax.tree_util.tree_leaves(xs))

null

When it happens

Trigger: lax.scan(f, init, xs=(1, array)) — a Python scalar leaf inside xs; or xs being a plain scalar; or an object without .shape in the pytree.

Common situations: Passing hyperparameters or scalars inside xs by accident; mixing a per-step constant into the scanned sequence instead of closing over it.

Related errors

• scan body output must be a pair, got {}.
• {} function carry input and carry output must have the same
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

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