jax.enable_checks

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

Context manager for jax_enable_checks config option.

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