unify configuration state handling

jax/__init__.py

@@ -30,7 +30,8 @@
 del _cloud_tpu_init
 
 # flake8: noqa: F401
-from .config import config
+from .config import (config, enable_checks, check_tracer_leaks, checking_leaks,
+                     debug_nans, debug_infs, log_compiles)
 from .api import (
   ad,  # TODO(phawkins): update users to avoid this.
   argnums_partial,  # TODO(phawkins): update Haiku to not use this.

jax/_src/lax/control_flow.py

@@ -1113,7 +1113,7 @@ def _cond_typecheck(*avals, branches, linear):
       f'called with operands of type {_avals_short(op_avals)}')
 
 def cond_bind(*args, branches, linear):
-  if not core.skip_checks:
+  if config.jax_enable_checks:
     avals = _map(core.get_aval, args)
     _cond_typecheck(*avals, branches=branches, linear=linear)
     for jaxpr in branches:
@@ -1876,7 +1876,7 @@ def _scan_typecheck(bind_time, *avals, reverse, length, num_consts, num_carry,
       f'called with sequence of type\n{_avals_short(x_avals)}')
 
 def scan_bind(*args, **params):
-  if not core.skip_checks:
+  if config.jax_enable_checks:
     avals = _map(core.get_aval, args)
     _scan_typecheck(True, *avals, **params)
     core.check_jaxpr(params['jaxpr'].jaxpr)

jax/_src/util.py

@@ -21,7 +21,6 @@
 
 import numpy as np
 
-import jax
 from jax.config import config
 
 partial = functools.partial
@@ -192,7 +191,7 @@ def cached(_, *args, **kwargs):
 
     @functools.wraps(f)
     def wrapper(*args, **kwargs):
-      if jax.core.debug_state.check_leaks:
+      if config.jax_check_tracer_leaks:
         return f(*args, **kwargs)
       else:
         return cached(bool(config.x64_enabled), *args, **kwargs)

jax/api.py

@@ -41,7 +41,7 @@
 from . import linear_util as lu
 from . import ad_util
 from . import dtypes
