{"record":{"id":"731b48a34672f5f5","repo":"jax-ml/jax","slug":"lax-fori-loop-body-fun-argument-should-be-callabl","errorCode":null,"errorMessage":"lax.fori_loop: body_fun argument should be callable.","messagePattern":"lax\\.fori_loop: body_fun argument should be callable\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2601,"sourceCode":"  Args:\n    lower: an integer representing the loop index lower bound (inclusive)\n    upper: an integer representing the loop index upper bound (exclusive)\n    body_fun: function of type ``(int, a) -> a``.\n    init_val: initial loop carry value of type ``a``.\n    unroll: An optional integer or boolean that determines how much to unroll\n      the loop. If an integer is provided, it determines how many unrolled\n      loop iterations to run within a single rolled iteration of the loop. If a\n      boolean is provided, it will determine if the loop is completely unrolled\n      (i.e. `unroll=True`) or left completely unrolled (i.e. `unroll=False`).\n      This argument is only applicable if the loop bounds are statically known.\n\n  Returns:\n    Loop value from the final iteration, of type ``a``.\n\n  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature\n  \"\"\"\n  if not callable(body_fun):\n    raise TypeError(\"lax.fori_loop: body_fun argument should be callable.\")\n\n  # TODO(phawkins): perhaps do more type checking here, better error messages.\n  lower_dtype = lax.dtype(lower)\n  upper_dtype = lax.dtype(upper)\n  if lower_dtype == upper_dtype:\n    dtype = lower_dtype\n  else:\n    # As a special case: allow promotion of weak integers (e.g., Python scalars)\n    # This improves the ergonomics if one but not both of the loop bounds is a\n    # scalar.\n    dtype = None\n    if (np.issubdtype(lower_dtype, np.signedinteger) and\n        np.issubdtype(upper_dtype, np.signedinteger)):\n      lower_weak = dtypes.is_weakly_typed(lower)\n      upper_weak = dtypes.is_weakly_typed(upper)\n      if lower_weak and not upper_weak:\n        dtype = upper_dtype\n      elif not lower_weak and upper_weak:","sourceCodeStart":2583,"sourceCodeEnd":2619,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2583-L2619","documentation":"lax.fori_loop validates that its body_fun argument is callable (a Python/JAX function). Passing anything else — an array result, a module, None, or the result of calling body_fun instead of the function itself — raises this TypeError immediately at trace time.","triggerScenarios":"Calling lax.fori_loop(lower, upper, body_fun, init) where body_fun is not callable, e.g. body_fun(i, x) (invoked instead of passed), or passing a jitted result, None, or a non-function object as the third positional argument.","commonSituations":"Accidentally invoking the body function instead of passing it (missing lambda), passing keyword args in the wrong order, or refactoring code so a variable holding a function is shadowed by its return value.","solutions":["Pass the function itself, not its result: fori_loop(0, n, body_fun, init) where body_fun is def body_fun(i, carry): ...","If using a lambda, ensure the signature is lambda i, carry: ... (two args, uninvoked)","Check for shadowing: verify no earlier assignment replaced body_fun with its call result"],"exampleFix":"// before\nout = lax.fori_loop(0, 10, step(i, x), x)  # called!\n// after\nout = lax.fori_loop(0, 10, step, x)","handlingStrategy":"type-guard","validationCode":"import typing\ndef validate_fori_args(lower, upper, body_fun, init):\n    if not callable(body_fun):\n        raise TypeError('body_fun must be callable, got %r' % type(body_fun))\n    return True","typeGuard":"def is_body_fun(x) -> bool:\n    return callable(x) and not isinstance(x, (jnp.ndarray, np.ndarray))","tryCatchPattern":"try:\n    lax.fori_loop(0, n, body, init)\nexcept TypeError as e:\n    if 'body_fun argument should be callable' in str(e):\n        raise  # programmer error: fix at call site","preventionTips":["Pass functions, never their results: fori_loop(0, n, step, x)","Use named keyword-adjacent signatures: def step(i, carry): ...","Lint for lax.fori_loop calls whose 3rd arg is not a lambda/function name"],"tags":["jax","fori-loop","typeerror","argument-validation"],"backgroundTag":"argument-not-callable","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}