{"record":{"id":"88db1fc2d2f17ed0","repo":"jax-ml/jax","slug":"cond-fun-must-return-a-boolean-scalar-but-got-out","errorCode":null,"errorMessage":"cond_fun must return a boolean scalar, but got output type(s) {}.","messagePattern":"cond_fun must return a boolean scalar, but got output type\\(s\\) (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":1684,"sourceCode":"      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\n  # may not match the output despite being compatible by virtue of their weak type.\n  # To do this, we compute the jaxpr in two passes: first with the raw inputs, and if\n  # necessary, a second time with modified init values.\n  cond_jaxpr, body_jaxpr, body_out_avals = _create_jaxpr(init_aval)\n  if len(body_out_avals) != len(init_aval):\n    _check_carry_type('while_loop body', body_fun, init_aval, body_out_avals)","sourceCodeStart":1666,"sourceCodeEnd":1702,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L1666-L1702","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","Squeeze/reshape array-valued conditions to scalar, e.g. cond.reshape(()) or bool(arr[0]) when semantically correct","For batched loops, use jnp.all()/jnp.any() to reduce a vector predicate to a scalar"],"exampleFix":"// before\ndef cond(v):\n    return v['count']  # int scalar, not bool\njax.lax.while_loop(cond, body, init)\n\n// after\ndef cond(v):\n    return v['count'] > 0  # scalar bool\njax.lax.while_loop(cond, body, init)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\npred = cond_fun(init_val)\nassert hasattr(pred, 'shape') and getattr(pred, 'shape', ()) == () and jnp.asarray(pred).dtype == jnp.bool_, 'cond must be scalar bool'","typeGuard":"def cond_is_scalar_bool(cond_fun, init_val) -> bool:\n    p = jnp.asarray(cond_fun(init_val))\n    return p.shape == () and p.dtype == jnp.bool_","tryCatchPattern":"try:\n    jax.lax.while_loop(cond, body, init)\nexcept TypeError as e:\n    if 'output type(s)' in str(e):\n        # coerce: cond = lambda v: jnp.asarray(cond(v), jnp.bool_)\n        raise","preventionTips":["Always write conditions as comparisons (>, <, ==) which yield bool, never bare counters","Reduce vector predicates with jnp.all/jnp.any","Cast ambiguous predicates with .astype(jnp.bool_) and reshape to () in batched code"],"tags":["jax","while-loop","cond-fun","dtype","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"}