JAX configuration#

config

check_tracer_leaks

Context manager for jax_check_tracer_leaks config option.

checking_leaks

Context manager for jax_check_tracer_leaks config option.

debug_nans

Context manager for jax_debug_nans config option.

debug_infs

Context manager for jax_debug_infs 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.

enable_checks

Context manager for jax_enable_checks config option.

enable_custom_prng

Context manager for jax_enable_custom_prng config option (transient).

enable_custom_vjp_by_custom_transpose

Context manager for jax_enable_custom_vjp_by_custom_transpose config option (transient).

log_compiles

Context manager for jax_log_compiles config option.

numpy_rank_promotion

Context manager for jax_numpy_rank_promotion config option.

transfer_guard(new_val)

Set up thread-local state and return a contextmanager for managing it.