{"record":{"id":"ee5bc5a28bdedf09","repo":"jax-ml/jax","slug":"lax-while-loop-body-fun-and-cond-fun-arguments-sh","errorCode":null,"errorMessage":"lax.while_loop: body_fun and cond_fun arguments should be callable.","messagePattern":"lax\\.while_loop: body_fun and cond_fun arguments should be callable\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":1660,"sourceCode":"\n  .. note::\n    :py:func:`while_loop` compiles ``cond_fun`` and ``body_fun``, so while it\n    can be combined with :py:func:`jit`, it's usually unnecessary.\n\n  Args:\n    cond_fun: function of type ``a -> Bool``.\n    body_fun: function of type ``a -> a``.\n    init_val: value of type ``a``, a type that can be a scalar, array, or any\n      pytree (nested Python tuple/list/dict) thereof, representing the initial\n      loop carry value.\n\n  Returns:\n    The output from the final iteration of body_fun, of type ``a``.\n\n  .. _Haskell-like type signature: https://wiki.haskell.org/Type_signature\n  \"\"\"\n  if not (callable(body_fun) and callable(cond_fun)):\n    raise TypeError(\"lax.while_loop: body_fun and cond_fun arguments should be callable.\")\n  if config.disable_jit.value:\n    try:\n      val = tree_map(lax.asarray, init_val)\n      while cond_fun(val):\n        val = tree_map(lax.asarray, body_fun(val))\n      return val\n    except core.ConcretizationTypeError:\n      # Can't run this while_loop in Python (e.g. because there's a vmap\n      # transformation on it), so we fall back to the primitive version.\n      pass\n\n  def _create_jaxpr(init_avals):\n    args_avals = ft.pack(((init_avals,), {}))\n    cond_jaxpr, cond_out_avals = pe.trace_to_jaxpr(cond_fun, args_avals, cond_dbg)\n    body_jaxpr, body_out_avals = pe.trace_to_jaxpr(body_fun, args_avals, body_dbg)\n    if not treedef_is_leaf(cond_out_avals.tree) or len(cond_jaxpr.out_avals) != 1:\n      msg = \"cond_fun must return a boolean scalar, but got pytree {}.\"\n      raise TypeError(msg.format(cond_out_avals.tree))","sourceCodeStart":1642,"sourceCodeEnd":1678,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L1642-L1678","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\njax.lax.while_loop(lambda v: v < 10, step(0.1), v)  # step(0.1) called, returns value\n\n// after\njax.lax.while_loop(lambda v: v < 10, step, v)  # pass the callable","handlingStrategy":"type-guard","validationCode":"assert callable(cond_fun) and callable(body_fun), 'cond_fun and body_fun must be callables'","typeGuard":"def is_valid_while_loop_args(cond_fun, body_fun) -> bool:\n    return callable(cond_fun) and callable(body_fun)","tryCatchPattern":"try:\n    jax.lax.while_loop(cond_fun, body_fun, init)\nexcept TypeError as e:\n    if 'should be callable' in str(e):\n        raise TypeError('pass function objects, not results: while_loop(cond, body, init)') from e\n    raise","preventionTips":["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"],"tags":["jax","while-loop","typeerror","callable","api-misuse"],"backgroundTag":"non-callable-argument-passed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}