make flax.core.copy add_or_replace optional

.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]
 
@@ -223,7 +227,7 @@ def unfreeze(x: Union[FrozenDict, Dict[str, Any]]) -> Dict[Any, Any]:
 
 def copy(
     x: Union[FrozenDict, Dict[str, Any]],
-    add_or_replace: Union[FrozenDict, Dict[str, Any]],
+    add_or_replace: Optional[Union[FrozenDict, Dict[str, Any]]] = None,
 ) -> Union[FrozenDict, Dict[str, Any]]:
   """Create a new dict with additional and/or replaced entries. This is a utility
   function that can act on either a FrozenDict or regular dict and mimics the
@@ -244,7 +248,8 @@ def copy(
     return x.copy(add_or_replace)
   elif isinstance(x, dict):
     new_dict = jax.tree_map(lambda x: x, x)  # make a deep copy of dict x
-    new_dict.update(add_or_replace)
+    if add_or_replace is not None:
+      new_dict.update(add_or_replace)
     return new_dict
   raise TypeError(f'Expected FrozenDict or dict, got {type(x)}')
 

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}},