jax-ml/jax · error · ValueError

scan got `length` argument of {} which disagrees with leadin

Error message

scan got `length` argument of {} which disagrees with leading axis sizes {}.

What it means

scan cross-checks the explicit length argument against the leading-axis size of every xs element. If they disagree (e.g. length=10 but xs.shape[0]==8), it raises ValueError showing both values.

Source

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

  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):
        msg = ("scan got `length` argument of {} which disagrees with "
              "leading axis sizes {}.")
        raise ValueError(msg.format(length, [x.shape[0] for x in xs_flat]))
      return length
  else:
    unique_lengths = set(lengths)
    if len(unique_lengths) > 1:
      msg = "scan got values with different leading axis sizes: {}."
      raise ValueError(msg.format(', '.join(str(x.shape[0]) for x in xs_flat)))
    elif len(unique_lengths) == 0:
      msg = "scan got no values to scan over and `length` not provided."
      raise ValueError(msg)
    else:
      return list(unique_lengths)[0]

def _capitalize(s):
  # s.capitalize() converts s[1:] to lowercase which we don't want.
  return s[0].capitalize() + s[1:]

def _check_carry_type(name, body_fun, in_carry, out_carry):
  try:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass length=None and let scan infer from xs
Recompute length from data: length=xs.shape[0]
Fix the off-by-one in how the length is derived (range(len(data)) vs len(data)+1)

Example fix

// before
lax.scan(body, init, xs, length=num_steps)  # xs has len num_steps+1
// after
lax.scan(body, init, xs[:num_steps], length=num_steps)
# or simply
lax.scan(body, init, xs)

Defensive patterns

Strategy: validation

Validate before calling

flat = jax.tree_util.tree_leaves(xs)
if length is not None:
    assert all(x.shape[0] == length for x in flat), \
        [(x.shape[0], length) for x in flat]

Type guard

def length_matches_xs(xs, length) -> bool:
    return length is None or all(x.shape[0] == length for x in jax.tree_util.tree_leaves(xs))

Try / catch

null

Prevention

Derive length from the data: length=xs.shape[0]
Omit length and let scan infer
Add a unit assert before scan in data-loading code

When it happens

Trigger: lax.scan(f, init, xs, length=n) where n != xs.shape[0] for any xs leaf; or xs of differing leading sizes with an explicit length matching only some.

Common situations: Off-by-one loop counts; slicing xs after computing length; passing a stale length constant after data preprocessing changed the batch size.

Related errors

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