make config.FLAGS.jax_enable_foo an error

jax/config.py

@@ -51,6 +51,7 @@ def __init__(self):
  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)
@@ -71,6 +72,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:
@@ -193,23 +201,17 @@ def define_bool_state(self, name: str, default: bool, help: str):
  with enable_foo(True):
  ...
 
- Accessing ``config.FLAGS.jax_enable_foo`` is different from accessing the
- thread-local state value via ``config.jax_enable_foo``: the former reads the
- flag value determined set by the environment variable or command-line flag
- and does not read the thread-local state, whereas the latter reads the
- thread-local state value managed by the contextmanager. Think of the
- contextmanager state as a layer on top of the flag value: if no
- contextmanager is in use then ``config.jax_enable_foo`` reflects the flag
- value ``config.FLAGS.jax_enable_foo``, whereas if a contextmanager is in use
- then only ``config.jax_enable_foo`` is updated. So in general using
- ``config.jax_enable_foo`` is best.
+ 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)
+ return val if val is not unset else self._read(name)
  setattr(Config, name, property(get_state))
 
  @contextlib.contextmanager

tests/debug_nans_test.py

@@ -30,7 +30,7 @@
 class DebugNaNsTest(jtu.JaxTestCase):
 
  def setUp(self):
- self.cfg = config.read("jax_debug_nans")
+ self.cfg = config._read("jax_debug_nans")
  config.update("jax_debug_nans", True)
 
  def tearDown(self):
@@ -144,7 +144,7 @@ def testPjit(self):
 class DebugInfsTest(jtu.JaxTestCase):
 
  def setUp(self):
- self.cfg = config.read("jax_debug_infs")
+ self.cfg = config._read("jax_debug_infs")
  config.update("jax_debug_infs", True)
 
  def tearDown(self):