rename graph.py -> graphlib.py

PiperOrigin-RevId: 877411273

Author: Cristian Garcia

flax/nnx/__init__.py

@@ -28,9 +28,9 @@
 from .filterlib import Not as Not
 from .filterlib import Everything as Everything
 from .filterlib import Nothing as Nothing
-from .graph import GraphDef as GraphDef
-from .graph import GraphState as GraphState
-from .graph import PureState as PureState
+from .graphlib import GraphDef as GraphDef
+from .graphlib import GraphState as GraphState
+from .graphlib import PureState as PureState
 from . import pytreelib as object
 from .pytreelib import Pytree as Pytree
 from .pytreelib import Object as Object
@@ -52,31 +52,31 @@
 from .module import view as view
 from .module import view_info as view_info
 from .module import iter_children as iter_children, iter_modules as iter_modules
-from .graph import merge as merge
-from .graph import UpdateContext as UpdateContext
-from .graph import update_context as update_context
-from .graph import current_update_context as current_update_context
-from .graph import split as split
-from .graph import update as update
-from .graph import clone as clone
-from .graph import pop as pop
-from .graph import state as state
-from .graph import graphdef as graphdef
-from .graph import iter_graph as iter_graph
-from .graph import recursive_map as recursive_map
-from .graph import find_duplicates as find_duplicates
-from .graph import call as call
-from .graph import set_metadata as set_metadata
-from .graph import SplitContext as SplitContext
-from .graph import split_context as split_context
-from .graph import MergeContext as MergeContext
-from .graph import merge_context as merge_context
-from .graph import variables as variables
-from .graph import vars_as as vars_as
-from .graph import pure as pure
-from .graph import cached_partial as cached_partial
-from .graph import flatten as flatten
-from .graph import unflatten as unflatten
+from .graphlib import merge as merge
+from .graphlib import UpdateContext as UpdateContext
+from .graphlib import update_context as update_context
+from .graphlib import current_update_context as current_update_context
+from .graphlib import split as split
+from .graphlib import update as update
+from .graphlib import clone as clone
+from .graphlib import pop as pop
+from .graphlib import state as state
+from .graphlib import graphdef as graphdef
+from .graphlib import iter_graph as iter_graph
+from .graphlib import recursive_map as recursive_map
+from .graphlib import find_duplicates as find_duplicates
+from .graphlib import call as call
+from .graphlib import set_metadata as set_metadata
+from .graphlib import SplitContext as SplitContext
+from .graphlib import split_context as split_context
+from .graphlib import MergeContext as MergeContext
+from .graphlib import merge_context as merge_context
+from .graphlib import variables as variables
+from .graphlib import vars_as as vars_as
+from .graphlib import pure as pure
+from .graphlib import cached_partial as cached_partial
+from .graphlib import flatten as flatten
+from .graphlib import unflatten as unflatten
 from .nn import initializers as initializers
 from .nn.activations import celu as celu
 from .nn.activations import elu as elu
@@ -210,7 +210,9 @@
 from .extract import NodeStates as NodeStates
 from .summary import tabulate as tabulate
 from . import traversals as traversals
-
+from . import graphlib as graphlib
+# import last to prevent potential import cycles
+from . import graph as graph
 
 import typing as _tp
 

flax/nnx/bridge/interop.py

@@ -15,7 +15,7 @@
 import typing as tp
 
 from flax.linen import module as nn_module
-from flax.nnx import graph, rnglib
+from flax.nnx import graphlib, rnglib
 from flax.nnx.bridge import wrappers
 from flax.nnx.bridge import module as bdg_module
 import flax.nnx.module as nnx_module
