[nnx] refactor GraphDef

Author: cgarciae

benchmarks/nnx_graph_overhead.py

@@ -97,11 +97,9 @@ def main(argv):
     def step_nnx(model: MLP, optimizer: nnx.Optimizer):
       pass
 
-    cached_step_nnx = nnx.cached_partial(step_nnx, model, optimizer)
-
     t0 = time()
     for _ in range(total_steps):
-      cached_step_nnx()
+      step_nnx(model, optimizer)
 
     total_time = time() - t0
     time_per_step = total_time / total_steps

flax/nnx/bridge/variables.py

@@ -154,7 +154,7 @@ def nnx_attrs_to_linen_vars(nnx_attrs: dict) -> dict:
     elif isinstance(v, variablelib.VariableState):
       col_name = variablelib.variable_name_from_type(v.type)
       v = to_linen_var(v)
-    elif isinstance(v, graph.NodeDef) or isinstance(v, graph.NodeRef):
+    elif isinstance(v, graph.GraphDef):
       col_name = 'nnx'  # an nnx.GraphDef for some ToLinen submodule
     else:
       raise ValueError(f'Cannot infer collection name from value: {v}')

flax/nnx/bridge/wrappers.py

@@ -38,7 +38,7 @@
 @dataclasses.dataclass
 class Functional(tp.Generic[M]):
   module_type: tp.Type[M]
-  graphdef: tp.Optional[graph.NodeDef[M]]
+  graphdef: tp.Optional[graph.GraphDef[M]]
   args: tuple[tp.Any, ...]
   kwargs: dict[str, tp.Any]
 
@@ -48,7 +48,6 @@ def init(self, *, rngs: tp.Optional[Rngs] = None) -> State:
       kwargs['rngs'] = rngs
     module = self.module_type(*self.args, **self.kwargs, **kwargs)
     graphdef, state = nnx.split(module)
-    assert type(graphdef) is graph.NodeDef
     self.graphdef = graphdef
     return state  # type: ignore
 
@@ -217,7 +216,7 @@ class ToLinen(linen.Module):
     >>> variables.keys()
     dict_keys(['nnx', 'params'])
     >>> type(variables['nnx']['graphdef'])
-    <class 'flax.nnx.graph.NodeDef'>
+    <class 'flax.nnx.graph.GraphDef'>
 
   Args:
     nnx_class: The NNX Module class (not instance!).