Make FlatState a Mapping instead of a dict

Also, make traverse_util.flatten_dict accept mappings.

Fixes #3879

flax/experimental/nnx/nnx/state.py

@@ -27,6 +27,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+from collections.abc import Mapping
 import typing as tp
 import typing_extensions as tpe
 
@@ -42,7 +43,7 @@
 A = tp.TypeVar('A')
 
 StateLeaf = tp.Union[VariableState[tp.Any], np.ndarray, jax.Array]
-FlatState = dict[PathParts, StateLeaf]
+FlatState = Mapping[PathParts, StateLeaf]
 
 
 def is_state_leaf(x: tp.Any) -> tpe.TypeGuard[StateLeaf]:
@@ -66,8 +67,8 @@ class State(tp.MutableMapping[Key, tp.Any], reprlib.Representable):
   def __init__(
     self,
     mapping: tp.Union[
-      tp.Mapping[Key, tp.Mapping | StateLeaf],
-      tp.Iterator[tuple[Key, tp.Mapping | StateLeaf]],
+      Mapping[Key, Mapping | StateLeaf],
+      tp.Iterator[tuple[Key, Mapping | StateLeaf]],
     ],
     /,
   ):

flax/traverse_util.py

@@ -1,3 +1,5 @@
+# pytype: skip-file
+# Skipped until https://github.com/google/pytype/issues/1619 is resolved.
 # Copyright 2024 The Flax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -52,18 +54,19 @@
 Traversals never mutate the original data. Therefore, an update essentially
 returns a copy of the data including the provided updates.
 """
+from __future__ import annotations
 
 import abc
+from collections.abc import Callable, Mapping
 import copy
 import dataclasses
 import warnings
-from typing import Any, Callable
+from typing import Any, overload
 
 import jax
 
 import flax
-from flax.core.scope import VariableDict
-from flax.typing import PathParts
+from flax.typing import PathParts, VariableDict
 
 from . import struct
 
@@ -77,7 +80,37 @@ class _EmptyNode:
 empty_node = _EmptyNode()
 
 
-def flatten_dict(xs, keep_empty_nodes=False, is_leaf=None, sep=None):
+# TODO: In Python 3.10, use TypeAlias.
+IsLeafCallable = Callable[[tuple[Any, ...], Mapping[Any, Any]], bool]
+
+
+@overload
+def flatten_dict(xs: Mapping[Any, Any],
+                 /,
+                 *,
+                 keep_empty_nodes: bool = False,
+                 is_leaf: None | IsLeafCallable = None,
+                 sep: None = None
+                 ) -> dict[tuple[Any, ...], Any]:
+  ...
+
+@overload
+def flatten_dict(xs: Mapping[Any, Any],
+                 /,
+                 *,
+                 keep_empty_nodes: bool = False,
+                 is_leaf: None | IsLeafCallable = None,
+                 sep: str,
+                 ) -> dict[str, Any]:
+  ...
+
+def flatten_dict(xs: Mapping[Any, Any],
+                 /,
+                 *,
+                 keep_empty_nodes: bool = False,
+                 is_leaf: None | IsLeafCallable = None,
+                 sep: None | str = None
+                 ) -> dict[Any, Any]:
   """Flatten a nested dictionary.
 
   The nested keys are flattened to a tuple.
@@ -111,16 +144,16 @@ def flatten_dict(xs, keep_empty_nodes=False, is_leaf=None, sep=None):
     The flattened dictionary.
   """
   assert isinstance(
-    xs, (flax.core.FrozenDict, dict)
-  ), f'expected (frozen)dict; got {type(xs)}'
+    xs, Mapping
+  ), f'expected Mapping; got {type(xs).__qualname__}'
 
-  def _key(path):
+  def _key(path: tuple[Any, ...]) -> tuple[Any, ...] | str:
     if sep is None:
       return path
     return sep.join(path)
 
-  def _flatten(xs, prefix):
-    if not isinstance(xs, (flax.core.FrozenDict, dict)) or (
+  def _flatten(xs: Any, prefix: tuple[Any, ...]) -> dict[Any, Any]:
+    if not isinstance(xs, Mapping) or (
       is_leaf and is_leaf(prefix, xs)
     ):
       return {_key(prefix): xs}