jax.checking_leaks = functools.partial(<jax._src.config._StateContextManager object>, True)#

Context manager for jax_check_tracer_leaks config option.

Turn on checking for leaked tracers as soon as a trace completes. Enabling leak checking may have performance impacts: some caching is disabled, and other overheads may be added. Additionally, be aware that some Python debuggers can cause false positives, so it is recommended to disable any debuggers while leak checking is enabled.