{"record":{"id":"66bd15b275baf658","repo":"jax-ml/jax","slug":"lax-scan-f-argument-should-be-a-callable","errorCode":null,"errorMessage":"lax.scan: f argument should be a callable.","messagePattern":"lax\\.scan: f argument should be a callable\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":346,"sourceCode":"      many unrolled loop iterations to run within a single rolled iteration of\n      the loop. `unroll=0` unrolls the entire loop.\n      If a boolean is provided, it will determine if the loop is\n      completely unrolled (i.e. `unroll=True`) or left completely rolled (i.e.\n      `unroll=False`).\n\n  Returns:\n    A pair of type ``(c, [b])`` where the first element represents the final\n    loop carry value and the second element represents the stacked outputs of\n    the second output of ``f`` when scanned over the leading axis of the inputs.\n\n  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature\n  \"\"\"\n\n  if config.scan3.value:\n    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 = []","sourceCodeStart":328,"sourceCodeEnd":364,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L328-L364","documentation":"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.","triggerScenarios":"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.","commonSituations":"Forgetting that scan takes the function itself; accidentally writing f(...) in the argument; passing a functools.partial of a non-callable.","solutions":["Pass the function reference: lax.scan(f, init, xs), not lax.scan(f(init, x), ...)","Ensure f is callable: assert callable(f)","Make sure f has signature f(carry, x) -> (carry, y)"],"exampleFix":"// before\nresult = lax.scan(step(carry, x), carry, xs)\n// after\nresult = lax.scan(step, carry, xs)","handlingStrategy":"type-guard","validationCode":"assert callable(f), 'lax.scan expects the body function itself, not its result'","typeGuard":"def is_scan_callable(f) -> bool:\n    return callable(f)","tryCatchPattern":"null","preventionTips":["Pass function references, never invocations","Standardize body signature f(carry, x) -> (carry, y)","Code-review for 'step(' appearing inside scan args"],"tags":["jax","scan","typeerror","callable","api-misuse"],"backgroundTag":"callback-must-be-callable","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}