{"record":{"id":"d585178d832517a8","repo":"jax-ml/jax","slug":"must-provide-length-to-scan-since-the-leading","errorCode":null,"errorMessage":"must provide `length` to `scan`, since the leading-axis size of non-array (hijax) types cannot be inferred","messagePattern":"must provide `length` to `scan`, since the leading-axis size of non-array \\(hijax\\) types cannot be inferred","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":484,"sourceCode":"  ys = ys_g.update(ys).unfilter().map2(\n      ext_to_ext_fwd, lambda y, f: y if f is None else _maybe_put(args_flat[f]))\n  out = [*carry_out, *ys]\n\n  if any(move_to_const):\n    out = pe.merge_lists(move_to_const + [False] * num_ys, out, new_consts)\n\n  return out_avals.update(out).unflatten()\n\ndef _infer_scan_length(\n    xs_flat: list[Any], xs_avals: list[AbstractValue],\n    length: Any | None) -> int:\n\n  # TODO(dougalm): put this in some sort of `scannable` typeclass\n  from jax._src.hijax import HiType\n  is_hi = [isinstance(a, HiType) for a in xs_avals]\n  if xs_flat and all(is_hi):\n    if length is None:\n      raise ValueError(\n          \"must provide `length` to `scan`, since the leading-axis size of \"\n          \"non-array (hijax) types cannot be inferred\")\n    return length\n  xs_flat = [x for x, h in zip(xs_flat, is_hi) if not h]\n  xs_avals = [a for a, h in zip(xs_avals, is_hi) if not h]\n\n  try:\n    lengths: list[int] = [x.shape[0] for x in xs_flat]\n  except AttributeError as err:\n    msg = \"scan got value with no leading axis to scan over: {}.\"\n    raise ValueError(\n      msg.format(', '.join(str(x) for x in xs_flat\n                           if not hasattr(x, 'shape')))) from err\n\n  xs_shaped_avals = lax_utils.ensure_shaped(*xs_avals)\n  if not all(a.sharding.spec.partitions[0] is None for a in xs_shaped_avals):\n    raise ValueError('0th dimension of all xs should be replicated. Got '\n                     f'{\", \".join(str(a.sharding.spec) for a in xs_shaped_avals)}')","sourceCodeStart":466,"sourceCodeEnd":502,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L466-L502","documentation":"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.","triggerScenarios":"lax.scan(f, init, xs) where every element of xs is a HiType instance and length=None.","commonSituations":"Using experimental hijax types or future non-array avals inside scans; internal/jax-dev usage more than user code.","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"],"exampleFix":"// before\nlax.scan(body, init, hijax_xs)\n// after\nlax.scan(body, init, hijax_xs, length=n_steps)","handlingStrategy":"validation","validationCode":"if length is None and all_hi_types(xs): raise ValueError('pass length explicitly')\n# or simply always pass length when xs is non-array","typeGuard":"def needs_explicit_length(xs) -> bool:\n    from jax._src.hijax import HiType\n    return xs and all(isinstance(a, HiType) for a in xs)","tryCatchPattern":"null","preventionTips":["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"],"tags":["jax","scan","hijax","length","experimental"],"backgroundTag":"missing-required-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}