Copybara import of the project:

--
ff85248 by IvyZX <[email protected]>:

Correctly expose flax.config.temp_flip_flag

PiperOrigin-RevId: 811006892

Author: IvyZX

docs/api_reference/flax.traverse_util.rst

@@ -5,4 +5,4 @@ flax.config package
 .. automodule:: flax.configurations
     :members:
     :undoc-members:
-    :exclude-members: FlagHolder, bool_flag, temp_flip_flag, static_bool_env
+    :exclude-members: FlagHolder, bool_flag, static_bool_env

docs/api_reference/flax.config.rst renamed to docs_nnx/api_reference/flax.config.rst

@@ -0,0 +1,15 @@
+flax.traverse_util package
+============================
+
+.. currentmodule:: flax.traverse_util
+
+.. automodule:: flax.traverse_util
+
+Dict utils
+------------
+
+.. autofunction:: flatten_dict
+
+.. autofunction:: unflatten_dict
+
+.. autofunction:: path_aware_map

docs_nnx/api_reference/flax.traverse_util.rst

@@ -72,6 +72,21 @@ def __repr__(self):
     values_repr = ', '.join(f'\n  {k}={v!r}' for k, v in self._values.items())
     return f'Config({values_repr}\n)'
 
+  @contextmanager
+  def temp_flip_flag(self, var_name: str, var_value: bool):
+    """Context manager to temporarily flip feature flags for test functions.
+
+    Args:
+      var_name: the config variable name (without the 'flax_' prefix)
+      var_value: the boolean value to set var_name to temporarily
+    """
+    old_value = getattr(self, f'flax_{var_name}')
+    try:
+      self.update(f'flax_{var_name}', var_value)
+      yield
+    finally:
+      self.update(f'flax_{var_name}', old_value)
+
 
 config = Config()
 
@@ -207,22 +222,6 @@ def static_int_env(varname: str, default: int | None) -> int | None:
     ) from None
 
 
-@contextmanager
-def temp_flip_flag(var_name: str, var_value: bool):
-  """Context manager to temporarily flip feature flags for test functions.
-
-  Args:
-    var_name: the config variable name (without the 'flax_' prefix)
-    var_value: the boolean value to set var_name to temporarily
-  """
-  old_value = getattr(config, f'flax_{var_name}')
-  try:
-    config.update(f'flax_{var_name}', var_value)
-    yield
-  finally:
-    config.update(f'flax_{var_name}', old_value)
-
-
 # Flax Global Configuration Variables:
 
 flax_filter_frames = bool_flag(

flax/configurations.py

@@ -23,30 +23,36 @@ class MyTestCase(absltest.TestCase):
   def setUp(self):
     super().setUp()
     self.enter_context(mock.patch.object(config, '_values', {}))
-    self._flag = bool_flag('test', default=False, help='Just a test flag.')
+    self._flag = bool_flag('flax_test', default=False, help='Just a test flag.')
 
   def test_duplicate_flag(self):
     with self.assertRaisesRegex(RuntimeError, 'already defined'):
       bool_flag(self._flag.name, default=False, help='Another test flag.')
 
   def test_default(self):
     self.assertFalse(self._flag.value)
-    self.assertFalse(config.test)
+    self.assertFalse(config.flax_test)
 
   def test_typed_update(self):
     config.update(self._flag, True)
     self.assertTrue(self._flag.value)
-    self.assertTrue(config.test)
+    self.assertTrue(config.flax_test)
 
   def test_untyped_update(self):
     config.update(self._flag.name, True)
     self.assertTrue(self._flag.value)
-    self.assertTrue(config.test)
+    self.assertTrue(config.flax_test)
 
   def test_update_unknown_flag(self):
     with self.assertRaisesRegex(LookupError, 'Unrecognized config option'):
       config.update('unknown', True)
 
+  def test_temp_flip(self):
+    self.assertFalse(self._flag.value)
+    with config.temp_flip_flag('test', True):
+      self.assertTrue(self._flag.value)
+    self.assertFalse(self._flag.value)
+
 
 if __name__ == '__main__':
   absltest.main()

tests/configurations_test.py

@@ -19,7 +19,7 @@
 from jax import random
 
 from flax import errors
-from flax.configurations import temp_flip_flag
+from flax import config
 from flax.core import Scope, apply, freeze, init, lazy_init, nn, scope
 from flax.core.scope import LazyRng
 
@@ -283,7 +283,7 @@ def f(scope, x):
     with self.assertRaises(errors.LazyInitError):
       init_fn(random.key(0), jax.ShapeDtypeStruct((8, 4), jnp.float32))
 
-  @temp_flip_flag('fix_rng_separator', True)
+  @config.temp_flip_flag('fix_rng_separator', True)
   def test_fold_in_static_seperator(self):
     x = LazyRng(random.key(0), ('ab', 'c'))
     y = LazyRng(random.key(0), ('a', 'bc'))