jax-ml/jax · error · ValueError

scan got values with different leading axis sizes: {}.

Error message

scan got values with different leading axis sizes: {}.

What it means

When length is not given, scan infers it from the leading axes of xs and requires all leaves to agree. Mixed leading axis sizes (e.g. xs=(a of shape (10,...), b of shape (12,...))) raise this ValueError.

Source

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

  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:
    sig = inspect.signature(body_fun)
  except (ValueError, TypeError):
    sig = None
  carry_name = sig and list(sig.parameters)[0]
  if carry_name:
    component = lambda p: (f'the input carry component {carry_name}{keystr(p)}'

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Align all xs leading axes before scan (truncate/pad to common length)
Move non-conforming leaves out of xs into constants closed over by f
Pass explicit length and slice every leaf to it

Example fix

# before
xs = (a, b)  # a.shape==(10,..), b.shape==(12,..)
lax.scan(body, init, xs)
# after
xs = (a[:10], b[:10])
lax.scan(body, init, xs)

Defensive patterns

Strategy: validation

Validate before calling

sizes = {x.shape[0] for x in jax.tree_util.tree_leaves(xs)}
assert len(sizes) <= 1, f'inconsistent leading axes: {sizes}'

Type guard

def consistent_leading_axes(xs) -> bool:
    return len({x.shape[0] for x in jax.tree_util.tree_leaves(xs)}) <= 1

Try / catch

null

Prevention

Slice/pad all xs leaves to a common length in preprocessing
Keep per-step data in one aligned struct
Assert alignment after any concatenation or resampling

When it happens

Trigger: Passing a pytree of xs whose arrays have different first-dimension sizes; accidentally including a broadcastable array of different length.

Common situations: Tuple of sequences misaligned in preprocessing; concatenating features along the wrong axis so lengths drift; including a static array in xs that has another size.

Related errors

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