Merge pull request #2898 from chiamp:frozen_dict_tests

PiperOrigin-RevId: 512190368

Author: Flax Authors

flax/configurations.py

@@ -24,6 +24,8 @@
 import os
 from jax import config as jax_config
 
+from contextlib import ContextDecorator
+
 # Keep a wrapper at the flax namespace, in case we make our implementation
 # in the future.
 config = jax_config
@@ -66,6 +68,23 @@ def static_bool_env(varname: str, default: bool) -> bool:
         'invalid truth value {!r} for environment {!r}'.format(val, varname))
 
 
+class use_regular_dict(ContextDecorator):
+  """Context decorator for test functions to temporarily use regular dicts
+  instead of FrozenDicts.
+
+  This is a temporary feature flag to help migrate to FrozenDicts. Returning
+  FrozenDicts will be deprecated and removed in the future.
+  """
+  def __enter__(self):
+    self._old_value = config.flax_return_frozendict # save current env value
+    config.update('flax_return_frozendict', False) # return regular dicts
+    return self
+
+  def __exit__(self, *exc):
+    config.update('flax_return_frozendict', self._old_value) # switch back to old env value
+    return False
+
+
 # Flax Global Configuration Variables:
 
 # Whether to use the lazy rng implementation.
@@ -96,6 +115,7 @@ def static_bool_env(varname: str, default: bool) -> bool:
     default=False,
     help=("When adopting outside modules, don't clobber existing names."))
 
+#TODO(marcuschiam): remove this feature flag once regular dict migration is complete
 flax_return_frozendict = define_bool_state(
     name='return_frozendict',
     default=True,

flax/linen/summary.py

@@ -310,15 +310,15 @@ def _get_path_variables(path: Tuple[str, ...], variables: FrozenVariableDict) ->
   path_variables = {}
 
   for collection in variables:
-    collection_variables = variables[collection]
+    collection_variables = jax.tree_util.tree_map(lambda x: x, variables[collection]) # make a deep copy
     for name in path:
       if name not in collection_variables:
         collection_variables = None
         break
       collection_variables = collection_variables[name]
 
     if collection_variables is not None:
-      path_variables[collection] = collection_variables.unfreeze()
+      path_variables[collection] = collection_variables
 
   return path_variables
 

tests/core/core_lift_test.py

@@ -14,7 +14,8 @@
 
 import operator
 from flax import errors
-from flax.core import Scope, init, apply, lift, nn, FrozenDict, unfreeze
+from flax.core import Scope, init, apply, lift, nn, FrozenDict, unfreeze, copy
+from flax.configurations import use_regular_dict
 
 import jax
 from jax import random
@@ -25,6 +26,7 @@
 
 from absl.testing import absltest
 
+
 class LiftTest(absltest.TestCase):
 
   def test_aliasing(self):
@@ -155,11 +157,12 @@ def false_fn(scope, x):
 
     x = jnp.ones((1, 3))
     y1, vars = init(f)(random.PRNGKey(0), x, True)
-    self.assertEqual(vars['state'].unfreeze(), {'true_count': 1, 'false_count': 0})
+    self.assertEqual(vars['state'], {'true_count': 1, 'false_count': 0})
     y2, vars = apply(f, mutable="state")(vars, x, False)
-    self.assertEqual(vars['state'].unfreeze(), {'true_count': 1, 'false_count': 1})
+    self.assertEqual(vars['state'], {'true_count': 1, 'false_count': 1})
     np.testing.assert_allclose(y1, -y2)
 
+  @use_regular_dict()
   def test_switch(self):
     def f(scope, x, index):
       scope.variable('state', 'a_count', lambda: 0)
@@ -182,14 +185,14 @@ def c_fn(scope, x):
 
     x = jnp.ones((1, 3))
     y1, vars = init(f)(random.PRNGKey(0), x, 0)
-    self.assertEqual(vars['state'].unfreeze(), {'a_count': 1, 'b_count': 0, 'c_count': 0})
+    self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 0, 'c_count': 0})
     y2, updates = apply(f, mutable="state")(vars, x, 1)
-    vars = vars.copy(updates)
-    self.assertEqual(vars['state'].unfreeze(), {'a_count': 1, 'b_count': 1, 'c_count': 0})
+    vars = copy(vars, updates)
+    self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 1, 'c_count': 0})
     np.testing.assert_allclose(y1, -y2)
     y3, updates = apply(f, mutable="state")(vars, x, 2)
