make flax.core.copy add_or_replace optional

ignore .envrc (direnv files)


format


add test


fix

.gitignore

@@ -12,4 +12,7 @@ build/
 .pytype
 .vscode/*
 /.devcontainer
-docs/**/tmp
+docs/**/tmp
+
+# used by direnv
+.envrc

flax/core/frozen_dict.py

@@ -15,7 +15,7 @@
 """Frozen Dictionary."""
 
 import collections
-from typing import Any, TypeVar, Mapping, Dict, Tuple, Union, Hashable
+from typing import Any, Dict, Hashable, Optional, Mapping, Tuple, TypeVar, Union
 
 from flax import serialization
 import jax
@@ -111,7 +111,11 @@ def __hash__(self):
       self._hash = h
     return self._hash
 
-  def copy(self, add_or_replace: Mapping[K, V]) -> 'FrozenDict[K, V]':
+  def copy(
+      self, add_or_replace: Optional[Mapping[K, V]] = None
+  ) -> 'FrozenDict[K, V]':
+    if add_or_replace is None:
+      add_or_replace = {}
     """Create a new FrozenDict with additional or replaced entries."""
     return type(self)({**self, **unfreeze(add_or_replace)})  # type: ignore[arg-type]
 

tests/core/core_frozen_dict_test.py

@@ -121,6 +121,18 @@ def test_utility_copy(self, x, add_or_replace, actual_new_x):
         new_x == actual_new_x and isinstance(new_x, type(actual_new_x))
     )
 
+  @parameterized.parameters(
+      {
+          'x': {'a': 1, 'b': {'c': 2}},
+      },
+      {
+          'x': FrozenDict({'a': 1, 'b': {'c': 2}}),
+      },
+  )
+  def test_utility_copy_singlearg(self, x):
+    new_x = copy(x)
+    self.assertTrue(new_x == x and isinstance(new_x, type(x)))
+
   @parameterized.parameters(
       {
           'x': {'a': 1, 'b': {'c': 2}},