jax-ml/jax · error · TypeError

lax.while_loop: body_fun and cond_fun arguments should be ca

Error message

lax.while_loop: body_fun and cond_fun arguments should be callable.

What it means

jax.lax.while_loop (and fori_loop, which builds on it) requires both cond_fun and body_fun to be Python callables. This guard catches passing non-callable values such as None, arrays, ints, or the result of calling a function instead of the function itself, before tracing begins.

Source

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

  .. note::
    :py:func:`while_loop` compiles ``cond_fun`` and ``body_fun``, so while it
    can be combined with :py:func:`jit`, it's usually unnecessary.

  Args:
    cond_fun: function of type ``a -> Bool``.
    body_fun: function of type ``a -> a``.
    init_val: value of type ``a``, a type that can be a scalar, array, or any
      pytree (nested Python tuple/list/dict) thereof, representing the initial
      loop carry value.

  Returns:
    The output from the final iteration of body_fun, of type ``a``.

  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature
  """
  if not (callable(body_fun) and callable(cond_fun)):
    raise TypeError("lax.while_loop: body_fun and cond_fun arguments should be callable.")
  if config.disable_jit.value:
    try:
      val = tree_map(lax.asarray, init_val)
      while cond_fun(val):
        val = tree_map(lax.asarray, body_fun(val))
      return val
    except core.ConcretizationTypeError:
      # Can't run this while_loop in Python (e.g. because there's a vmap
      # transformation on it), so we fall back to the primitive version.
      pass

  def _create_jaxpr(init_avals):
    args_avals = ft.pack(((init_avals,), {}))
    cond_jaxpr, cond_out_avals = pe.trace_to_jaxpr(cond_fun, args_avals, cond_dbg)
    body_jaxpr, body_out_avals = pe.trace_to_jaxpr(body_fun, args_avals, body_dbg)
    if not treedef_is_leaf(cond_out_avals.tree) or len(cond_jaxpr.out_avals) != 1:
      msg = "cond_fun must return a boolean scalar, but got pytree {}."
      raise TypeError(msg.format(cond_out_avals.tree))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass the function objects themselves: jax.lax.while_loop(cond, body, init) with no parentheses after cond/body
If you meant counted iteration, use jax.lax.fori_loop(lower, upper, body_fun, init_val) instead
Check for variables shadowing cond_fun/body_fun with non-callable values

Example fix

// before
jax.lax.while_loop(lambda v: v < 10, step(0.1), v)  # step(0.1) called, returns value

// after
jax.lax.while_loop(lambda v: v < 10, step, v)  # pass the callable

Defensive patterns

Strategy: type-guard

Validate before calling

assert callable(cond_fun) and callable(body_fun), 'cond_fun and body_fun must be callables'

Type guard

def is_valid_while_loop_args(cond_fun, body_fun) -> bool:
    return callable(cond_fun) and callable(body_fun)

Try / catch

try:
    jax.lax.while_loop(cond_fun, body_fun, init)
except TypeError as e:
    if 'should be callable' in str(e):
        raise TypeError('pass function objects, not results: while_loop(cond, body, init)') from e
    raise

Prevention

Never add parentheses when passing cond/body
Prefer fori_loop for counted loops to avoid signature confusion
Lint for lambda/function-typed parameters in loop wrappers

When it happens

Trigger: Calling jax.lax.while_loop(cond_fun, body_fun, init_val) with body_fun=None (e.g. using fori-style arguments by mistake); passing a jitted/called result like body_fun(init) instead of body_fun; typos where a variable shadowing the function holds a number or array.

Common situations: Confusing fori_loop(lower, upper, body_fun, init_val) with while_loop's signature; refactoring that accidentally invokes the function; passing a bound method result or partial object incorrectly; copy-paste between APIs.

Related errors

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