jax-ml/jax · error · TypeError

check_error takes an Error as argument, got type

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.

Solutions

Use the first element of the tuple returned by checkify: error, checked_fn = checkify(f); ...; checkify.check_error(error)
Verify you're not passing the function output, None, or the partial function
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)

Defensive patterns

Strategy: type-guard

Type guard

from jax._src.checkify import Error, check_error

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

Try / catch

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

Prevention

Always unpack: error, checked_fn = checkify(f) and call check_error(error)
Name the tuple element 'error' at every call site to avoid mixing it up with outputs

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

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

Appendix: 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)