From 59ffe8d722906db7c96477b17e648be16f0c31ab Mon Sep 17 00:00:00 2001 From: Cristian Garcia Date: Sat, 30 Mar 2024 11:22:12 +0000 Subject: [PATCH] [nnx] non-str State keys --- flax/experimental/nnx/nnx/filterlib.py | 22 +- flax/experimental/nnx/nnx/graph_utils.py | 193 +++++++++--------- flax/experimental/nnx/nnx/helpers.py | 21 ++ flax/experimental/nnx/nnx/module.py | 4 +- flax/experimental/nnx/nnx/rnglib.py | 3 +- flax/experimental/nnx/nnx/state.py | 25 +-- .../nnx/tests/nn/test_attention.py | 7 +- .../nnx/tests/test_graph_utils.py | 18 +- flax/experimental/nnx/tests/test_module.py | 6 +- .../nnx/tests/test_partitioning.py | 8 +- .../experimental/nnx/tests/test_transforms.py | 16 +- flax/typing.py | 11 +- 12 files changed, 179 insertions(+), 155 deletions(-) diff --git a/flax/experimental/nnx/nnx/filterlib.py b/flax/experimental/nnx/nnx/filterlib.py index dfb3e4792d..6bb3d23be5 100644 --- a/flax/experimental/nnx/nnx/filterlib.py +++ b/flax/experimental/nnx/nnx/filterlib.py @@ -14,7 +14,7 @@ import builtins import dataclasses -from flax.typing import Path +from flax.typing import PathParts import typing as tp if tp.TYPE_CHECKING: @@ -22,7 +22,7 @@ else: ellipsis = tp.Any -Predicate = tp.Callable[[Path, tp.Any], bool] +Predicate = tp.Callable[[PathParts, tp.Any], bool] FilterLiteral = tp.Union[type, str, Predicate, bool, ellipsis, None] Filter = tp.Union[FilterLiteral, tuple[FilterLiteral, ...], list[FilterLiteral]] @@ -48,17 +48,17 @@ def to_predicate(filter: Filter) -> Predicate: @dataclasses.dataclass class AtPath: - path: str + str_key: str - def __call__(self, path: Path, x: tp.Any): - return self.path == path + def __call__(self, path: PathParts, x: tp.Any): + return self.str_key in path @dataclasses.dataclass class OfType: type: type - def __call__(self, path: Path, x: tp.Any): + def __call__(self, path: PathParts, x: tp.Any): return isinstance(x, self.type) @@ -68,7 +68,7 @@ def __init__(self, *filters: Filter): to_predicate(collection_filter) for collection_filter in filters ) - def __call__(self, path: Path, x: tp.Any): + def __call__(self, path: PathParts, x: tp.Any): return any(predicate(path, x) for predicate in self.predicates) @@ -78,7 +78,7 @@ def __init__(self, *filters: Filter): to_predicate(collection_filter) for collection_filter in filters ) - def __call__(self, path: Path, x: tp.Any): + def __call__(self, path: PathParts, x: tp.Any): return all(predicate(path, x) for predicate in self.predicates) @@ -86,15 +86,15 @@ class Not: def __init__(self, collection_filter: Filter): self.predicate = to_predicate(collection_filter) - def __call__(self, path: Path, x: tp.Any): + def __call__(self, path: PathParts, x: tp.Any): return not self.predicate(path, x) class Everything: - def __call__(self, path: Path, x: tp.Any): + def __call__(self, path: PathParts, x: tp.Any): return True class Nothing: - def __call__(self, path: Path, x: tp.Any): + def __call__(self, path: PathParts, x: tp.Any): return False diff --git a/flax/experimental/nnx/nnx/graph_utils.py b/flax/experimental/nnx/nnx/graph_utils.py index b398b4eba2..d56b913502 100644 --- a/flax/experimental/nnx/nnx/graph_utils.py +++ b/flax/experimental/nnx/nnx/graph_utils.py @@ -37,11 +37,16 @@ CallableProxy, DelayedAccessor, ) -from flax.experimental.nnx.nnx.state import State, StateLeaf, is_state_leaf +from flax.experimental.nnx.nnx.state import ( + FlatState, + State, + StateLeaf, + is_state_leaf, +) from flax.experimental.nnx.nnx.rnglib import Rngs from flax.experimental.nnx.nnx.state import State from flax.experimental.nnx.nnx.variables import EMPTY, Empty, Variable -from flax.typing import Path, PathParts +from flax.typing import PathParts, Key A = tp.TypeVar('A') B = tp.TypeVar('B') @@ -129,28 +134,28 @@ def __str__(self) -> str: @dataclasses.dataclass(frozen=True) class NodeImplBase(tp.Generic[Node, Leaf, AuxData]): type: type - flatten: tp.Callable[[Node], tuple[tp.Sequence[tuple[str, Leaf]], AuxData]] + flatten: tp.Callable[[Node], tuple[tp.Sequence[tuple[Key, Leaf]], AuxData]] - def node_dict(self, node: Node) -> dict[str, Leaf]: + def node_dict(self, node: Node) -> dict[Key, Leaf]: nodes, _ = self.flatten(node) return dict(nodes) @dataclasses.dataclass(frozen=True) class GraphNodeImpl(NodeImplBase[Node, Leaf, AuxData]): - set_key: tp.Callable[[Node, str, Leaf], None] - pop_key: tp.Callable[[Node, str], Leaf] + set_key: tp.Callable[[Node, Key, Leaf], None] + pop_key: tp.Callable[[Node, Key], Leaf] create_empty: tp.Callable[[AuxData], Node] clear: tp.Callable[[Node, AuxData], None] - def init(self, node: Node, items: tuple[tuple[str, Leaf], ...]): + def init(self, node: Node, items: tuple[tuple[Key, Leaf], ...]): for key, value in items: self.set_key(node, key, value) @dataclasses.dataclass(frozen=True) class PytreeNodeImpl(NodeImplBase[Node, Leaf, AuxData]): - unflatten: tp.Callable[[tuple[tuple[str, Leaf], ...], AuxData], Node] + unflatten: tp.Callable[[tuple[tuple[Key, Leaf], ...], AuxData], Node] NodeImpl = tp.Union[ @@ -160,9 +165,9 @@ class PytreeNodeImpl(NodeImplBase[Node, Leaf, AuxData]): def register_graph_node_type( type: type, - flatten: tp.Callable[[Node], tuple[tp.Sequence[tuple[str, Leaf]], AuxData]], - set_key: tp.Callable[[Node, str, Leaf], None], - pop_key: tp.Callable[[Node, str], Leaf], + flatten: tp.Callable[[Node], tuple[tp.Sequence[tuple[Key, Leaf]], AuxData]], + set_key: tp.Callable[[Node, Key, Leaf], None], + pop_key: tp.Callable[[Node, Key], Leaf], create_empty: tp.Callable[[AuxData], Node], clear: tp.Callable[[Node, AuxData], None], ): @@ -243,7 +248,7 @@ def __repr__(self) -> str: @dataclasses.dataclass(repr=False) class _MappingRepr(reprlib.Representable): - mapping: tp.Mapping[str, tp.Any] + mapping: tp.Mapping[Key, tp.Any] def __nnx_repr__(self): yield reprlib.Object(type='', value_sep=': ', start='{', end='}') @@ -263,7 +268,7 @@ def __init__( self, type: tp.Type[Variable[tp.Any]], index: int, - metadata: dict[str, tp.Any], + metadata: dict[Key, tp.Any], ): self._type = type self._index = index @@ -332,10 +337,10 @@ def __init__( self, type: tp.Type[Node], index: int, - attributes: tuple[str, ...], - subgraphs: tp.Iterable[tuple[str, tp.Union['GraphDef[tp.Any]', int]]], - static_fields: tp.Iterable[tuple[str, tp.Any]], - variables: tp.Iterable[tuple[str, VariableDef | int]], + attributes: tuple[Key, ...], + subgraphs: tp.Iterable[tuple[Key, tp.Union['GraphDef[tp.Any]', int]]], + static_fields: tp.Iterable[tuple[Key, tp.Any]], + variables: tp.Iterable[tuple[Key, VariableDef | int]], metadata: tp.Any, ): self._type: type[Node] = type @@ -433,10 +438,10 @@ def _graphdef_unflatten( metadata: tuple[ tp.Type[Node], int, - tuple[str, ...], - tuple[tuple[str, GraphDef[Node] | int], ...], - tuple[tuple[str, tp.Any], ...], - tuple[tuple[str, Variable[Empty] | int], ...], + tuple[Key, ...], + tuple[tuple[Key, GraphDef[Node] | int], ...], + tuple[tuple[Key, tp.Any], ...], + tuple[tuple[Key, Variable[Empty] | int], ...], tp.Any, ], _, @@ -454,7 +459,7 @@ def graph_flatten( /, ) -> tuple[GraphDef[Node], State, tp.Mapping[tp.Any, Index]]: ref_to_index = RefMap[tp.Any, Index]() - flat_state: dict[Path, StateLeaf] = {} + flat_state: dict[PathParts, StateLeaf] = {} graphdef = _graph_flatten((), ref_to_index, flat_state, x) assert not isinstance(graphdef, int) return graphdef, State.from_flat_path(flat_state), ref_to_index @@ -463,7 +468,7 @@ def graph_flatten( def _graph_flatten( path: PathParts, ref_to_index: RefMap[tp.Any, Index], - flat_state: dict[Path, StateLeaf], + flat_state: dict[PathParts, StateLeaf], node: Node, ) -> GraphDef[Node] | int: if not is_node(node): @@ -481,33 +486,26 @@ def _graph_flatten( else: index = -1 - subgraphs: list[tuple[str, tp.Union[GraphDef[Node], int]]] = [] - static_fields: list[tuple[str, tp.Any]] = [] - variables: list[tuple[str, VariableDef | int]] = [] + subgraphs: list[tuple[Key, tp.Union[GraphDef[Node], int]]] = [] + static_fields: list[tuple[Key, tp.Any]] = [] + variables: list[tuple[Key, VariableDef | int]] = [] values, metadata = node_impl.flatten(node) for key, value in values: - if not isinstance(key, str): - raise TypeError( - f'Node (of type {type(node).__name__}) has a key of non-string ' - f'type {type(key).__name__}.' - ) if is_node(value): graphdef = _graph_flatten((*path, key), ref_to_index, flat_state, value) subgraphs.append((key, graphdef)) elif isinstance(value, Variable): - str_path = '/'.join((*path, key)) if value in ref_to_index: variables.append((key, ref_to_index[value])) else: - flat_state[str_path] = value.copy() + flat_state[(*path, key)] = value.copy() variable_index = ref_to_index[value] = len(ref_to_index) variables.append( (key, VariableDef.from_variable(value, variable_index)) ) elif is_state_leaf(value): - str_path = '/'.join((*path, key)) - flat_state[str_path] = value + flat_state[(*path, key)] = value else: static_fields.append((key, value)) @@ -702,16 +700,16 @@ def graph_pop( id_to_index: dict[int, Index] = {} path_parts: PathParts = () predicates = tuple(filterlib.to_predicate(filter) for filter in filters) - states = tuple({} for _ in predicates) - _graph_pop(node, id_to_index, path_parts, states, predicates) - return tuple(State(x) for x in states) + flat_states: tuple[FlatState, ...] = tuple({} for _ in predicates) + _graph_pop(node, id_to_index, path_parts, flat_states, predicates) + return tuple(State.from_flat_path(flat_state) for flat_state in flat_states) def _graph_pop( node: tp.Any, id_to_index: dict[int, Index], path_parts: PathParts, - states: tuple[dict[Path, tp.Any], ...], + flat_states: tuple[FlatState, ...], predicates: tuple[filterlib.Predicate, ...], ) -> None: if not is_node(node): @@ -726,17 +724,19 @@ def _graph_pop( for name, value in node_dict.items(): if is_node(value): - _graph_pop(value, id_to_index, (*path_parts, name), states, predicates) + _graph_pop( + value, id_to_index, (*path_parts, name), flat_states, predicates + ) continue elif not is_state_leaf(value): continue elif id(value) in id_to_index: continue - path = '/'.join((*path_parts, name)) + node_path = (*path_parts, name) node_impl = get_node_impl(node) - for state, predicate in zip(states, predicates): - if predicate(path, value): + for state, predicate in zip(flat_states, predicates): + if predicate(node_path, value): if isinstance(node_impl, PytreeNodeImpl): raise ValueError( f'Cannot pop key {name!r} from node of type {type(node).__name__}' @@ -745,7 +745,7 @@ def _graph_pop( node_impl.pop_key(node, name) if isinstance(value, Variable): value = value.copy() - state[path] = value + state[node_path] = value break else: # NOTE: should we raise an error here? @@ -1009,7 +1009,7 @@ def clone(node: Node) -> Node: return static.merge(state) -def iter_nodes(node: tp.Any) -> tp.Iterator[tuple[Path, tp.Any]]: +def iter_nodes(node: tp.Any) -> tp.Iterator[tuple[PathParts, tp.Any]]: visited: set[int] = set() path_parts: PathParts = () yield from _iter_nodes(node, visited, path_parts) @@ -1017,14 +1017,13 @@ def iter_nodes(node: tp.Any) -> tp.Iterator[tuple[Path, tp.Any]]: def _iter_nodes( node: tp.Any, visited: set[int], path_parts: PathParts -) -> tp.Iterator[tuple[Path, tp.Any]]: +) -> tp.Iterator[tuple[PathParts, tp.Any]]: if not is_node(node): return if id(node) in visited: return visited.add(id(node)) - path = '/'.join(path_parts) - yield path, node + yield path_parts, node node_impl = get_node_impl(node) node_dict = node_impl.node_dict(node) for key, value in node_dict.items(): @@ -1158,6 +1157,18 @@ class GraphNode(reprlib.Representable, metaclass=GraphNodeMeta): if tp.TYPE_CHECKING: _graph_node__state: ModuleState + def __init_subclass__(cls) -> None: + super().__init_subclass__() + + graph_utils.register_graph_node_type( + type=cls, + flatten=cls._graph_node_flatten, + set_key=cls._graph_node_set_key, + pop_key=cls._graph_node_pop_key, + create_empty=cls._graph_node_create_empty, + clear=cls._graph_node_clear, + ) + if not tp.TYPE_CHECKING: def __setattr__(self, name: str, value: Any) -> None: @@ -1204,55 +1215,43 @@ def __nnx_repr__(self): if clear_seen: CONTEXT.seen_modules_repr = None - def __init_subclass__(cls) -> None: - super().__init_subclass__() - - graph_utils.register_graph_node_type( - type=cls, - flatten=_graph_node_flatten, - set_key=_graph_node_set_key, - pop_key=_graph_node_pop_key, - create_empty=_graph_node_create_empty, - clear=_graph_node_clear, + # Graph Definition + def _graph_node_flatten(self): + nodes = sorted( + (key, value) + for key, value in vars(self).items() + if key != '_graph_node__state' ) + return nodes, type(self) + def _graph_node_set_key(self, key: Key, value: tp.Any): + if not isinstance(key, str): + raise KeyError(f'Invalid key: {key!r}') + elif ( + hasattr(self, key) + and isinstance(variable := getattr(self, key), Variable) + and isinstance(value, Variable) + ): + variable.copy_from(value) + else: + setattr(self, key, value) -# Graph Definition -def _graph_node_flatten(node: GraphNode): - nodes = tuple( - (name, value) - for name, value in sorted(vars(node).items()) - if name != '_graph_node__state' - ) - return nodes, type(node) - - -def _graph_node_set_key(node: GraphNode, name: str, value: tp.Any): - if ( - hasattr(node, name) - and isinstance(variable := getattr(node, name), Variable) - and isinstance(value, Variable) - ): - variable.copy_from(value) - else: - setattr(node, name, value) - - -def _graph_node_pop_key(node: GraphNode, name: str): - return vars(node).pop(name) - - -def _graph_node_create_empty(cls: tp.Type[G]) -> G: - node = object.__new__(cls) - vars(node).update(_graph_node__state=ModuleState()) - return node + def _graph_node_pop_key(self, key: Key): + if not isinstance(key, str): + raise KeyError(f'Invalid key: {key!r}') + return vars(self).pop(key) + @staticmethod + def _graph_node_create_empty(node_type: tp.Type[G]) -> G: + node = object.__new__(node_type) + vars(node).update(_graph_node__state=ModuleState()) + return node -def _graph_node_clear(node: GraphNode, cls: tp.Type[G]): - module_state = node._graph_node__state - module_vars = vars(node) - module_vars.clear() - module_vars['_graph_node__state'] = module_state + def _graph_node_clear(self, cls: tp.Type[G]): + module_state = self._graph_node__state + module_vars = vars(self) + module_vars.clear() + module_vars['_graph_node__state'] = module_state # --------------------------------------------------------- @@ -1266,13 +1265,13 @@ def is_pytree_node(x: tp.Any) -> bool: return not jax.tree_util.all_leaves([x]) -def _key_path_to_str(key: tp.Any) -> str: +def _key_path_to_key(key: tp.Any) -> Key: if isinstance(key, jax.tree_util.SequenceKey): - return str(key.idx) + return key.idx elif isinstance( key, (jax.tree_util.DictKey, jax.tree_util.FlattenedIndexKey) ): - return str(key.key) + return key.key elif isinstance(key, jax.tree_util.GetAttrKey): return key.name else: @@ -1283,13 +1282,13 @@ def _flatten_pytree(pytree: tp.Any): leaves, treedef = jax.tree_util.tree_flatten_with_path( pytree, is_leaf=lambda x: x is not pytree ) - nodes = tuple((_key_path_to_str(path[0]), value) for path, value in leaves) + nodes = tuple((_key_path_to_key(path[0]), value) for path, value in leaves) return nodes, treedef def _unflatten_pytree( - nodes: tuple[tuple[str, tp.Any], ...], treedef: jax.tree_util.PyTreeDef + nodes: tuple[tuple[Key, tp.Any], ...], treedef: jax.tree_util.PyTreeDef ): pytree = treedef.unflatten(value for _, value in nodes) return pytree diff --git a/flax/experimental/nnx/nnx/helpers.py b/flax/experimental/nnx/nnx/helpers.py index 0b4c86dd5b..137fe11bd2 100644 --- a/flax/experimental/nnx/nnx/helpers.py +++ b/flax/experimental/nnx/nnx/helpers.py @@ -34,6 +34,7 @@ import jax.numpy as jnp import optax +from flax.experimental.nnx.nnx.graph_utils import Key from flax.experimental.nnx.nnx.module import GraphDef, Module from flax.experimental.nnx.nnx.proxy_caller import ApplyCaller from flax.experimental.nnx.nnx.rnglib import Rngs @@ -105,6 +106,26 @@ def __iter__(self) -> tp.Iterator[A]: def __len__(self) -> int: return self._length + def _graph_node_flatten(self): + nodes: list[tuple[Key, tp.Any]] = sorted( + (int(key), value) + for key, value in vars(self).items() + if key not in ('_graph_node__state', '_length') + ) + nodes.append(('_length', self._length)) + return nodes, type(self) + + def _graph_node_set_key(self, key: Key, value: tp.Any): + if isinstance(key, int): + key = str(key) + return super()._graph_node_set_key(key, value) + + def _graph_node_pop_key(self, key: Key): + if isinstance(key, int): + key = str(key) + return super()._graph_node_pop_key(key) + + class Sequential(List): def __call__(self, *args, rngs: tp.Optional[Rngs] = None, **kwargs) -> tp.Any: output: tp.Any = None diff --git a/flax/experimental/nnx/nnx/module.py b/flax/experimental/nnx/nnx/module.py index fb507e9eee..ae92cdce48 100644 --- a/flax/experimental/nnx/nnx/module.py +++ b/flax/experimental/nnx/nnx/module.py @@ -35,7 +35,7 @@ from flax.experimental.nnx.nnx.rnglib import Rngs from flax.experimental.nnx.nnx.state import State from flax.experimental.nnx.nnx.variables import Variable -from flax.typing import Path +from flax.typing import Path, PathParts A = tp.TypeVar('A') B = tp.TypeVar('B') @@ -265,7 +265,7 @@ def sow( reduced_value = reduce_fn(init_fn(), value) setattr(self, name, variable_type(reduced_value)) - def modules(self) -> tp.Iterator[tuple[Path, Module]]: + def modules(self) -> tp.Iterator[tuple[PathParts, Module]]: for path, value in graph_utils.iter_nodes(self): if isinstance(value, Module): yield path, value diff --git a/flax/experimental/nnx/nnx/rnglib.py b/flax/experimental/nnx/nnx/rnglib.py index 70307b67d4..99776ac47b 100644 --- a/flax/experimental/nnx/nnx/rnglib.py +++ b/flax/experimental/nnx/nnx/rnglib.py @@ -183,7 +183,8 @@ def fork( for name, stream in self._rngs.items(): for predicate, pattern in predicate_pattern: - if predicate(name, stream): + stream_path = (name,) + if predicate(stream_path, stream): fork = stream.fork(pattern) if pattern is None: broadcasts[name] = fork diff --git a/flax/experimental/nnx/nnx/state.py b/flax/experimental/nnx/nnx/state.py index 1aca5bb7ba..f287391a45 100644 --- a/flax/experimental/nnx/nnx/state.py +++ b/flax/experimental/nnx/nnx/state.py @@ -37,13 +37,12 @@ from flax import traverse_util from flax.experimental.nnx.nnx import filterlib, reprlib from flax.experimental.nnx.nnx.variables import Variable -from flax.typing import Path +from flax.typing import Key, PathParts A = tp.TypeVar('A') -Key = str StateLeaf = tp.Union[Variable[tp.Any], np.ndarray, jax.Array] -FlatState = dict[Path, StateLeaf] +FlatState = dict[PathParts, StateLeaf] def is_state_leaf(x: tp.Any) -> tpe.TypeGuard[StateLeaf]: @@ -81,9 +80,7 @@ def __init__( def raw_mapping(self) -> dict[Key, dict[str, tp.Any] | tp.Any]: return self._mapping - def __getitem__(self, key: Key | int) -> State | StateLeaf: - if isinstance(key, int): - key = str(key) + def __getitem__(self, key: Key) -> State | StateLeaf: value = self._mapping[key] if is_state_leaf(value): return value @@ -94,10 +91,7 @@ def __getattr__(self, key: Key) -> State | StateLeaf: raise AttributeError(f"No attribute '{key}' in State") return self[key] - def __setitem__(self, key: Key | int, value: State | StateLeaf) -> None: - if isinstance(key, int): - key = str(key) - + def __setitem__(self, key: Key, value: State | StateLeaf) -> None: if isinstance(value, State): self._mapping[key] = value._mapping else: @@ -122,12 +116,12 @@ def __nnx_repr__(self): v = NestedStateRepr(v) yield reprlib.Attr(repr(k), v) - def flat_state(self) -> dict[Key, Variable[Variable]]: - return traverse_util.flatten_dict(self._mapping, sep='/') # type: ignore + def flat_state(self) -> dict[PathParts, StateLeaf]: + return traverse_util.flatten_dict(self._mapping) # type: ignore @classmethod def from_flat_path(cls, flat_state: FlatState, /) -> State: - nested_state = traverse_util.unflatten_dict(flat_state, sep='/') + nested_state = traverse_util.unflatten_dict(flat_state) return cls(nested_state) @tp.overload @@ -226,7 +220,6 @@ def __sub__(self, other: 'State') -> 'State': return State.from_flat_path(diff) - def _state_flatten_with_keys(x: State): items = sorted(x._mapping.items(), key=lambda item: item[0]) children = tuple((jtu.DictKey(key), value) for key, value in items) @@ -234,8 +227,8 @@ def _state_flatten_with_keys(x: State): def _state_unflatten( - static: tuple[Path, ...], - leaves: tuple[Variable, ...] | tuple[dict[str, Variable]], + static: tuple[Key, ...], + leaves: tuple[StateLeaf, ...] | tuple[dict[Key, StateLeaf]], ): return State(zip(static, leaves)) diff --git a/flax/experimental/nnx/tests/nn/test_attention.py b/flax/experimental/nnx/tests/nn/test_attention.py index f312db7321..6ffca788ce 100644 --- a/flax/experimental/nnx/tests/nn/test_attention.py +++ b/flax/experimental/nnx/tests/nn/test_attention.py @@ -69,11 +69,12 @@ def __call__(self, x, sow_weights=False): _ = module(x, True) intermediates = module.pop(nnx.Intermediate) - assert intermediates['attention_layers/0/attention_weights'].raw_value[ + # assert intermediates['attention_layers/0/attention_weights'].raw_value[ + assert intermediates['attention_layers'][0]['attention_weights'].raw_value[ 0 ].shape == (4, 8, 6, 6) - assert 'attention_layers/1/attention_weights' not in intermediates - assert intermediates['attention_layers/2/attention_weights'].raw_value[ + assert 1 not in intermediates['attention_layers'] + assert intermediates['attention_layers'][2]['attention_weights'].raw_value[ 0 ].shape == (4, 8, 6, 6) diff --git a/flax/experimental/nnx/tests/test_graph_utils.py b/flax/experimental/nnx/tests/test_graph_utils.py index 85a30f4c5e..80ad57a9fe 100644 --- a/flax/experimental/nnx/tests/test_graph_utils.py +++ b/flax/experimental/nnx/tests/test_graph_utils.py @@ -27,8 +27,8 @@ def test_flatten(self): static, state, ref_idx = nnx.graph_utils.graph_flatten(g) - state['0']['b'].raw_value = 2 - state['3'].raw_value = 4 + state[0]['b'].raw_value = 2 + state[3].raw_value = 4 assert len(ref_idx) == 2 assert a['b'] in ref_idx @@ -69,7 +69,7 @@ def test_update_dynamic(self): static, state, _ = nnx.graph_utils.graph_flatten(g) - state['0']['b'].raw_value = 3 + state[0]['b'].raw_value = 3 nnx.graph_utils.graph_update_dynamic(g, state) assert g[0]['b'].raw_value == 3 @@ -125,12 +125,12 @@ def test_module_list(self): static, state, _ = nnx.graph_utils.graph_flatten(ls) - assert state['0']['kernel'].raw_value.shape == (2, 2) - assert state['0']['bias'].raw_value.shape == (2,) - assert state['1']['scale'].raw_value.shape == (2,) - assert state['1']['bias'].raw_value.shape == (2,) - assert state['1']['mean'].raw_value.shape == (2,) - assert state['1']['var'].raw_value.shape == (2,) + assert state[0]['kernel'].raw_value.shape == (2, 2) + assert state[0]['bias'].raw_value.shape == (2,) + assert state[1]['scale'].raw_value.shape == (2,) + assert state[1]['bias'].raw_value.shape == (2,) + assert state[1]['mean'].raw_value.shape == (2,) + assert state[1]['var'].raw_value.shape == (2,) def test_shared_variables(self): v = nnx.Param(1) diff --git a/flax/experimental/nnx/tests/test_module.py b/flax/experimental/nnx/tests/test_module.py index 42832a308c..199b038cfc 100644 --- a/flax/experimental/nnx/tests/test_module.py +++ b/flax/experimental/nnx/tests/test_module.py @@ -649,11 +649,11 @@ def __init__(self, *, rngs: nnx.Rngs): modules = list(module.modules()) assert len(modules) == 3 - assert modules[0][0] == '' + assert modules[0][0] == () assert isinstance(modules[0][1], Foo) - assert modules[1][0] == 'submodules/0/a' + assert modules[1][0] == ('submodules', 0, 'a') assert isinstance(modules[1][1], nnx.Linear) - assert modules[2][0] == 'submodules/1/b' + assert modules[2][0] == ('submodules', 1, 'b') assert isinstance(modules[2][1], nnx.Conv) def test_array_in_module(self): diff --git a/flax/experimental/nnx/tests/test_partitioning.py b/flax/experimental/nnx/tests/test_partitioning.py index 53d03d824b..7ed98c1e7b 100644 --- a/flax/experimental/nnx/tests/test_partitioning.py +++ b/flax/experimental/nnx/tests/test_partitioning.py @@ -33,11 +33,11 @@ def test_partition(self): assert len(rest) == 1 # check params - assert params['a']['0'].raw_value == m.a[0].value + assert params['a'][0].raw_value == m.a[0].value assert params['b'].raw_value == m.b.value # check rest - assert rest['a']['1'].raw_value == m.a[1].value + assert rest['a'][1].raw_value == m.a[1].value m2 = graphdef.merge(params, rest) @@ -152,8 +152,8 @@ def test_get_paritition(self): assert vars(m.a)['0'] is not vars(m)['b'] state = m.extract(nnx.Variable) - assert state['a']['0'].raw_value == m.a[0].value - assert state['a']['1'].raw_value == m.a[1].value + assert state['a'][0].raw_value == m.a[0].value + assert state['a'][1].raw_value == m.a[1].value assert state['b'].raw_value == m.b.value assert state.b is not state.a[0] assert len(state.flat_state()) == 3 diff --git a/flax/experimental/nnx/tests/test_transforms.py b/flax/experimental/nnx/tests/test_transforms.py index 46d541cac0..4a079e0979 100644 --- a/flax/experimental/nnx/tests/test_transforms.py +++ b/flax/experimental/nnx/tests/test_transforms.py @@ -340,10 +340,10 @@ def f(m: nnx.Dict): assert m.a[0] is m.b assert isinstance(grads, nnx.State) - assert grads['a']['0'].raw_value == 2.0 - assert isinstance(grads.a['0'], nnx.Variable) - assert grads['a']['1'].raw_value == 1.0 - assert isinstance(grads.a['1'], nnx.Variable) + assert grads['a'][0].raw_value == 2.0 + assert isinstance(grads.a[0], nnx.Variable) + assert grads['a'][1].raw_value == 1.0 + assert isinstance(grads.a[1], nnx.Variable) assert len(grads.flat_state()) == 2 m.update(grads) @@ -371,8 +371,8 @@ def f(m: nnx.Dict): grads = f(m) assert isinstance(grads, nnx.State) - assert grads['a']['0'].raw_value == 1.0 - assert isinstance(grads.a['0'], nnx.Param) + assert grads['a'][0].raw_value == 1.0 + assert isinstance(grads.a[0], nnx.Param) assert len(grads) == 2 m.update(grads) @@ -399,8 +399,8 @@ def f(m: nnx.Dict): grads = f(m) assert isinstance(grads, nnx.State) - assert grads['a']['1'].raw_value == 1.0 - assert isinstance(grads.a['1'], nnx.BatchStat) + assert grads['a'][1].raw_value == 1.0 + assert isinstance(grads.a[1], nnx.BatchStat) assert len(grads) == 1 m.update(grads) diff --git a/flax/typing.py b/flax/typing.py index 70cd5c3f01..d566a31c8a 100644 --- a/flax/typing.py +++ b/flax/typing.py @@ -17,8 +17,10 @@ Callable, Dict, Generic, + Hashable, Mapping, Optional, + Protocol, Sequence, Tuple, TypeVar, @@ -38,9 +40,16 @@ RNGSequences = Dict[str, PRNGKey] Dtype = Union[jax.typing.DTypeLike, Any] Shape = Sequence[int] +K = TypeVar('K') + + +class Key(Hashable, Protocol): + def __lt__(self: K, value: K, /) -> bool: + ... + Path = str -PathParts = Tuple[str, ...] +PathParts = Tuple[Key, ...] Leaf = Any