Merge pull request #4912 from google:no-copy-merge

PiperOrigin-RevId: 800115486

Author: Flax Authors

flax/nnx/__init__.py

@@ -208,7 +208,4 @@ def __getattr__(name):
           DeprecationWarning,
           stacklevel=2,
       )
-    if name not in globals():
-      raise AttributeError(f"Module {__name__} has no attribute '{name}'")
-
-    return globals()[name]
+    raise AttributeError(f"Module {__name__} has no attribute '{name}'")

flax/nnx/graph.py

@@ -1072,7 +1072,7 @@ def unflatten(  # type: ignore[invalid-annotation]
   *,
   index_ref: IndexMap | None = None,
   outer_index_outer_ref: IndexMap | None = None,
-  copy_variables: bool = True,
+  copy_variables: bool = False,
 ) -> Node:
   """Unflattens a graphdef into a node with the given state.
 
@@ -1835,6 +1835,7 @@ def merge(  # type: ignore[invalid-annotation]
       _state,
       index_ref=self.index_ref,
       outer_index_outer_ref=outer_index_outer_ref,
+      copy_variables=True,
     )
     return node
 
@@ -2307,6 +2308,7 @@ def merge(  # type: ignore[invalid-annotation]
   state: tp.Any,
   /,
   *states: tp.Any,
+  copy: bool = False,
 ) -> A:
   """The inverse of :func:`flax.nnx.split`.
 
@@ -2348,6 +2350,7 @@ def merge(  # type: ignore[invalid-annotation]
     graphdef: A :class:`flax.nnx.GraphDef` object.
     state: A :class:`flax.nnx.State` object.
     *states: Additional :class:`flax.nnx.State` objects.
+    copy: Whether to create new copies of the Variables in the states, defaults to ``False``.
   Returns:
     The merged :class:`flax.nnx.Module`.
   """
@@ -2357,7 +2360,7 @@ def merge(  # type: ignore[invalid-annotation]
     _state = state
   else:
     _state = _merge_to_flat_state((state, *states))
-  node = unflatten(graphdef, _state)
+  node = unflatten(graphdef, _state, copy_variables=copy)
   return node
 
 
@@ -2592,7 +2595,7 @@ def clone(node: Node) -> Node:
     A deep copy of the :class:`Module` object.
   """
   graphdef, state = split(node)
-  return merge(graphdef, state)
+  return merge(graphdef, state, copy=True)
 
 
 def _mutable_like(path, x):

tests/nnx/graph_utils_test.py

@@ -292,7 +292,7 @@ def __init__(self):
 
     assert isinstance(m2.a, nnx.Param)
     assert isinstance(state['a'], nnx.Param)
-    assert m2.a is not state['a']
+    assert m2.a is state['a']
     assert m2.a.value == state['a'].value
 
   def test_shared_state_variables_shared_with_graph(self):
@@ -319,8 +319,8 @@ def __init__(self):
     assert isinstance(m2.a, nnx.Param)
     assert isinstance(m2.b, nnx.Param)
     assert isinstance(state['a'], nnx.Param)
-    assert m2.a is not state['a']
-    assert m2.b is not state['a']
+    assert m2.a is state['a']
+    assert m2.b is state['a']
     assert m2.a.value == state['a'].value
     assert m2.b.value == state['a'].value
     assert m2.a is m2.b
@@ -367,7 +367,7 @@ def __init__(self):
     assert isinstance(m2.tree, Tree)
     assert m2.tree.a.raw_value == 1
     assert m2.tree.b == 'a'
-    assert m2.tree.a is not m.tree.a
+    assert m2.tree.a is m.tree.a
     assert m2.tree is not m.tree
 
   def test_cached_unflatten(self):

tests/nnx/integration_test.py

@@ -204,9 +204,9 @@ def __call__(self, x):
     @jax.jit
     def train_step(params, counts, x, y):
       def loss_fn(params):
-        y_pred, (_, updates) = graphdef.apply(params, counts)(x)
-        loss = jax.numpy.mean((y_pred - y) ** 2)
-        return loss, nnx.filter_state(updates, Count)
+        model = nnx.merge(graphdef, params, counts, copy=True)
+        loss = jax.numpy.mean((model(x) - y) ** 2)
+        return loss, nnx.state(model, Count)
 
       # compute gradient
       grads, counts = jax.grad(loss_fn, has_aux=True)(params)

tests/nnx/mutable_array_test.py

@@ -200,7 +200,7 @@ def __init__(self):
     self.assertTrue(m3.a.has_ref)
     self.assertTrue(m3.b.has_ref)
     self.assertIsNot(m2, m3)
-    self.assertIsNot(m.a, m3.a)
+    self.assertIs(m.a, m3.a)
 
   def test_freeze_duplicate_error(self):
     class Foo(nnx.Module):