jax-ml/jax · error · ValueError

must provide `length` to `scan`, since the leading-axis size

Error message

must provide `length` to `scan`, since the leading-axis size of non-array (hijax) types cannot be inferred

What it means

When all xs inputs to scan are non-array 'hijax' types (JAX's higher-level, non-ndarray values), the scan length cannot be inferred from a leading axis, so an explicit length must be passed.

Source

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

  ys = ys_g.update(ys).unfilter().map2(
      ext_to_ext_fwd, lambda y, f: y if f is None else _maybe_put(args_flat[f]))
  out = [*carry_out, *ys]

  if any(move_to_const):
    out = pe.merge_lists(move_to_const + [False] * num_ys, out, new_consts)

  return out_avals.update(out).unflatten()

def _infer_scan_length(
    xs_flat: list[Any], xs_avals: list[AbstractValue],
    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)}')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the length argument explicitly: lax.scan(f, init, xs, length=N)
2. Convert hijax-typed inputs to arrays so the leading axis is inferable
3. Restructure to fori_loop(length=N, ...) which takes length natively

Example fix

// before
lax.scan(body, init, hijax_xs)
// after
lax.scan(body, init, hijax_xs, length=n_steps)

if length is None and all_hi_types(xs): raise ValueError('pass length explicitly')
# or simply always pass length when xs is non-array

def needs_explicit_length(xs) -> bool:
    from jax._src.hijax import HiType
    return xs and all(isinstance(a, HiType) for a in xs)

null

When it happens

Trigger: lax.scan(f, init, xs) where every element of xs is a HiType instance and length=None.

Common situations: Using experimental hijax types or future non-array avals inside scans; internal/jax-dev usage more than user code.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

• Effects not supported in `scan`: {disallowed_effects}
• State effect not supported in vmap-of-cond.
• `unroll` must be a `bool` or a non-negative `int`.
• lax.scan: f argument should be a callable.
• zero-length scan is not supported in disable_jit() mode beca

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