jax-ml/jax · error · TypeError

cond_fun must return a boolean scalar, but got pytree {}.

Error message

cond_fun must return a boolean scalar, but got pytree {}.

What it means

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.

Source

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

    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))

    pred_aval = cond_jaxpr.out_avals[0]
    if (not isinstance(pred_aval, ShapedArray)
        or ShapedArray(pred_aval.shape, pred_aval.dtype) != ShapedArray((), np.bool_)):
      msg = "cond_fun must return a boolean scalar, but got output type(s) {}."
      raise TypeError(msg.format(cond_jaxpr.out_avals))

    return cond_jaxpr, body_jaxpr, body_out_avals

  cond_dbg = api_util.debug_info("while_cond", cond_fun, (init_val,), {})
  body_dbg = api_util.debug_info("while_body", body_fun, (init_val,), {})
  init_val_flat = ft.flatten(init_val)
  check_no_transformed_refs_args(lambda: body_dbg, init_val_flat.vals)
  del init_val
  init_aval = init_val_flat.map(core.typeof)

  # The body input and output avals must match exactly. However, we want to account for
  # the case when init contains weakly-typed values (e.g. Python scalars), with avals that

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return a single boolean leaf: combine conditions with jnp.logical_and / & , e.g. lambda v: (v < 10) & (v > 0)
2. If cond_fun computes several things, keep only the boolean as the return and move other outputs into the body or close over them
3. Verify with a quick call: assert cond_jaxpr-style check via jax.make_jaxpr(cond)(init) has exactly one output aval

Example fix

// before
def cond(v):
    return v < 10, v > 0  # tuple, not scalar
jax.lax.while_loop(cond, body, init)

// after
def cond(v):
    return (v < 10) & (v > 0)  # single scalar boolean
jax.lax.while_loop(cond, body, init)

import jax
out = cond_fun(init_val)
assert jax.tree_util.tree_structure(out).num_leaves == 1 and not hasattr(out, '__len__') or jnp.isscalar(out), 'cond must return one leaf'

def cond_returns_scalar(cond_fun, init_val) -> bool:
    t = jax.tree_util.tree_structure(cond_fun(init_val))
    return t.num_leaves == 1 and jax.tree_util.treedef_is_leaf(t)

try:
    jax.lax.while_loop(cond, body, init)
except TypeError as e:
    if 'boolean scalar' in str(e) and 'pytree' in str(e):
        # combine multiple conditions with & and retry
        raise

When it happens

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

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

Related errors

• {} function carry input and carry output must have the same
• cond_fun must return a boolean scalar, but got output type(s
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

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