Checkify

Reference-aware version of jax.experimental.checkify. Checks for errors (like NaN/Inf) during execution.

Abstract Signature:

Checkify(f: Callable, errors: frozenset = frozenset({<class 'jax._src.checkify.FailedCheckError'>}))

PyTorchJAX (Core)KerasTensorFlowApple MLXFlax NNXPaxML / Praxis

PyTorch

API: —

Strategy: Custom / Partial

JAX (Core)

API: jax.experimental.checkify.checkify

Strategy: Direct Mapping

Official Docs ↗

Keras

API: —

Strategy: Custom / Partial

TensorFlow

API: tf.debugging.check_numerics

Strategy: Direct Mapping

Official Docs ↗

Apple MLX

API: —

Strategy: Custom / Partial

Flax NNX

API: flax.nnx.checkify

Strategy: Direct Mapping

Official Docs ↗

PaxML / Praxis

API: jax.experimental.checkify.checkify

Strategy: Direct Mapping

Official Docs ↗