jax-ml/jax · error · TypeError

cond_fun must return a boolean scalar, but got output type(s

Error message

cond_fun must return a boolean scalar, but got output type(s) {}.

What it means

Companion to the structural check: cond_fun of jax.lax.while_loop returned a single leaf, but that leaf is not a scalar boolean — its shape is non-empty or its dtype is not bool (e.g. float32, int32). Since the loop predicate must compile to a concrete branch, JAX requires an aval exactly equal to ShapedArray((), bool).

Source

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

      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
  # may not match the output despite being compatible by virtue of their weak type.
  # To do this, we compute the jaxpr in two passes: first with the raw inputs, and if
  # necessary, a second time with modified init values.
  cond_jaxpr, body_jaxpr, body_out_avals = _create_jaxpr(init_aval)
  if len(body_out_avals) != len(init_aval):
    _check_carry_type('while_loop body', body_fun, init_aval, body_out_avals)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the predicate explicitly a boolean scalar: lambda v: (v > 0) yields bool; use jnp.asarray(cond, jnp.bool_) or .astype(jnp.bool_) if needed
2. Squeeze/reshape array-valued conditions to scalar, e.g. cond.reshape(()) or bool(arr[0]) when semantically correct
3. For batched loops, use jnp.all()/jnp.any() to reduce a vector predicate to a scalar

Example fix

// before
def cond(v):
    return v['count']  # int scalar, not bool
jax.lax.while_loop(cond, body, init)

// after
def cond(v):
    return v['count'] > 0  # scalar bool
jax.lax.while_loop(cond, body, init)

import jax.numpy as jnp
pred = cond_fun(init_val)
assert hasattr(pred, 'shape') and getattr(pred, 'shape', ()) == () and jnp.asarray(pred).dtype == jnp.bool_, 'cond must be scalar bool'

def cond_is_scalar_bool(cond_fun, init_val) -> bool:
    p = jnp.asarray(cond_fun(init_val))
    return p.shape == () and p.dtype == jnp.bool_

try:
    jax.lax.while_loop(cond, body, init)
except TypeError as e:
    if 'output type(s)' in str(e):
        # coerce: cond = lambda v: jnp.asarray(cond(v), jnp.bool_)
        raise

When it happens

Trigger: cond_fun returning a numeric value like lambda v: n - v (int/float, not bool); returning a boolean array of shape (1,) or (batch,) instead of shape (); returning np.where(...) that yields non-bool dtype; returning a traced Python int.

Common situations: Writing cond as a countdown counter (non-negative int) instead of counter > 0; vmap-ed or batched predicates producing vector booleans; mixing NumPy scalars that keep float dtype; porting Python while n: idioms directly.

Related errors

• cond_fun must return a boolean scalar, but got pytree {}.
• Indexer must have integer or boolean type, got indexer with
• k must be a scalar integer; got {k}
• Only arrays with 32-bit element types can be converted to sc
• primal and tangent arguments to jax.jvp do not match; dtypes

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