Merge pull request #3796 from chiamp:graphnode

PiperOrigin-RevId: 620320193
google · Mar 29, 2024 · 514c111 · 514c111
2 parents 4671ddd + 1eb988d
commit 514c111
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flax/experimental/nnx/__init__.py b/flax/experimental/nnx/__init__.py
@@ -24,10 +24,10 @@
 from .nnx.filterlib import All as All
 from .nnx.filterlib import Not as Not
 from .nnx.graph_utils import GraphDef as GraphDef
+from .nnx.graph_utils import GraphNode as GraphNode
 from .nnx.helpers import Dict as Dict
 from .nnx.helpers import List as List
 from .nnx.helpers import TrainState as TrainState
-from .nnx.module import GraphDef as GraphDef
 from .nnx.module import M as M
 from .nnx.module import Module as Module
 from .nnx.graph_utils import merge as merge