-from .core import eval_jaxpr, checking_leaks
+from .core import eval_jaxpr
 from .api_util import (flatten_fun, apply_flat_fun, flatten_fun_nokwargs,
                        flatten_fun_nokwargs2, argnums_partial,
                        argnums_partial_except, flatten_axes, donation_vector,
@@ -362,7 +362,7 @@ def f_jitted(*args, **kwargs):
       context = (getattr(core.thread_local_state.trace_state.trace_stack,
                          "dynamic", None), config.x64_enabled)
       # TODO(jblespiau): Move this to C++.
-      if (FLAGS.jax_debug_nans or FLAGS.jax_debug_infs) and not _jit_is_disabled():
+      if (config.jax_debug_nans or config.jax_debug_infs) and not _jit_is_disabled():
         device_arrays = cpp_jitted_f(context, *args, **kwargs)
         try:
           xla.check_special(xla.xla_call_p, [
@@ -372,7 +372,7 @@ def f_jitted(*args, **kwargs):
           ])
           return device_arrays
         except FloatingPointError:
-          assert FLAGS.jax_debug_nans or FLAGS.jax_debug_infs  # compiled_fun can only raise in this case
+          assert config.jax_debug_nans or config.jax_debug_infs  # compiled_fun can only raise in this case
           print("Invalid nan value encountered in the output of a C++-jit "
                 "function. Calling the de-optimized version.")
           return cache_miss(*args, **kwargs)[0]  # probably won't return
@@ -389,7 +389,7 @@ def f_jitted(*args, **kwargs):
     @api_boundary
     def f_jitted(*args, **kwargs):
       # TODO(jblespiau): Move this to C++.
-      if (FLAGS.jax_debug_nans or FLAGS.jax_debug_infs) and not _jit_is_disabled():
+      if (config.jax_debug_nans or config.jax_debug_infs) and not _jit_is_disabled():
         device_arrays = cpp_jitted_f(*args, **kwargs)
         try:
           xla.check_special(xla.xla_call_p, [
@@ -399,7 +399,7 @@ def f_jitted(*args, **kwargs):
           ])
           return device_arrays
         except FloatingPointError:
-          assert FLAGS.jax_debug_nans or FLAGS.jax_debug_infs  # compiled_fun can only raise in this case
+          assert config.jax_debug_nans or config.jax_debug_infs  # compiled_fun can only raise in this case
           print("Invalid nan value encountered in the output of a C++-jit "
                 "function. Calling the de-optimized version.")
           return cache_miss(*args, **kwargs)[0]  # probably won't return

jax/config.py

@@ -12,8 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
+import functools
 import os
 import sys
+import threading
+
 from jax import lib
 
 def bool_env(varname: str, default: bool) -> bool:
@@ -42,11 +46,16 @@ def int_env(varname: str, default: int) -> int:
 
 
 class Config:
+  _HAS_DYNAMIC_ATTRIBUTES = True
+
   def __init__(self):
     self.values = {}
     self.meta = {}
     self.FLAGS = NameSpace(self.read)
     self.use_absl = False
+    self._contextmanager_flags = set()
+
+    # TODO(mattjj): delete these when only omnistaging is available
     self.omnistaging_enabled = bool_env('JAX_OMNISTAGING', True)
     self._omnistaging_disablers = []
 
@@ -65,6 +74,13 @@ def update(self, name, val):
       lib.jax_jit.global_state().enable_x64 = val
 
   def read(self, name):
+    if name in self._contextmanager_flags:
+      raise AttributeError(
+          "For flags with a corresponding contextmanager, read their value "
+          f"via e.g. `config.{name}` rather than `config.FLAGS.{name}`.")
+    return self._read(name)
+
+  def _read(self, name):
     if self.use_absl:
       return getattr(self.absl_flags.FLAGS, name)
     else:
@@ -143,14 +159,82 @@ def disable_omnistaging(self):
         disabler()
       self.omnistaging_enabled = False
 
-  @property
-  def x64_enabled(self):
-    return lib.jax_jit.get_enable_x64()
-
-  # TODO(jakevdp): make this public when thread-local x64 is fully implemented.
-  def _set_x64_enabled(self, state):
-    lib.jax_jit.thread_local_state().enable_x64 = bool(state)
-
+#   # TODO(jakevdp, mattjj): unify this with `define_bool_state` stuff below
+#   @property
+#   def x64_enabled(self):
+#     return lib.jax_jit.get_enable_x64()
+
+#   def _set_x64_enabled(self, state):
+#     lib.jax_jit.thread_local_state().enable_x64 = bool(state)
+
+  def define_bool_state(self, name: str, default: bool, help: str):
+    """Set up thread-local state and return a contextmanager for managing it.
+
+    This function is a convenience wrapper. It defines a flag and corresponding
+    thread-local state, which can be managed via the contextmanager it returns.
+
+    The thread-local state value can be read via the ``config.<option_name>``
+    attribute, where ``config`` is the singleton ``Config`` instance.
+
+    Args:
+      name: string, converted to lowercase to define the name of the config
+        option (and absl flag). It is converted to uppercase to define the
+        corresponding shell environment variable.
+      default: boolean, a default value for the option.
+      help: string, used to populate the flag help information as well as the
+        docstring of the returned context manager.
+
+    Returns:
+      A contextmanager to control the thread-local state value.
+
+    Example:
+
+      enable_foo = config.define_bool_state(
+          name='jax_enable_foo',
+          default=False,
+          help='Enable foo.')
+
+      # Now the JAX_ENABLE_FOO shell environment variable and --jax_enable_foo
+      # command-line flag can be used to control the process-level value of
+      # the configuration option, in addition to using e.g.
+      # ``config.update("jax_enable_foo", True)`` directly. We can also use a
+      # context manager:
+
+      with enable_foo(True):
+        ...
+
+    The value of the thread-local state or flag can be accessed via
+    ``config.jax_enable_foo``. Reading it via ``config.FLAGS.jax_enable_foo`` is
+    an error.
+    """
+    name = name.lower()
+    self.DEFINE_bool(name, bool_env(name.upper(), default), help)
+    self._contextmanager_flags.add(name)
+
+    def get_state(self):
+      val = getattr(_thread_local_state, name, unset)
+      return val if val is not unset else self._read(name)
+    setattr(Config, name, property(get_state))
+
+    @contextlib.contextmanager
+    def set_state(new_val: bool):
+      prev_val = getattr(_thread_local_state, name, unset)
+      setattr(_thread_local_state, name, new_val)
+      try:
+        yield
+      finally:
+        if prev_val is unset:
+          delattr(_thread_local_state, name)
+        else:
+          setattr(_thread_local_state, name, prev_val)
+    set_state.__name__ = name[4:] if name.startswith('jax_') else name
+    set_state.__doc__ = f"Context manager for `{name}` config option.\n\n{help}"
+    return set_state
+
+_thread_local_state = threading.local()
+
+class Unset: pass
+unset = Unset()
 
 class NameSpace(object):
   def __init__(self, getter):
@@ -166,11 +250,6 @@ def __getattr__(self, name):
 
 already_configured_with_absl = False
 
-flags.DEFINE_bool(
-    'jax_enable_checks',
-    bool_env('JAX_ENABLE_CHECKS', False),
-    help='Turn on invariant checking (core.skip_checks = False)'
-)
 
 flags.DEFINE_bool(
     'jax_omnistaging',
@@ -184,14 +263,6 @@ def __getattr__(self, name):
     help='Set the number of stack frames in JAX tracer error messages.'
 )
 
-flags.DEFINE_bool(
-    'jax_check_tracer_leaks',
-    bool_env('JAX_CHECK_TRACER_LEAKS', False),
-    help=('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.'),
-)
-
 flags.DEFINE_bool(
     'jax_host_callback_inline',
     bool_env('JAX_HOST_CALLBACK_INLINE', False),
@@ -206,3 +277,72 @@ def __getattr__(self, name):
           'until the Python callback consume more outfeeds.'),
     lower_bound=int(16 * 1e6)
 )
