Merge pull request google#2958 from chiamp:pop_doc_str

Flax Authors · Flax Authors · commit 8957389db412 · 2023-03-27T13:46:58.000-07:00
PiperOrigin-RevId: 519817606
diff --git a/flax/core/frozen_dict.py b/flax/core/frozen_dict.py
@@ -245,7 +245,8 @@ def copy(x: Union[FrozenDict, Dict[str, Any]], add_or_replace: Union[FrozenDict,
 
 def pop(x: Union[FrozenDict, Dict[str, Any]], key: str) -> Tuple[Union[FrozenDict, Dict[str, Any]], Any]:
   """Create a new dict where one entry is removed. This is a utility
-  function for regular dicts that mimics the behavior of `FrozenDict.pop`.
+  function that can act on either a FrozenDict or regular dict and
+  mimics the behavior of `FrozenDict.pop`.
 
   Example::
 
