{"record":{"id":"8d7f48e621bc3698","repo":"jax-ml/jax","slug":"prim-name-takes-a-scalar-pred-as-argument-got","errorCode":null,"errorMessage":"{prim_name} takes a scalar pred as argument, got {pred}","messagePattern":"(.+?) takes a scalar pred as argument, got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/checkify.py","lineNumber":1295,"sourceCode":"    >>> import jax.numpy as jnp\n    >>> from jax.experimental import checkify\n    >>> def f(x):\n    ...   checkify.check(x>0, \"{x} needs to be positive!\", x=x)\n    ...   return 1/x\n    >>> checked_f = checkify.checkify(f)\n    >>> err, out = jax.jit(checked_f)(-3.)\n    >>> err.throw()  # doctest: +IGNORE_EXCEPTION_DETAIL\n    Traceback (most recent call last):\n      ...\n    jax._src.checkify.JaxRuntimeError: -3. needs to be positive!\n\n  \"\"\"\n  _check(pred, msg, debug, *fmt_args, **fmt_kwargs)\n\ndef _check(pred, msg, debug, *fmt_args, **fmt_kwargs):\n  if not is_scalar_pred(pred):\n    prim_name = 'debug_check' if debug else 'check'\n    raise TypeError(f'{prim_name} takes a scalar pred as argument, got {pred}')\n  for arg in jtu.tree_leaves((fmt_args, fmt_kwargs)):\n    if not isinstance(arg, (Array, np.ndarray)):\n      raise TypeError('Formatting arguments to checkify.check need to be '\n                      'PyTrees of arrays, but got '\n                      f'{arg!r} of type {type(arg)}.')\n  new_error = FailedCheckError(get_traceback(), msg, *fmt_args, **fmt_kwargs)\n  error = assert_func(init_error, jnp.logical_not(pred), new_error)\n  _check_error(error, debug=debug)\n\ndef _check_error(error, *, debug=False):\n  if any(map(np.shape, error._pred.values())):\n    error = _reduce_any_error(error)\n  err_args, tree_def = tree_flatten(error)\n\n  return check_p.bind(*err_args, err_tree=tree_def, debug=debug)\n\n\ndef is_scalar_pred(pred) -> bool:","sourceCodeStart":1277,"sourceCodeEnd":1313,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/checkify.py#L1277-L1313","documentation":"checkify.check / debug_check require a scalar boolean predicate; anything else (array, tracer with shape, non-bool) is rejected with this TypeError.","triggerScenarios":"Passing a non-scalar pred to checkify.check, e.g. an array of booleans, an int, or a shaped tracer instead of a scalar bool.","commonSituations":"Writing check(x > 0) where x is an array (author intended .all() or .all(axis=...)); passing Python bools of arrays; using debug_check with vectorized conditions.","solutions":["Reduce the predicate to a scalar: check((x > 0).all()) or pick a specific element","Verify the pred is boolean — add an explicit cast like jnp.asarray(pred).astype(bool) if needed","For per-element diagnostics, check each invariant separately or format values into the message"],"exampleFix":"# before\ncheckify.check(x > 0, 'x positive', x)  # x is an array\n# after\ncheckify.check((x > 0).all(), 'x positive: {}', x)","handlingStrategy":"type-guard","validationCode":"def is_scalar_pred(pred) -> bool:\n    return isinstance(pred, (bool, np.bool_)) or (hasattr(pred, 'shape') and getattr(pred, 'shape', None) == () and str(getattr(pred, 'dtype', '')).startswith('bool'))","typeGuard":"def is_scalar_pred(pred) -> bool:\n    a = jnp.asarray(pred)\n    return a.shape == () and a.dtype == jnp.bool_","tryCatchPattern":null,"preventionTips":["Always reduce array predicates with .all()/.any()","Write a small assert helper check_scalar(pred) in test suites"],"tags":["jax","checkify","type-error","scalar-predicate"],"backgroundTag":"type-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}