+
+
+enable_checks = config.define_bool_state(
+    name='jax_enable_checks',
+    default=False,
+    help='Turn on invariant checking for JAX internals. Makes things slower.')
+
+check_tracer_leaks = config.define_bool_state(
+    name='jax_check_tracer_leaks',
+    default=False,
+    help=('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.'))
+checking_leaks = functools.partial(check_tracer_leaks, True)
+
+debug_nans = config.define_bool_state(
+    name='jax_debug_nans',
+    default=False,
+    help=('Add nan checks to every operation. When a nan is detected on the '
+          'output of a jit-compiled computation, call into the un-compiled '
+          'version in an attempt to more precisely identify the operation '
+          'which produced the nan.'))
+
+debug_infs = config.define_bool_state(
+    name='jax_debug_infs',
+    default=False,
+    help=('Add inf checks to every operation. When an inf is detected on the '
+          'output of a jit-compiled computation, call into the un-compiled '
+          'version in an attempt to more precisely identify the operation '
+          'which produced the inf.'))
+
+log_compiles = config.define_bool_state(
+    name='jax_log_compiles',
+    default=False,
+    help=('Log a message each time every time `jit` or `pmap` compiles an XLA '
+          'computation. Logging is performed with `absl.logging`. When this '
+          'option is set, the log level is WARNING; otherwise the level is '
+          'DEBUG.'))
+
+# Because jax_enable_x64 is managed by C++ code, we don't reuse the
+# config.define_bool_state mechanism, though conceptually it is the same.
+config.DEFINE_bool('jax_enable_x64', bool_env('JAX_ENABLE_X64', False),
+                   help='Enable 64-bit types to be used')
+lib.jax_jit.global_state().enable_x64 = bool_env('JAX_ENABLE_X64', False)
+
+@contextlib.contextmanager
+def enable_x64(new_val: bool = True):
+  """Experimental context manager to temporarily enable X64 mode.
+
+  Usage::
+
+    >>> import jax.numpy as jnp
+    >>> with enable_x64(True):
+    ...   print(jnp.arange(10.0).dtype)
+    ...
+    float64
+  """
+  prev_val = config.jax_enable_x64
+  lib.jax_jit.thread_local_state().enable_x64 = bool(new_val)
+  try:
+    yield
+  finally:
+    lib.jax_jit.thread_local_state().enable_x64 = prev_val
+Config.jax_enable_x64 = property(lambda self: lib.jax_jit.get_enable_x64())
+# config._contextmanager_flags.add('jax_enable_x64')  # TODO(mattjj): remove footgun
+
+# The `x64_enabled` property doesn't fit the naming scheme, but we use it for
+# backward compatibility.
+Config.x64_enabled = Config.jax_enable_x64