jax-ml/jax · error · TypeError

lax.fori_loop: body_fun argument should be callable.

Error message

lax.fori_loop: body_fun argument should be callable.

What it means

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.

Source

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

  Args:
    lower: an integer representing the loop index lower bound (inclusive)
    upper: an integer representing the loop index upper bound (exclusive)
    body_fun: function of type ``(int, a) -> a``.
    init_val: initial loop carry value of type ``a``.
    unroll: An optional integer or boolean that determines how much to unroll
      the loop. If an integer is provided, it determines how many unrolled
      loop iterations to run within a single rolled iteration of the loop. If a
      boolean is provided, it will determine if the loop is completely unrolled
      (i.e. `unroll=True`) or left completely unrolled (i.e. `unroll=False`).
      This argument is only applicable if the loop bounds are statically known.

  Returns:
    Loop value from the final iteration, of type ``a``.

  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature
  """
  if not callable(body_fun):
    raise TypeError("lax.fori_loop: body_fun argument should be callable.")

  # TODO(phawkins): perhaps do more type checking here, better error messages.
  lower_dtype = lax.dtype(lower)
  upper_dtype = lax.dtype(upper)
  if lower_dtype == upper_dtype:
    dtype = lower_dtype
  else:
    # As a special case: allow promotion of weak integers (e.g., Python scalars)
    # This improves the ergonomics if one but not both of the loop bounds is a
    # scalar.
    dtype = None
    if (np.issubdtype(lower_dtype, np.signedinteger) and
        np.issubdtype(upper_dtype, np.signedinteger)):
      lower_weak = dtypes.is_weakly_typed(lower)
      upper_weak = dtypes.is_weakly_typed(upper)
      if lower_weak and not upper_weak:
        dtype = upper_dtype
      elif not lower_weak and upper_weak:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the function itself, not its result: fori_loop(0, n, body_fun, init) where body_fun is def body_fun(i, carry): ...
2. If using a lambda, ensure the signature is lambda i, carry: ... (two args, uninvoked)
3. Check for shadowing: verify no earlier assignment replaced body_fun with its call result

Example fix

// before
out = lax.fori_loop(0, 10, step(i, x), x)  # called!
// after
out = lax.fori_loop(0, 10, step, x)

import typing
def validate_fori_args(lower, upper, body_fun, init):
    if not callable(body_fun):
        raise TypeError('body_fun must be callable, got %r' % type(body_fun))
    return True

def is_body_fun(x) -> bool:
    return callable(x) and not isinstance(x, (jnp.ndarray, np.ndarray))

try:
    lax.fori_loop(0, n, body, init)
except TypeError as e:
    if 'body_fun argument should be callable' in str(e):
        raise  # programmer error: fix at call site

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• lower and upper arguments to fori_loop must have equal types
• iteration over a 0-d array
• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error
• `compute_on`'s compute_type argument must be a string.

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