{"record":{"id":"acbc8cc6c3f72c63","repo":"jax-ml/jax","slug":"formatting-arguments-to-checkify-check-need-to-be","errorCode":null,"errorMessage":"Formatting arguments to checkify.check need to be PyTrees of arrays, but got {arg!r} of type {type(arg)}.","messagePattern":"Formatting arguments to checkify\\.check need to be PyTrees of arrays, but got (.+?) of type (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/checkify.py","lineNumber":1298,"sourceCode":"    ...   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:\n  return (isinstance(pred, bool) or\n          isinstance(pred, Array) and pred.shape == () and\n          pred.dtype == np.dtype('bool'))","sourceCodeStart":1280,"sourceCodeEnd":1316,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/checkify.py#L1280-L1316","documentation":"checkify.check formats an error message from extra positional/keyword arguments when the predicate fails, but every leaf of those argument PyTrees must be a JAX Array or numpy ndarray. The source iterates jtu.tree_leaves over (fmt_args, fmt_kwargs) and raises TypeError for any leaf that isn't an array. This keeps the check traceable/JIT-compatible, since only arrays can flow through JAX primitives.","triggerScenarios":"Calling checkify.check(pred, \"msg {}\", value) (or passing fmt kwargs) where value (or any leaf of a nested PyTree argument) is a Python scalar, str, list of non-arrays, or any non-array object. Example: check(False, 'x={}', 'hello') or check(pred, '{v}', v=some_python_int).","commonSituations":"Developers migrating from Python assert or absl assertions pass plain ints/strings as format args. Also passing a Python bool or float computed outside jit, or a dict containing mixed array/non-array leaves.","solutions":["Convert the value to an array before passing: jnp.asarray(value) or np.asarray(value)","Pre-format the message yourself: check(pred, f'x={value}') with no fmt args","Ensure every leaf of any PyTree argument (dicts/tuples) is an Array or np.ndarray","If you truly need a Python-side check, use checkify.check only for array-valued data and hoist scalar metadata into the f-string"],"exampleFix":"// before\ncheckify.check(x > 0, 'x must be > 0, got {}', x_scalar_python)\n// after\ncheckify.check(x > 0, 'x must be > 0, got {}', jnp.asarray(x_scalar_python))","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nimport numpy as np\nfrom jax._src import checkify\n\ndef all_array_leaves(tree) -> bool:\n    return all(\n        isinstance(leaf, (jnp.ndarray.__mro__[0], np.ndarray))\n        if hasattr(jnp, 'ndarray') else isinstance(leaf, (np.ndarray,))\n        for leaf in jax.tree_util.tree_leaves(tree)\n    )\n\n# simpler: try conversion first\ndef safe_args(args):\n    return jax.tree_util.tree_map(\n        lambda x: x if isinstance(x, np.ndarray) else jnp.asarray(x), args)","typeGuard":"def is_checkify_fmt_safe(*args, **kwargs) -> bool:\n    import numpy as np, jax\n    return all(\n        isinstance(l, (np.ndarray,)) or type(l).__name__ == 'Array'\n        for l in jax.tree_util.tree_leaves((args, kwargs))\n    )","tryCatchPattern":null,"preventionTips":["Pre-format messages with f-strings when values are Python scalars/strings","Always wrap non-array values in jnp.asarray before passing as fmt args","Remember bools/ints/strs are not arrays even though they're valid PyTree leaves"],"tags":["jax","checkify","typeerror","assertions","jit"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}