{"record":{"id":"c86b078fe3d6c475","repo":"jax-ml/jax","slug":"cond-fun-must-return-a-boolean-scalar-but-got-pyt","errorCode":null,"errorMessage":"cond_fun must return a boolean scalar, but got pytree {}.","messagePattern":"cond_fun must return a boolean scalar, but got pytree (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":1678,"sourceCode":"    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))\n\n    pred_aval = cond_jaxpr.out_avals[0]\n    if (not isinstance(pred_aval, ShapedArray)\n        or ShapedArray(pred_aval.shape, pred_aval.dtype) != ShapedArray((), np.bool_)):\n      msg = \"cond_fun must return a boolean scalar, but got output type(s) {}.\"\n      raise TypeError(msg.format(cond_jaxpr.out_avals))\n\n    return cond_jaxpr, body_jaxpr, body_out_avals\n\n  cond_dbg = api_util.debug_info(\"while_cond\", cond_fun, (init_val,), {})\n  body_dbg = api_util.debug_info(\"while_body\", body_fun, (init_val,), {})\n  init_val_flat = ft.flatten(init_val)\n  check_no_transformed_refs_args(lambda: body_dbg, init_val_flat.vals)\n  del init_val\n  init_aval = init_val_flat.map(core.typeof)\n\n  # The body input and output avals must match exactly. However, we want to account for\n  # the case when init contains weakly-typed values (e.g. Python scalars), with avals that","sourceCodeStart":1660,"sourceCodeEnd":1696,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L1660-L1696","documentation":"The condition function of jax.lax.while_loop (and fori_loop's internal while) must return exactly one value: a scalar boolean. This error is the structural check — cond_fun returned a pytree that is not a single leaf, such as a tuple, list, dict, or multiple values, so there is no single scalar predicate to branch on.","triggerScenarios":"cond_fun like lambda v: (v < 10, v > 0) or lambda v: {'stop': v < n} or returning two arrays; a cond that returns (done, state) designed for a different loop API; unpacking/returns added during refactoring.","commonSituations":"Combining multiple stopping criteria with a comma instead of & ; returning a 1-element list from a helper; migrating code from a framework whose step functions return (done, obs) tuples (e.g. RL environments); early-return refactorings that wrap the boolean in a tuple.","solutions":["Return a single boolean leaf: combine conditions with jnp.logical_and / & , e.g. lambda v: (v < 10) & (v > 0)","If cond_fun computes several things, keep only the boolean as the return and move other outputs into the body or close over them","Verify with a quick call: assert cond_jaxpr-style check via jax.make_jaxpr(cond)(init) has exactly one output aval"],"exampleFix":"// before\ndef cond(v):\n    return v < 10, v > 0  # tuple, not scalar\njax.lax.while_loop(cond, body, init)\n\n// after\ndef cond(v):\n    return (v < 10) & (v > 0)  # single scalar boolean\njax.lax.while_loop(cond, body, init)","handlingStrategy":"validation","validationCode":"import jax\nout = cond_fun(init_val)\nassert jax.tree_util.tree_structure(out).num_leaves == 1 and not hasattr(out, '__len__') or jnp.isscalar(out), 'cond must return one leaf'","typeGuard":"def cond_returns_scalar(cond_fun, init_val) -> bool:\n    t = jax.tree_util.tree_structure(cond_fun(init_val))\n    return t.num_leaves == 1 and jax.tree_util.treedef_is_leaf(t)","tryCatchPattern":"try:\n    jax.lax.while_loop(cond, body, init)\nexcept TypeError as e:\n    if 'boolean scalar' in str(e) and 'pytree' in str(e):\n        # combine multiple conditions with & and retry\n        raise","preventionTips":["Write cond as a single boolean expression; combine criteria with jnp.logical_and","Avoid returning dicts/tuples from cond helpers","Unit-test cond(init) returns a Python/DeviceArray bool scalar"],"tags":["jax","while-loop","cond-fun","pytree","boolean-scalar"],"backgroundTag":"loop-condition-not-scalar-boolean","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}