jax-ml/jax · error · ValueError

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

Error message

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

What it means

scan infers the trip count from the leading axis of array inputs, but non-array (hijax/hi-primitive) inputs have no inferable leading size. When all inputs are of such types and length was not given, scan refuses to guess and requires an explicit length.

Solutions

Pass the length argument explicitly: lax.scan(f, init, xs, length=N)
Convert hijax-typed inputs to arrays so the leading axis is inferable
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)

Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Try / catch

null

Prevention

Always pass length when scanning non-array inputs
Prefer array xs or fori_loop for data-free loops
Treat hijax types as experimental; pin jax versions

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

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

Appendix: 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)