jax.enable_checks

jax.enable_checks = <jax._src.config.State object>

Context manager for jax_enable_checks config option.

Turn on invariant checking for JAX internals. Makes things slower.

Parameters:

new_val (Any)