jax-ml/jax · error · TypeError

check_error takes an Error as argument, got type {type(error

Error message

check_error takes an Error as argument, got type {type(error)} instead.

What it means

checkify.check_error(error) is the post-run step that surfaces any failed checks collected by checkify; it only accepts an instance of checkify.Error (the first element of the tuple returned by checkify()). The source isinstance-checks the argument and raises TypeError otherwise. This commonly happens when callers pass the function result or the second tuple element instead of the Error object.

Source

Thrown at jax/_src/checkify.py:1418

  >>> def f(x):
  ...   checkify.check(x>0, "must be positive!")
  ...   return x
  >>> def with_inner_jit(x):
  ...   checked_f = checkify.checkify(f)
  ...   # a checkified function can be jitted
  ...   error, out = jax.jit(checked_f)(x)
  ...   checkify.check_error(error)
  ...   return out
  >>> _ = with_inner_jit(1)  # no failed check
  >>> with_inner_jit(-1)  # doctest: +IGNORE_EXCEPTION_DETAIL
  Traceback (most recent call last):
    ...
  jax._src.JaxRuntimeError: must be positive!
  >>> # can re-checkify
  >>> error, _ = checkify.checkify(with_inner_jit)(-1)
  """
  if not isinstance(error, Error):
    raise TypeError('check_error takes an Error as argument, '
                     f'got type {type(error)} instead.')
  _check_error(error, debug=False)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the first element of the tuple returned by checkify: error, checked_fn = checkify(f); ...; checkify.check_error(error)
2. Verify you're not passing the function output, None, or the partial function
3. If re-checkifying a jitted function, thread the outer and inner Error values through explicitly as shown in the checkify docs

Example fix

// before
error, checked_fn = checkify(f)
out = checked_fn(args)
checkify.check_error(out)  # wrong: that's f's return
// after
error, checked_fn = checkify(f)
out = checked_fn(args)
checkify.check_error(error)

from jax._src.checkify import Error, check_error

def is_error_obj(x) -> bool:
    return isinstance(x, Error)

try:
    checkify.check_error(error)
except TypeError as e:
    raise TypeError(f'pass the Error from checkify(), not {type(error)}') from e

When it happens

Trigger: Calling check_error on something that isn't an Error: e.g. err, _ = checkify(f); check_error(f_result) — passing the function's return value instead of err, or unpacking the tuple in the wrong order, or passing None after forgetting that checkify returns (Error, functools.partial).

Common situations: New users assume check_error takes the checked function's output. Also re-checkifying: after error is transformed, passing the wrapped function or a jax.Array where the Error should be.

Related errors

• Formatting arguments to checkify.check need to be PyTrees of
• lax.scan: f argument should be a callable.
• lax.while_loop: body_fun and cond_fun arguments should be ca
• {name} wrapped function must be passed at least one argument
• primal and tangent arguments to jax.jvp must be tuples or li

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