jax-ml/jax · error · ValueError

zero-length scan is not supported in disable_jit() mode beca

Error message

zero-length scan is not supported in disable_jit() mode because the output type is unknown.

What it means

With jax.disable_jit() (or config.disable_jit), scan runs eagerly in Python. A zero-length scan has no iterations to run, so the output type/shape cannot be determined, and eager execution cannot fabricate it — hence ValueError.

Source

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

    return scan3(f, init, xs, length, reverse, unroll)

  if not callable(f):
    raise TypeError("lax.scan: f argument should be a callable.")

  dbg_body = api_util.debug_info("scan", f, (init, xs), {})
  init_flat = ft.flatten(init)
  xs_flat = ft.flatten(xs)
  args = ft.pack((init_flat, xs_flat))
  check_no_transformed_refs_args(lambda: dbg_body, args.vals)
  del init, xs

  args_avals = args.map(core.typeof)
  init_avals, xs_avals = args_avals.unpack()
  length = _infer_scan_length(list(xs_flat), list(xs_avals), length)

  if config.disable_jit.value:
    if length == 0:
      raise ValueError("zero-length scan is not supported in disable_jit() "
                       "mode because the output type is unknown.")
    carry = init_flat.unflatten()
    ys = []
    maybe_reversed = reversed if reverse else lambda x: x
    for i in maybe_reversed(range(length)):
      xs_slice = xs_flat.map(lambda x: slicing.index_in_dim(x, i, keepdims=False))
      carry, y = f(carry, xs_slice.unflatten())
      ys.append(y)
    stack = lambda *ys: _stack(ys)
    stacked_y = tree_map(stack, *maybe_reversed(ys))
    return carry, stacked_y

  if config.mutable_array_checks.value:
    check_no_aliased_ref_args(lambda: dbg_body, list(args_avals), list(args))

  x_avals = xs_avals.map(lambda aval: core.mapped_leading_aval(length, aval))
  def _create_jaxpr(carry_avals):
    new_arg_avals = ft.pack(((carry_avals, x_avals), {}))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass an explicit length=0-compatible structure by providing xs with correct trailing shapes even when length 0, under jit keep it compiled
2. Guard in Python: skip the scan when the batch is empty and return init plus an empty stacked y
3. Re-enable jit for this section
4. Construct ys explicitly: jnp.zeros((0, *trailing_shape)) instead of relying on scan output

Example fix

// before
with jax.disable_jit():
  carry, ys = lax.scan(step, init, xs[:0])
// after
with jax.disable_jit():
  if xs.shape[0] == 0:
    ys = jnp.zeros((0,) + xs.shape[1:])
    carry, ys_out = init, ys
  else:
    carry, ys_out = lax.scan(step, init, xs)

n = xs.shape[0] if hasattr(xs, 'shape') else length
if jax.config.disable_jit.value and n == 0:
    return init, jnp.zeros((0,) + trailing_shape)  # bypass scan

def safe_to_scan_eagerly(xs) -> bool:
    return not jax.config.disable_jit.value or xs.shape[0] > 0

try: with jax.disable_jit(): carry, ys = lax.scan(step, init, xs)
except ValueError as e:
    if 'zero-length' in str(e): carry, ys = init, jnp.zeros((0,) + xs.shape[1:])
    else: raise

When it happens

Trigger: lax.scan(f, init, xs) where xs has leading axis 0 (e.g. jnp.zeros((0, ...))) while disable_jit is active, or length=0 with no xs.

Common situations: Debugging with disable_jit() or JAX_DEBUG_JITS on empty batches; edge-case dataset with zero elements; tests exercising empty inputs.

Related errors

• Effects not supported in `scan`: {disallowed_effects}
• `unroll` must be a `bool` or a non-negative `int`.
• lax.scan: f argument should be a callable.
• scan body output must be a pair, got {}.
• must provide `length` to `scan`, since the leading-axis size

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