fixes

google · Aug 1, 2023 · 1ff1890 · 1ff1890
1 parent a30f09f
commit 1ff1890
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flax/core/frozen_dict.py b/flax/core/frozen_dict.py
@@ -16,6 +16,7 @@
 
 import collections
 from typing import Any, Dict, Hashable, Optional, Mapping, Tuple, TypeVar, Union
+from types import MappingProxyType
 
 from flax import serialization
 import jax
@@ -246,8 +247,7 @@ 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
-    if add_or_replace is not None:
-      new_dict.update(add_or_replace)
+    new_dict.update(add_or_replace)
     return new_dict
   raise TypeError(f'Expected FrozenDict or dict, got {type(x)}')