Repository: google/flax

Issue:
Make `self.make_rng()` callable by using a default RNG stream and have `.init()` and `.apply()` use the default RNG stream if no explicit RNG mapping is passed.
There has been some internal and external (#3506) discussion on making `self.make_rng()` callable by using a default RNG stream and having `.init()` and `.apply()` use the default RNG stream if no explicit RNG mapping is passed. 

In the current design:
* a string name must be passed into `self.make_rng(name)` to explicitly specify the RNG stream to use in the Module
* an explicit RNG mapping (e.g. `{'params': random.key(0), 'other_rng_stream': random.key(1)}`) must be passed to `.init()` and `.apply()` for each RNG stream that's used in the Module

A proposed change:
* Make `self.make_rng(name='default')` callable without explicitly passing in a string name, in which it would default to the `'default'` RNG stream
  * This would allow the user to pass in a single RNG mapping (e.g. `{'default': random.key(0)}`) to `.init()` and `.apply()` that `self.make_rng('other_rng_stream')` will use as the RNG seed if, for example, `{'other_rng_stream': random.key(1)}` wasn't in the RNG mapping that was passed into `.init()` or `.apply()`

Discussion:

Patch:
diff --git a/docs/api_reference/flax.linen/module.rst b/docs/api_reference/flax.linen/module.rst
index 200bd124..302546dc 100644
--- a/docs/api_reference/flax.linen/module.rst
+++ b/docs/api_reference/flax.linen/module.rst
@@ -5,7 +5,7 @@ Module
 .. currentmodule:: flax.linen
 
 .. autoclass:: Module
-   :members: setup, variable, param, bind, unbind, apply, init, init_with_output, copy, make_rng, sow, variables, Variable, __setattr__, tabulate, module_paths, is_initializing, perturb, put_variable, has_variable, has_rng, lazy_init, get_variable, path, is_mutable_collection
+   :members: setup, variable, param, bind, unbind, apply, init, init_with_output, copy, make_rng, sow, variables, Variable, __setattr__, tabulate, is_initializing, perturb, put_variable, has_variable, has_rng, lazy_init, get_variable, path, is_mutable_collection
 
 .. autofunction:: apply
 .. autofunction:: init
diff --git a/flax/core/scope.py b/flax/core/scope.py
index 4ea8caa4..1b0e081b 100644
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -750,10 +750,13 @@ class Scope:
     """Returns true if a PRNGSequence with name `name` exists."""
     return name in self.rngs
 
-  def make_rng(self, name: str) -> PRNGKey:
+  def make_rng(self, name: str = 'default') -> PRNGKey:
     """Generates A PRNGKey from a PRNGSequence with name `name`."""
     if not self.has_rng(name):
-      raise errors.InvalidRngError(f'{self.name} needs PRNG for "{name}"')
+      if self.has_rng('default'):
+        name = 'default'
+      else:
+        raise errors.InvalidRngError(f'{self.name} needs PRNG for "{name}"')
     self._check_valid()
     self._validate_trace_level()
     self.rng_counters[name] += 1
diff --git a/flax/linen/module.py b/flax/linen/module.py
index 0abafc03..db2d8472 100644
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -508,10 +508,7 @@ def nowrap(fun: _CallableT) -> _CallableT:
    - if you're subclassing a method like Module.param and don't want this
      overriden core function decorated with the state management wrapper.
    - If you want a method to be callable from an unbound Module (e.g.: a
-     function of construction of arguments that doesn't depend on params/RNGs).
-     If you want to learn more about how Flax Modules manage their state read the
-     [The Flax Module lifecycle](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html)
-     guide.
+     function of construction of arguments that doesn't depend on params/RNGs)
 
   For instance::
 
@@ -1758,7 +1755,7 @@ class Module(ModuleBase):
       raise ValueError("Can't query for RNGs on unbound modules")
     return self.scope.has_rng(name)
 
-  def make_rng(self, name: str) -> KeyArray:
+  def make_rng(self, name: str = 'default') -> KeyArray:
     """Returns a new RNG key from a given RNG sequence for this Module.
 
     The new RNG key is split from the previous one. Thus, every call to
diff --git a/flax/linen/normalization.py b/flax/linen/normalization.py
index 62a0ae02..1535a7e8 100644
--- a/flax/linen/normalization.py
+++ b/flax/linen/normalization.py
@@ -380,31 +380,18 @@ class LayerNorm(Module):
   i.e. applies a transformation that maintains the mean activation within
   each example close to 0 and the activation standard deviation close to 1.
 
-  NOTE: This normalization operation is identical to InstanceNorm and GroupNorm;
-  the difference is simply which axes are reduced and the shape of the feature axes
-  (i.e. the shape of the learnable scale and bias parameters).
-
   Example usage::
 
     >>> import flax.linen as nn
-    >>> import jax
-    >>> import numpy as np
+    >>> import jax, jax.numpy as jnp
 
-    >>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6))
+    >>> x = jax.random.normal(jax.random.key(0), (5, 6))
     >>> layer = nn.LayerNorm()
     >>> variables = layer.init(jax.random.key(1), x)
     >>> variables
     {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}}
     >>> y = layer.apply(variables, x)
 
-    >>> y = nn.LayerNorm(reduction_axes=(1, 2, 3)).apply(variables, x)
-    >>> y2 = nn.GroupNorm(num_groups=1).apply(variables, x)
-    >>> np.testing.assert_allclose(y, y2)
-
-    >>> y = nn.LayerNorm(reduction_axes=(1, 2), feature_axes=-1).apply(variables, x)
-    >>> y2 = nn.InstanceNorm(channel_axes=-1).apply(variables, x)
-    >>> np.testing.assert_allclose(y, y2)
-
   Attributes:
     epsilon: A small float added to variance to avoid dividing by zero.
     dtype: the dtype of the result (default: infer from input and params).
@@ -589,30 +576,18 @@ class GroupNorm(Module):
   The user should either specify the total number of channel groups or the
   number of channels per group.
 
-  NOTE: LayerNorm is a special case of GroupNorm where `num_groups=1`, and
-  InstanceNorm is a special case of GroupNorm where `group_size=1`.
-
   Example usage::
 
     >>> import flax.linen as nn
-    >>> import jax
-    >>> import numpy as np
+    >>> import jax, jax.numpy as jnp
 
-    >>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6))
+    >>> x = jax.random.normal(jax.random.key(0), (5, 6))
     >>> layer = nn.GroupNorm(num_groups=3)
     >>> variables = layer.init(jax.random.key(1), x)
     >>> variables
     {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}}
     >>> y = layer.apply(variables, x)
 
-    >>> y = nn.GroupNorm(num_groups=1).apply(variables, x)
-    >>> y2 = nn.LayerNorm(reduction_axes=(1, 2, 3)).apply(variables, x)
-    >>> np.testing.assert_allclose(y, y2)
-
-    >>> y = nn.GroupNorm(num_groups=None, group_size=1).apply(variables, x)
-    >>> y2 = nn.InstanceNorm(channel_axes=-1).apply(variables, x)
-    >>> np.testing.assert_allclose(y, y2)
-
   Attributes:
     num_groups: the total number of channel groups. The default value of 32 is
       proposed by the original group normalization paper.
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
index c0d794ba..4aadf57b 100644
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -54,7 +54,6 @@ from typing import (
 )
 
 from flax import errors, struct, traceback_util
-from flax import serialization
 from flax.core import Scope, lift, meta
 from flax.core.frozen_dict import FrozenDict
 from flax.ids import FlaxId
@@ -443,6 +442,25 @@ def decorator_lift_transform(
   return wrapped_fn
 
 
+@dataclasses.dataclass(frozen=True)
+class _ModuleHash:
+  module: Module
+  hashable_module: Module
+
+  @classmethod
+  def from_module(cls, module: Module):
+    return cls(module, module.clone(_deep_clone=True, _reset_names=True))
+
+  def __hash__(self):
+    return hash(self.hashable_module)
+
+  def __eq__(self, other):
+    return (
+        isinstance(other, _ModuleHash)
+        and self.hashable_module == other.hashable_module
+    )
+
+
 @dataclasses.dataclass(frozen=True)
 class _HashableProxy:
   """A hashable proxy object that is use to define a hash for Modules.
@@ -454,84 +472,76 @@ class _HashableProxy:
   module: Module
 
   def __hash__(self):
-    key = _module_pytree_hash(self.module)
+    key = _hashable_pytree(self.module)
     return hash(key)
 
   def __eq__(self, other):
-    return isinstance(other, _HashableProxy) and _module_pytree_hash(
+    return isinstance(other, _HashableProxy) and _hashable_pytree(
         self.module
-    ) == _module_pytree_hash(other.module)
+    ) == _hashable_pytree(other.module)
 
 
-def _recursive_freeze(name: str, x: Any, seen_modules: dict[FlaxId, int]):
-  if isinstance(x, dict):
-    return tuple(
-        (k, _recursive_freeze(k, v, seen_modules)) for k, v in x.items()
-    )
-  elif isinstance(x, Module):
-    return _module_pytree_hash_recursive(x, seen_modules)
-  else:
-    _check_field_is_hashable(name, x)
-    return x
-
-
-def _module_pytree_hash(module: Module) -> tuple[type[Any], Any]:
-  return _module_pytree_hash_recursive(module, {})
+def _hashable_pytree(x: Any) -> tuple[type[Any], Any]:
+  return _hashable_pytree_recursive(x, {})
 
 
-def _module_pytree_hash_recursive(
-    module: Module, seen_modules: dict[FlaxId, int]
-) -> tuple[type[Any], Any]:
+def _hashable_pytree_recursive(
+    x: Any, modules_seen: dict[FlaxId, int]
+) -> tuple[type, Any]:
   """Creates a hashable representation for a Module by travering its structure recursively."""
   static: Any
-  if module._id in seen_modules:
-    # if we have already seen the module we just use the index
-    # as its static component
-    static = seen_modules[module._id]
+  if isinstance(x, Module):
+    if x._id in modules_seen:
+      # if we have already seen the module we just use the index
+      # as its static component
+      static = modules_seen[x._id]
+    else:
+      # if its a new module we add it to the cache and give it
+      # a new index
+      modules_seen[x._id] = len(modules_seen)
+      # TODO(cgarciae): define a way for the user of nn.jit to define
+      # what fields it wants to ignore per Module instance.
+      static = tuple(
+          (
+              field.name,
+              _hashable_pytree_recursive(getattr(x, field.name), modules_seen),
+          )
+          for field in dataclasses.fields(x)
+          if hasattr(x, field.name) and field.name not in ('parent', 'name')
+      )
+  elif isinstance(x, FrozenDict):
+    static = tuple(
+        (name, _hashable_pytree_recursive(value, modules_seen))
+        for name, value in sorted(x.items())
+    )
+  elif isinstance(x, tuple):
+    static = tuple(
+        _hashable_pytree_recursive(value, modules_seen) for value in x
+    )
   else:
-    # if its a new module we add it to the cache and give it
-    # a new index
-    seen_modules[module._id] = len(seen_modules)
-    # TODO(cgarciae): define a way for the user of nn.jit to define
-    # what fields it wants to ignore per Module instance.
-    values = []
-    for field in dataclasses.fields(module):
-      if field.init and field.name not in ('parent', 'name'):
-        value = getattr(module, field.name)
-        if isinstance(value, Module):
-          value = _module_pytree_hash_recursive(value, seen_modules)
-        elif serialization.is_serializable(value):
-          value = serialization.to_state_dict(value)
-          value = _recursive_freeze(field.name, value, seen_modules)
+    # test for hashability to catch errors early
+    try:
+      hash(x)
+    except Exception as e:
+      if dataclasses.is_dataclass(x):
+        if not hasattr(x, '__hash__'):
+          raise ValueError(
+              f"type '{type(x)}' is a dataclass but its not hashable, using"
+              ' `dataclass(frozen=True, eq=True)` to make it immutable and'
+              ' hashable, or `dataclass(unsafe_hash=True)` for mutable types.'
+          ) from e
         else:
-          _check_field_is_hashable(field.name, value)
-        values.append((field.name, value))
-    static = tuple(values)
-
-  return type(module), static
-
-
-def _check_field_is_hashable(name: str, x: Any):
-  """Checks if a field is hashable."""
-  try:
-    hash(x)
-  except Exception as e:
-    if dataclasses.is_dataclass(x):
-      if x.__hash__ is None:
-        raise ValueError(
-            f"field '{name}' of type '{type(x)}' is a dataclass but its not"
-            ' hashable, using `dataclass(frozen=True, eq=True)` to make it'
-            ' immutable and hashable, or `dataclass(unsafe_hash=True)` for'
-            ' mutable types.'
-        ) from e
+          raise ValueError(
+              f"type '{type(x)}' is a hashable dataclass but hashing failed,"
+              ' this probably means that at least one of its fields is not'
+              ' hashable.'
+          ) from e
       else:
-        raise ValueError(
-            f"field '{name}' of type '{type(x)}' is a hashable dataclass but"
-            ' hashing failed, this probably means that at least one of its'
-            ' fields is not hashable.'
-        ) from e
-    else:
-      raise ValueError(f"type '{type(x)}' is not hashable.") from e
+        raise ValueError(f"type '{type(x)}' is not hashable.") from e
+
+    static = x
+
+  return type(x), static
 
 
 def decorator_lift_transform_jit(class_fn, **trafo_kwargs):
@@ -604,6 +614,7 @@ def decorator_lift_transform_jit(class_fn, **trafo_kwargs):
 
     # get a hash for the Module by using its repr as a proxy
     hash_key = _HashableProxy(self)
+    # hash_key = _ModuleHash.from_module(self)
 
     return trafo_fn(module_scopes, hash_key, *args, **kwargs)
 
@@ -677,6 +688,7 @@ def module_class_lift_transform_jit(module_class, methods=None, **trafo_kwargs):
 
       # get a hash for the Module by using its repr as a proxy
       hash_key = _HashableProxy(self)
+      # hash_key = _ModuleHash.from_module(self)
 
       ret = trafo_fn(module_scopes, hash_key, *args, **kwargs)
       return ret
diff --git a/flax/serialization.py b/flax/serialization.py
index 485894f4..6fef6b88 100644
--- a/flax/serialization.py
+++ b/flax/serialization.py
@@ -110,12 +110,6 @@ def to_state_dict(target) -> Dict[str, Any]:
   return state_dict
 
 
-def is_serializable(target):
-  if not isinstance(target, type):
-    target = type(target)
-  return target in _STATE_DICT_REGISTRY
-
-
 def register_serialization_state(
   ty, ty_to_state_dict, ty_from_state_dict, override=False
 ):


Test patch:
diff --git a/tests/linen/linen_module_test.py b/tests/linen/linen_module_test.py
index 9150a937..2a2eb078 100644
--- a/tests/linen/linen_module_test.py
+++ b/tests/linen/linen_module_test.py
@@ -2522,6 +2522,145 @@ class ModuleTest(absltest.TestCase):
     self.assertIn('Bar_0/Dense_1', module_paths)
     self.assertEqual(type(module_paths['Bar_0/Dense_1']), nn.Dense)
 
+  def test_init_apply_default_rng(self):
+    class SubModel(nn.Module):
+      @nn.compact
+      def __call__(self, x, apply_dropout):
+        x = nn.Dense(8)(x)
+        x = nn.Dropout(0.8)(x, deterministic=not apply_dropout)
+        p = self.param(
+          'parameter', lambda key, shape: jax.random.normal(key, shape), x.shape
+        )
+        noise = jax.random.normal(self.make_rng('noise'), x.shape)
+        return x * p + noise
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x, apply_dropout):
+        x = nn.Dense(16)(x)
+        x = SubModel()(x, apply_dropout)
+        x = nn.Dropout(0.5)(x, deterministic=not apply_dropout)
+        v = self.variable(
+          'var_collection',
+          'variable',
+          lambda shape: jax.random.normal(self.make_rng('var_rng'), shape),
+          x.shape,
+        )
+        noise = jax.random.normal(self.make_rng('noise'), x.shape)
+        return x * v.value + noise
+
+    key0, key1, key2 = jax.random.split(jax.random.key(0), 3)
+    x = jax.random.normal(key0, (10, 4))
+    model = Model()
+
+    # test init equality
+    default_variables = model.init({'default': key1}, x, apply_dropout=False)
+    # adding 'default' rng shouldn't change anything
+    rngs = {'params': key1, 'var_rng': key1, 'noise': key1, 'default': key0}
+    explicit_variables = model.init(rngs, x, apply_dropout=False)
+    self.assertTrue(
+      jax.tree_util.tree_all(
+        jax.tree_map(
+          lambda v1, v2: (v1 == v2).all(), default_variables, explicit_variables
+        )
+      )
+    )
+
+    # test init inequality
+    rngs['default'] = key1  # adding 'default' rng shouldn't change anything
+    for rng_name in ('params', 'var_rng'):
+      rngs[rng_name] = key2
+      explicit_variables = model.init(rngs, x, apply_dropout=False)
+      self.assertFalse(
+        jax.tree_util.tree_all(
+          jax.tree_map(
+            lambda v1, v2: (v1 == v2).all(),
+            default_variables,
+            explicit_variables,
+          )
+        )
+      )
+      rngs[rng_name] = key1
+
+    # test apply equality
+    default_out = model.apply(
+      default_variables, x, apply_dropout=True, rngs={'default': key1}
+    )
+    # adding 'default' rng shouldn't change anything
+    rngs = {'dropout': key1, 'noise': key1, 'default': key0}
+    explicit_out = model.apply(
+      default_variables, x, apply_dropout=True, rngs=rngs
+    )
+    np.testing.assert_allclose(default_out, explicit_out)
+
+    # test apply inequality
+    rngs['default'] = key1  # adding 'default' rng shouldn't change anything
+    for rng_name in ('dropout', 'noise'):
+      rngs[rng_name] = key2
+      explicit_out = model.apply(
+        default_variables, x, apply_dropout=True, rngs=rngs
+      )
+      with self.assertRaises(AssertionError):
+        np.testing.assert_allclose(default_out, explicit_out, atol=1e-1)
+      rngs[rng_name] = key1
+
+  def test_default_make_rng(self):
+    class SubModel(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        noise = jax.random.normal(self.make_rng(), x.shape)
+        return x + noise
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        x = SubModel()(x)
+        noise = jax.random.normal(self.make_rng(), x.shape)
+        return x + noise
+
+    key0, key1 = jax.random.split(jax.random.key(0), 2)
+    x = jax.random.normal(key0, (10, 4))
+    default_out = Model().apply({}, x, rngs={'default': key1})
+
+    class SubModel(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        noise = jax.random.normal(self.make_rng('default'), x.shape)
+        return x + noise
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        x = SubModel()(x)
+        noise = jax.random.normal(self.make_rng('default'), x.shape)
+        return x + noise
+
+    explicit_out = Model().apply({}, x, rngs={'default': key1})
+    np.testing.assert_allclose(default_out, explicit_out)
+
+  def test_default_rng_error(self):
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return nn.Dense(2)(x)
+
+    model = Model()
+    with self.assertRaisesRegex(
+      errors.InvalidRngError, 'Dense_0 needs PRNG for "params"'
+    ):
+      model.init({'other_rng_stream': jax.random.key(0)}, jnp.ones((1, 3)))
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return x + jax.random.normal(self.make_rng('default'), x.shape)
+
+    model = Model()
+    with self.assertRaisesRegex(
+      errors.InvalidRngError, 'None needs PRNG for "default"'
+    ):
+      model.init(jax.random.key(0), jnp.ones((1, 3)))
+
 
 class LeakTests(absltest.TestCase):
   def test_tracer_leaks(self):
diff --git a/tests/linen/linen_test.py b/tests/linen/linen_test.py
index 9cc9b040..88735070 100644
--- a/tests/linen/linen_test.py
+++ b/tests/linen/linen_test.py
@@ -485,69 +485,6 @@ class NormalizationTest(parameterized.TestCase):
       layer = nn.InstanceNorm(channel_axes=channel_axes)
       _ = layer.init(jax.random.key(1), x)
 
-  @parameterized.parameters(
-    {
-      'layer1': nn.LayerNorm(feature_axes=(1, 2)),
-      'layer2': nn.InstanceNorm(channel_axes=(1, 2)),
-    },
-    {
-      'layer1': nn.LayerNorm(reduction_axes=(1, 2), feature_axes=-1),
-      'layer2': nn.InstanceNorm(channel_axes=-1),
-    },
-    {
-      'layer1': nn.LayerNorm(
-        reduction_axes=-2,
-        feature_axes=(1, 3),
-        bias_init=nn.initializers.uniform(),
-        scale_init=nn.initializers.uniform(),
-      ),
-      'layer2': nn.InstanceNorm(
-        channel_axes=(1, -1),
-        bias_init=nn.initializers.uniform(),
-        scale_init=nn.initializers.uniform(),
-      ),
-    },
-    {
-      'layer1': nn.LayerNorm(
-        reduction_axes=(1, 2, 3),
-        bias_init=nn.initializers.uniform(),
-        scale_init=nn.initializers.uniform(),
-      ),
-      'layer2': nn.GroupNorm(
-        num_groups=1,
-        bias_init=nn.initializers.uniform(),
-        scale_init=nn.initializers.uniform()
-      ),
-    },
-    {
-      'layer1': nn.InstanceNorm(
-        bias_init=nn.initializers.uniform(),
-        scale_init=nn.initializers.uniform(),
-      ),
-      'layer2': nn.GroupNorm(
-        num_groups=None,
-        group_size=1,
-        bias_init=nn.initializers.uniform(),
-        scale_init=nn.initializers.uniform(),
-      ),
-    },
-  )
-  def test_normalization_equivalence(self, layer1, layer2):
-    x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))
-    layer1_variables = layer1.init(jax.random.key(1), x)
-    layer2_variables = layer2.init(jax.random.key(1), x)
-    self.assertTrue(
-      jax.tree_util.tree_all(
-        jax.tree_map(
-          lambda v1, v2: (v1 == v2).all(), layer1_variables, layer2_variables
-        )
-      )
-    )
-
-    layer1_y = layer1.apply(layer1_variables, x)
-    layer2_y = layer2.apply(layer2_variables, x)
-    np.testing.assert_allclose(layer1_y, layer2_y, atol=1e-7)
-
   @parameterized.parameters(
     {
       'model_index': 0,
diff --git a/tests/linen/linen_transforms_test.py b/tests/linen/linen_transforms_test.py
index ad86881d..04b212d5 100644
--- a/tests/linen/linen_transforms_test.py
+++ b/tests/linen/linen_transforms_test.py
@@ -22,7 +22,6 @@ import unittest
 from absl.testing import absltest, parameterized
 from flax import errors
 from flax import linen as nn
-from flax import serialization
 from flax.core import copy, freeze
 from flax.linen.transforms import _HashableProxy
 import jax
@@ -1940,46 +1939,6 @@ class TransformTest(parameterized.TestCase):
     y = Parent([Foo(3), Foo(4)]).apply({}, x)
     self.assertEqual(n, 2)
 
-  def test_jit_hashes_serializable_types(self):
-    class Node:
-
-      def __init__(self, a: int):
-        self.a = a
-
-      def __hash__(self):
-        # test object is not being passed as static
-        raise Exception('immutable')
-
-    def to_dict(node: Node):
-      return {'a': node.a}
-
-    def from_dict(node: Node, d: dict[str, Any]):
-      node.a = d['a']
-      return node
-
-    serialization.register_serialization_state(Node, to_dict, from_dict)
-
-    try:
-      n = 0
-
-      class Foo(nn.Module):
-        node: Node
-
-        @nn.jit
-        @nn.compact
-        def __call__(self, x):
-          nonlocal n
-          n += 1
-          return self.node.a + nn.Dense(2)(x)
-
-      m = Foo(Node(1))
-      m.init_with_output(random.key(0), jnp.ones((2, 2)))
-      self.assertEqual(n, 1)
-      m.init_with_output(random.key(0), jnp.ones((2, 2)))
-      self.assertEqual(n, 1)
-    finally:
-      del serialization._STATE_DICT_REGISTRY[Node]
-
   def test_while_loop(self):
     class Foo(nn.Module):