JAX configuration

config

enable_checks

Context manager for jax_enable_checks config option.

check_tracer_leaks

Context manager for jax_check_tracer_leaks config option.

checking_leaks

Context manager for jax_check_tracer_leaks config option.

enable_custom_prng

Context manager for jax_enable_custom_prng config option (transient).

debug_nans

Context manager for jax_debug_nans config option.

debug_infs

Context manager for jax_debug_infs config option.

log_compiles

Context manager for jax_log_compiles config option.

default_matmul_precision

Context manager for jax_default_matmul_precision config option.

default_prng_impl

Context manager for jax_default_prng_impl config option.

numpy_rank_promotion

Context manager for jax_numpy_rank_promotion config option.