jax-ml/jax · error · TypeError

{prim_name} takes a scalar pred as argument, got {pred}

Error message

{prim_name} takes a scalar pred as argument, got {pred}

What it means

checkify.check / debug_check require a scalar boolean predicate; anything else (array, tracer with shape, non-bool) is rejected with this TypeError.

Source

Thrown at jax/_src/checkify.py:1295

    >>> import jax.numpy as jnp
    >>> from jax.experimental import checkify
    >>> def f(x):
    ...   checkify.check(x>0, "{x} needs to be positive!", x=x)
    ...   return 1/x
    >>> checked_f = checkify.checkify(f)
    >>> err, out = jax.jit(checked_f)(-3.)
    >>> err.throw()  # doctest: +IGNORE_EXCEPTION_DETAIL
    Traceback (most recent call last):
      ...
    jax._src.checkify.JaxRuntimeError: -3. needs to be positive!

  """
  _check(pred, msg, debug, *fmt_args, **fmt_kwargs)

def _check(pred, msg, debug, *fmt_args, **fmt_kwargs):
  if not is_scalar_pred(pred):
    prim_name = 'debug_check' if debug else 'check'
    raise TypeError(f'{prim_name} takes a scalar pred as argument, got {pred}')
  for arg in jtu.tree_leaves((fmt_args, fmt_kwargs)):
    if not isinstance(arg, (Array, np.ndarray)):
      raise TypeError('Formatting arguments to checkify.check need to be '
                      'PyTrees of arrays, but got '
                      f'{arg!r} of type {type(arg)}.')
  new_error = FailedCheckError(get_traceback(), msg, *fmt_args, **fmt_kwargs)
  error = assert_func(init_error, jnp.logical_not(pred), new_error)
  _check_error(error, debug=debug)

def _check_error(error, *, debug=False):
  if any(map(np.shape, error._pred.values())):
    error = _reduce_any_error(error)
  err_args, tree_def = tree_flatten(error)

  return check_p.bind(*err_args, err_tree=tree_def, debug=debug)


def is_scalar_pred(pred) -> bool:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the predicate to a scalar: check((x > 0).all()) or pick a specific element
2. Verify the pred is boolean — add an explicit cast like jnp.asarray(pred).astype(bool) if needed
3. For per-element diagnostics, check each invariant separately or format values into the message

Example fix

# before
checkify.check(x > 0, 'x positive', x)  # x is an array
# after
checkify.check((x > 0).all(), 'x positive: {}', x)

def is_scalar_pred(pred) -> bool:
    return isinstance(pred, (bool, np.bool_)) or (hasattr(pred, 'shape') and getattr(pred, 'shape', None) == () and str(getattr(pred, 'dtype', '')).startswith('bool'))

def is_scalar_pred(pred) -> bool:
    a = jnp.asarray(pred)
    return a.shape == () and a.dtype == jnp.bool_

When it happens

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

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

Related errors

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {str(exc)}
• Checkify does not support batched while-loops (checkify-of-v
• Mesh must be provided for shard_map with checkify.

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