-    vars = vars.copy(updates)
-    self.assertEqual(vars['state'].unfreeze(), {'a_count': 1, 'b_count': 1, 'c_count': 1})
+    vars = copy(vars, updates)
+    self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 1, 'c_count': 1})
     np.testing.assert_allclose(y1, y3)
 
   def test_subscope_var_aliasing(self):

tests/core/core_meta_test.py

@@ -110,7 +110,7 @@ def body(scope, x):
       return c
 
     _, variables = init(f)(random.PRNGKey(0), jnp.zeros((8, 3)))
-    boxed_shapes = jax.tree_map(jnp.shape, variables['params'].unfreeze())
+    boxed_shapes = jax.tree_map(jnp.shape, variables['params'])
     self.assertEqual(boxed_shapes, {
         'kernel': meta.Partitioned((8, 3, 3), ('layers', 'in', 'out')),
         'bias': (8, 3),

tests/linen/linen_attention_test.py

@@ -19,6 +19,8 @@
 
 from flax import linen as nn
 from flax import jax_utils
+from flax.core import pop
+from flax.configurations import use_regular_dict
 
 import jax
 from jax import lax
@@ -31,7 +33,6 @@
 # Parse absl flags test_srcdir and test_tmpdir.
 jax.config.parse_flags_with_absl()
 
-
 class AttentionTest(parameterized.TestCase):
 
   def test_multihead_self_attention(self):
@@ -102,6 +103,7 @@ def test_causal_mask_1d(self):
     np.testing.assert_allclose(mask_1d, mask_1d_simple,)
 
   @parameterized.parameters([((5,), (1,)), ((6, 5), (2,))])
+  @use_regular_dict()
   def test_decoding(self, spatial_shape, attn_dims):
     bs = 2
     num_heads = 3
@@ -119,7 +121,7 @@ def test_decoding(self, spatial_shape, attn_dims):
     decode_module = module.clone(decode=True)
 
     initial_vars = decode_module.init(key2, inputs)
-    state, params = initial_vars.pop('params')
+    state, params = pop(initial_vars, 'params')
     causal_mask = nn.attention.make_causal_mask(jnp.ones((bs,) + spatial_shape))
     y_ref = jax.jit(lambda x, y: module.apply(initial_vars, x, y))(
         inputs, causal_mask)

tests/linen/linen_meta_test.py

@@ -144,7 +144,7 @@ def body(_, c):
     x = jnp.ones((8, 128))
     spec = nn.get_partition_spec(
         jax.eval_shape(model.init, random.PRNGKey(0), x))
-    self.assertEqual(spec.unfreeze(), {
+    self.assertEqual(spec, {
         'params': {
             'MLP_0': {
                 'Dense_0': {

tests/linen/linen_module_test.py

@@ -31,6 +31,7 @@
 from flax import struct
 from flax.core import Scope, freeze, FrozenDict, tracers
 from flax.linen import compact
+from flax.configurations import use_regular_dict
 import jax
 from jax import random
 from jax.nn import initializers
@@ -41,7 +42,6 @@
 # Parse absl flags test_srcdir and test_tmpdir.
 jax.config.parse_flags_with_absl()
 
-
 def tree_equals(x, y):
   return jax.tree_util.tree_all(jax.tree_util.tree_map(operator.eq, x, y))
 
@@ -1140,6 +1140,7 @@ def test(self):
     A().test()
     self.assertFalse(setup_called)
 
+  @use_regular_dict()
   def test_module_pass_as_attr(self):
 
     class A(nn.Module):
@@ -1158,7 +1159,7 @@ def __call__(self, x):
 
     variables = A().init(random.PRNGKey(0), jnp.ones((1,)))
     var_shapes = jax.tree_util.tree_map(jnp.shape, variables)
-    ref_var_shapes = freeze({
+    ref_var_shapes = {
         'params': {
             'b': {
                 'foo': {
@@ -1167,9 +1168,10 @@ def __call__(self, x):
                 }
             },
         },
-    })
+    }
     self.assertTrue(tree_equals(var_shapes, ref_var_shapes))
 
+  @use_regular_dict()
   def test_module_pass_in_closure(self):
     a = nn.Dense(2)
 
@@ -1183,17 +1185,18 @@ def __call__(self, x):
 
     variables = B().init(random.PRNGKey(0), jnp.ones((1,)))
     var_shapes = jax.tree_util.tree_map(jnp.shape, variables)
-    ref_var_shapes = freeze({
+    ref_var_shapes = {
         'params': {
             'foo': {
                 'bias': (2,),
                 'kernel': (1, 2),
             }
         },
-    })
+    }
     self.assertTrue(tree_equals(var_shapes, ref_var_shapes))
     self.assertIsNone(a.name)
 
+  @use_regular_dict()
   def test_toplevel_submodule_adoption(self):
 
     class Encoder(nn.Module):
@@ -1233,7 +1236,7 @@ def __call__(self, x):
     self.assertEqual(y.shape, (4, 5))
 
     var_shapes = jax.tree_util.tree_map(jnp.shape, variables)
-    ref_var_shapes = freeze({
+    ref_var_shapes = {
         'params': {
             'dense_out': {
                 'bias': (5,),
@@ -1246,9 +1249,10 @@ def __call__(self, x):
                 },
             },
         },
-    })
+    }
     self.assertTrue(tree_equals(var_shapes, ref_var_shapes))
 
+  @use_regular_dict()
   def test_toplevel_submodule_adoption_pytree(self):
 
     class A(nn.Module):
@@ -1276,7 +1280,7 @@ def __call__(self, c, x):
 
     params = B(a_pytree).init(key, x, x)
     unused_y, counters = b.apply(params, x, x, mutable='counter')
-    ref_counters = freeze({
+    ref_counters = {
         'counter': {
             'A_bar': {
                 'i': jnp.array(2.0),
@@ -1285,13 +1289,14 @@ def __call__(self, c, x):
                 'i': jnp.array(2.0),
             },
         },
-    })
+    }
     self.assertTrue(
         jax.tree_util.tree_all(
             jax.tree_util.tree_map(
                 lambda x, y: np.testing.assert_allclose(x, y, atol=1e-7),
                 counters, ref_counters)))
 
+  @use_regular_dict()
   def test_toplevel_submodule_adoption_sharing(self):
     dense = functools.partial(nn.Dense, use_bias=False)
 
@@ -1323,7 +1328,7 @@ def __call__(self, x):
     c = C(a, b)
     p = c.init(key, x)
     var_shapes = jax.tree_util.tree_map(jnp.shape, p)
-    ref_var_shapes = freeze({
+    ref_var_shapes = {
         'params': {
             'Dense_0': {
                 'kernel': (2, 2),
@@ -1339,9 +1344,10 @@ def __call__(self, x):
                 },
             },
         },
-    })
+    }
     self.assertTrue(tree_equals(var_shapes, ref_var_shapes))
 
+  @use_regular_dict()
   def test_toplevel_named_submodule_adoption(self):
     dense = functools.partial(nn.Dense, use_bias=False)
 
@@ -1369,7 +1375,7 @@ def __call__(self, x):
     init_vars = b.init(k, x)
     var_shapes = jax.tree_util.tree_map(jnp.shape, init_vars)
     if config.flax_preserve_adopted_names:
-      ref_var_shapes = freeze({
+      ref_var_shapes = {
           'params': {
               'foo': {
                   'dense': {
@@ -1380,9 +1386,9 @@ def __call__(self, x):
                   'kernel': (4, 6),
               },
           },
-      })
+      }
     else:
-      ref_var_shapes = freeze({
+      ref_var_shapes = {
           'params': {
               'a': {
                   'dense': {
@@ -1393,9 +1399,10 @@ def __call__(self, x):
                   'kernel': (4, 6),
               },
           },
-      })
+      }
     self.assertTrue(tree_equals(var_shapes, ref_var_shapes))
 
+  @use_regular_dict()
   def test_toplevel_submodule_pytree_adoption_sharing(self):
 
     class A(nn.Module):
@@ -1423,13 +1430,13 @@ def __call__(self, x):
 
     params = b.init(key, x)
     _, counters = b.apply(params, x, mutable='counter')
-    ref_counters = freeze({
+    ref_counters = {
         'counter': {
             'A_bar': {
                 'i': jnp.array(6.0),
             },
         },
-    })
+    }
     self.assertTrue(tree_equals(counters, ref_counters))
 
   def test_inner_class_def(self):
@@ -1650,7 +1657,6 @@ def __call__(self, x):
 
     x = jnp.ones((3,))
     variables = Foo().init(random.PRNGKey(0), x)
-    variables = variables.unfreeze()
     y = Foo().apply(variables, x)
     self.assertEqual(y.shape, (2,))