add default values to config context managers

jax/_src/config.py

@@ -257,7 +257,8 @@ def get_state(self):
     setattr(Config, name, property(get_state))
 
     return _StateContextManager(name, help, update_thread_local_hook,
-                                extra_description=extra_description)
+                                extra_description=extra_description,
+                                default_value=True)
 
   def define_enum_state(
       self, name: str, enum_values: List[str], default: Optional[str],
@@ -352,18 +353,32 @@ def _trace_context(self):
     return (self.x64_enabled, self.jax_numpy_rank_promotion,
             self.jax_default_matmul_precision, self.jax_dynamic_shapes)
 
+class NoDefault: pass
+no_default = NoDefault()
+
 class _StateContextManager:
   def __init__(self, name, help, update_thread_local_hook,
                validate_new_val_hook: Optional[Callable[[Any], None]] = None,
-               extra_description: str = ""):
+               extra_description: str = "", default_value: Any = no_default):
     self._name = name
     self.__name__ = name[4:] if name.startswith('jax_') else name
-    self.__doc__ = f"Context manager for `{name}` config option{extra_description}.\n\n{help}"
+    self.__doc__ = (f"Context manager for `{name}` config option"
+                    f"{extra_description}.\n\n{help}")
     self._update_thread_local_hook = update_thread_local_hook
     self._validate_new_val_hook = validate_new_val_hook
+    self._default_value = default_value
 
   @contextlib.contextmanager
-  def __call__(self, new_val):
+  def __call__(self, new_val=no_default):
+    if new_val is no_default:
+      if self._default_value is not no_default:
+        new_val = self._default_value  # default_value provided to constructor
+      else:
+        # no default_value provided to constructor and no value provided as an
+        # argument, so we raise an error
+        raise TypeError(f"Context manager for {self.__name__} config option "
+                        "requires an argument representing the new value for "
+                        "the config option.")
     if self._validate_new_val_hook:
       self._validate_new_val_hook(new_val)
     prev_val = getattr(_thread_local_state, self._name, unset)