@@ -47,15 +47,15 @@ def nnx_in_bridge_mdl(factory: tp.Callable[[rnglib.Rngs], nnx_module.Module],
     module = factory(parent.scope.rngs)
   else:
     rngs = parent.scope.rngs if parent.scope.rngs else rnglib.Rngs(7)  # dummy
-    module = nnx_eval_shape(factory, rngs)
+    module = nnx_eval_shape(factory, rngs, graph=True)
 
     @nnx_jit
     def rng_state(rngs):
-      return graph.state(factory(rngs), rnglib.RngState)
+      return graphlib.state(factory(rngs), rnglib.RngState, graph=True)
 
     # Make sure the internal rng state is not abstract - other vars shall be
     if parent.scope.rngs:
-      graph.update(module, rng_state(parent.scope.rngs))
+      graphlib.update(module, rng_state(parent.scope.rngs))
 
   # Automatically set the attribute if compact. If setup, user is responsible
   # for setting the attribute of the superlayer.

flax/nnx/bridge/module.py

@@ -28,7 +28,7 @@
 from flax.core import meta
 from flax.core.scope import CollectionFilter
 from flax.core.frozen_dict import FrozenDict
-from flax.nnx import graph, rnglib, statelib, traversals
+from flax.nnx import graphlib, rnglib, statelib, traversals
 from flax.nnx import variablelib
 import flax.nnx.module as nnx_module
 from flax.nnx.pytreelib import Pytree
@@ -104,7 +104,7 @@ def _maybe_call_setup(module: Module):
 def _bind_module(parent: Module, module: Module) -> Module:
   assert parent.scope is not None
 
-  for _, value in reversed(list(graph.iter_graph(module, graph=True))):
+  for _, value in reversed(list(graphlib.iter_graph(module, graph=True))):
     if isinstance(value, Module):
       if module.scope is None:
         value.scope = parent.scope.copy()  # type: ignore[attribute-error]
@@ -245,8 +245,8 @@ def _setattr(self, name: str, value: tp.Any) -> None:
       if name in vars(self) and isinstance(
         state := vars(self)[name], ModuleState
       ):
-        graph.update(value, state)
-      for leaf in jax.tree.leaves(value, is_leaf=graph.is_graph_node):
+        graphlib.update(value, state)
+      for leaf in jax.tree.leaves(value, is_leaf=graphlib.is_graph_node):
         if isinstance(leaf, Module):
           leaf._pytree__state._initializing = self.is_initializing()
           _bind_module(self, leaf)
@@ -374,7 +374,7 @@ def variable(  # type: ignore[invalid-annotation]
     return variable
 
   def _get_variables(self) -> tp.Mapping:
-    state = graph.state(self)
+    state = graphlib.state(self, graph=True)
     _variables: dict = {}
 
     variable: variablelib.Variable
@@ -422,7 +422,7 @@ def apply(
     _initialize: bool = False,
     **kwargs,
   ) -> tp.Any:
-    module = graph.clone(self)
+    module = graphlib.clone(self, graph=True)
 
     # create variables
     real_variables = dict(variables)
@@ -440,7 +440,7 @@ def to_variable(value):
 
     states = ({},) if not real_variables else real_variables.values()
     state = statelib.merge_state(*states, cls=ModuleState)
-    graph.update(module, state)
+    graphlib.update(module, state)
 
     if rngs is None:
       rngs = rnglib.Rngs()
@@ -471,7 +471,7 @@ def to_variable(value):
     _method = _get_unbound_fn(_method)
 
     # set temporary state
-    for _, value in graph.iter_graph(module, graph=True):
+    for _, value in graphlib.iter_graph(module, graph=True):
       if isinstance(value, Pytree):
         value._pytree__state._initializing = _initialize
       if isinstance(value, Module):
@@ -486,7 +486,7 @@ def to_variable(value):
     finally:
       MODULE_CONTEXT.module_stack.pop()
       # reset temporary state
-      for _, value in graph.iter_graph(module, graph=True):
+      for _, value in graphlib.iter_graph(module, graph=True):
         if isinstance(value, Pytree):
           value._pytree__state._initializing = False
         if isinstance(value, Module):

flax/nnx/bridge/wrappers.py

@@ -23,7 +23,7 @@
 from flax import nnx
 from flax.core import FrozenDict
 from flax.core import meta
-from flax.nnx import graph
+from flax.nnx import graphlib
 from flax.nnx import variablelib
 from flax.nnx.bridge import variables as bv
 from flax.nnx.bridge import module as bdg_module
@@ -40,7 +40,7 @@
 @dataclasses.dataclass
 class Functional(tp.Generic[M]):
   module_type: tp.Type[M]
-  graphdef: tp.Optional[graph.GraphDef[M]]
+  graphdef: tp.Optional[graphlib.GraphDef[M]]
   args: tuple[tp.Any, ...]
   kwargs: dict[str, tp.Any]
 
@@ -66,7 +66,7 @@ def _functional_constructor(*args: tp.Any, **kwargs: tp.Any) -> Functional[M]:
 
 
 def _set_initializing(module: Module, initializing: bool):
-  for _, value in graph.iter_graph(module, graph=True):
+  for _, value in graphlib.iter_graph(module, graph=True):
     if isinstance(value, Pytree):
       value._pytree__state._initializing = initializing