jax-ml/jax · error · TypeError

Formatting arguments to checkify.check need to be PyTrees of

Error message

Formatting arguments to checkify.check need to be PyTrees of arrays, but got {arg!r} of type {type(arg)}.

What it means

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.

Source

Thrown at jax/_src/checkify.py:1298

    ...   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:
  return (isinstance(pred, bool) or
          isinstance(pred, Array) and pred.shape == () and
          pred.dtype == np.dtype('bool'))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the value to an array before passing: jnp.asarray(value) or np.asarray(value)
2. Pre-format the message yourself: check(pred, f'x={value}') with no fmt args
3. Ensure every leaf of any PyTree argument (dicts/tuples) is an Array or np.ndarray
4. If you truly need a Python-side check, use checkify.check only for array-valued data and hoist scalar metadata into the f-string

Example fix

// before
checkify.check(x > 0, 'x must be > 0, got {}', x_scalar_python)
// after
checkify.check(x > 0, 'x must be > 0, got {}', jnp.asarray(x_scalar_python))

import jax.numpy as jnp
import numpy as np
from jax._src import checkify

def all_array_leaves(tree) -> bool:
    return all(
        isinstance(leaf, (jnp.ndarray.__mro__[0], np.ndarray))
        if hasattr(jnp, 'ndarray') else isinstance(leaf, (np.ndarray,))
        for leaf in jax.tree_util.tree_leaves(tree)
    )

# simpler: try conversion first
def safe_args(args):
    return jax.tree_util.tree_map(
        lambda x: x if isinstance(x, np.ndarray) else jnp.asarray(x), args)

def is_checkify_fmt_safe(*args, **kwargs) -> bool:
    import numpy as np, jax
    return all(
        isinstance(l, (np.ndarray,)) or type(l).__name__ == 'Array'
        for l in jax.tree_util.tree_leaves((args, kwargs))
    )

When it happens

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

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

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• check_error takes an Error as argument, got type {type(error
• iteration over a 0-d array
• {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/acbc8cc6c3f72c63. Report an issue: GitHub.