diff --git a/flax/core/frozen_dict.py b/flax/core/frozen_dict.py
index 37a3ad76..5500a24a 100644
--- 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::