jax-ml/jax · error · TypeError

lax.scan: f argument should be a callable.

Error message

lax.scan: f argument should be a callable.

What it means

lax.scan requires its first argument to be a Python callable describing the loop body. Passing an array, a jaxpr, or the already-applied result of a function raises TypeError immediately.

Source

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

      many unrolled loop iterations to run within a single rolled iteration of
      the loop. `unroll=0` unrolls the entire loop.
      If a boolean is provided, it will determine if the loop is
      completely unrolled (i.e. `unroll=True`) or left completely rolled (i.e.
      `unroll=False`).

  Returns:
    A pair of type ``(c, [b])`` where the first element represents the final
    loop carry value and the second element represents the stacked outputs of
    the second output of ``f`` when scanned over the leading axis of the inputs.

  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature
  """

  if config.scan3.value:
    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 = []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the function reference: lax.scan(f, init, xs), not lax.scan(f(init, x), ...)
2. Ensure f is callable: assert callable(f)
3. Make sure f has signature f(carry, x) -> (carry, y)

Example fix

// before
result = lax.scan(step(carry, x), carry, xs)
// after
result = lax.scan(step, carry, xs)

assert callable(f), 'lax.scan expects the body function itself, not its result'

def is_scan_callable(f) -> bool:
    return callable(f)

null

When it happens

Trigger: lax.scan(f(carry, x), init, xs) (calling f instead of passing f), or lax.scan(some_array, init, xs), or passing a non-function object.

Common situations: Forgetting that scan takes the function itself; accidentally writing f(...) in the argument; passing a functools.partial of a non-callable.

Related errors

• lax.while_loop: body_fun and cond_fun arguments should be ca
• check_error takes an Error as argument, got type {type(error
• lax.platform_dependent: the '{pname}' branch must be a calla
• lax.platform_dependent: the 'default' branch must be a calla
• scan body output must be a pair, got {}.

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