Add __copy__ and __deepcopy__ methods to the _Sentinel object used to…

… define the default value for `parent`, to simply return the sentinel object.

PiperOrigin-RevId: 485122748

flax/linen/module.py

@@ -175,7 +175,15 @@ def __init__(self):
 
 
 class _Sentinel:
-  pass
+
+  def __copy__(self):
+    return self  # Do not copy singleton sentinel.
+
+  def __deepcopy__(self, memo):
+    del memo
+    return self  # Do not copy singleton sentinel.
+
+
 _unspecified_parent = _Sentinel()
 
 

tests/linen/linen_module_test.py

@@ -14,9 +14,11 @@
 
 """Tests for flax.linen."""
 
+import copy
 import dataclasses
 import functools
 import gc
+import inspect
 import operator
 from typing import (Any, Callable, Generic, Mapping, NamedTuple, Sequence,
                     Tuple, TypeVar)
@@ -1790,6 +1792,17 @@ def __call__(self, input):
     with self.assertRaises(errors.IncorrectPostInitOverrideError):
       r.init(jax.random.PRNGKey(2), jnp.ones(3))
 
+  def test_deepcopy_unspecified_parent(self):
+    parent_parameter = inspect.signature(DummyModule).parameters['parent']
+    unspecified_parent = parent_parameter.default
+
+    self.assertIs(unspecified_parent,
+                  copy.copy(unspecified_parent))
+
+    self.assertIs(unspecified_parent,
+                  copy.deepcopy(unspecified_parent))
+
+
 class LeakTests(absltest.TestCase):
 
   def test_tracer_leaks(self):