{"record":{"id":"f4d9af1a1a4d017b","repo":"jax-ml/jax","slug":"zero-length-scan-is-not-supported-in-disable-jit","errorCode":null,"errorMessage":"zero-length scan is not supported in disable_jit() mode because the output type is unknown.","messagePattern":"zero-length scan is not supported in disable_jit\\(\\) mode because the output type is unknown\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":361,"sourceCode":"    return scan3(f, init, xs, length, reverse, unroll)\n\n  if not callable(f):\n    raise TypeError(\"lax.scan: f argument should be a callable.\")\n\n  dbg_body = api_util.debug_info(\"scan\", f, (init, xs), {})\n  init_flat = ft.flatten(init)\n  xs_flat = ft.flatten(xs)\n  args = ft.pack((init_flat, xs_flat))\n  check_no_transformed_refs_args(lambda: dbg_body, args.vals)\n  del init, xs\n\n  args_avals = args.map(core.typeof)\n  init_avals, xs_avals = args_avals.unpack()\n  length = _infer_scan_length(list(xs_flat), list(xs_avals), length)\n\n  if config.disable_jit.value:\n    if length == 0:\n      raise ValueError(\"zero-length scan is not supported in disable_jit() \"\n                       \"mode because the output type is unknown.\")\n    carry = init_flat.unflatten()\n    ys = []\n    maybe_reversed = reversed if reverse else lambda x: x\n    for i in maybe_reversed(range(length)):\n      xs_slice = xs_flat.map(lambda x: slicing.index_in_dim(x, i, keepdims=False))\n      carry, y = f(carry, xs_slice.unflatten())\n      ys.append(y)\n    stack = lambda *ys: _stack(ys)\n    stacked_y = tree_map(stack, *maybe_reversed(ys))\n    return carry, stacked_y\n\n  if config.mutable_array_checks.value:\n    check_no_aliased_ref_args(lambda: dbg_body, list(args_avals), list(args))\n\n  x_avals = xs_avals.map(lambda aval: core.mapped_leading_aval(length, aval))\n  def _create_jaxpr(carry_avals):\n    new_arg_avals = ft.pack(((carry_avals, x_avals), {}))","sourceCodeStart":343,"sourceCodeEnd":379,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L343-L379","documentation":"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.","triggerScenarios":"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.","commonSituations":"Debugging with disable_jit() or JAX_DEBUG_JITS on empty batches; edge-case dataset with zero elements; tests exercising empty inputs.","solutions":["Pass an explicit length=0-compatible structure by providing xs with correct trailing shapes even when length 0, under jit keep it compiled","Guard in Python: skip the scan when the batch is empty and return init plus an empty stacked y","Re-enable jit for this section","Construct ys explicitly: jnp.zeros((0, *trailing_shape)) instead of relying on scan output"],"exampleFix":"// before\nwith jax.disable_jit():\n  carry, ys = lax.scan(step, init, xs[:0])\n// after\nwith jax.disable_jit():\n  if xs.shape[0] == 0:\n    ys = jnp.zeros((0,) + xs.shape[1:])\n    carry, ys_out = init, ys\n  else:\n    carry, ys_out = lax.scan(step, init, xs)","handlingStrategy":"type-guard","validationCode":"n = xs.shape[0] if hasattr(xs, 'shape') else length\nif jax.config.disable_jit.value and n == 0:\n    return init, jnp.zeros((0,) + trailing_shape)  # bypass scan","typeGuard":"def safe_to_scan_eagerly(xs) -> bool:\n    return not jax.config.disable_jit.value or xs.shape[0] > 0","tryCatchPattern":"try: with jax.disable_jit(): carry, ys = lax.scan(step, init, xs)\nexcept ValueError as e:\n    if 'zero-length' in str(e): carry, ys = init, jnp.zeros((0,) + xs.shape[1:])\n    else: raise","preventionTips":["Guard empty batches before scanning, especially in debug/nocache modes","Keep scans under jit to preserve shape inference","Make empty-case outputs explicit in pipeline design"],"tags":["jax","scan","disable-jit","zero-length","edge-case"],"backgroundTag":"empty-batch-dimension","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}