Add flag for variables in nnx.merge

Author: samanklesaria

flax/nnx/graphlib.py

@@ -1089,6 +1089,7 @@ def unflatten(  # type: ignore[invalid-annotation]
   index_ref: IndexMap | None = None,
   outer_index_outer_ref: IndexMap | None = None,
   copy_variables: bool = False,
+  auto_create_variables: bool = True,
 ) -> Node:
   """Unflattens a graphdef into a node with the given state.
 
@@ -1150,6 +1151,7 @@ def unflatten(  # type: ignore[invalid-annotation]
       index_ref,
       outer_index_outer_ref,
       copy_variables,
+      auto_create_variables
     )
 
     try:
@@ -1171,6 +1173,7 @@ def _graph_unflatten(
   index_ref: IndexMap,
   outer_index_outer_ref: IndexMap | None,
   copy_variables: bool,
+  auto_create_variables: bool
 ) -> Node:
   """Recursive helper for graph_unflatten.
 
@@ -1265,7 +1268,7 @@ def get_mutable_array(array_refdef: ArrayRefDef, leaf):
         variable.set_raw_value(value)
     else:  # variabledef.index not in index_ref_cache
       # variable reference does not exist outside, create a new one
-      if isinstance(value, Variable):
+      if isinstance(value, Variable) or not auto_create_variables:
         variable = value
       else:
         variable = variabledef.type.from_metadata(
@@ -1314,6 +1317,7 @@ def _get_children() -> list[tuple[Key, tp.Any]]:
             index_ref,
             outer_index_outer_ref,
             copy_variables,
+            auto_create_variables
           )
         else:
           raise RuntimeError(f'Unknown node definition: {node_def!r}')
@@ -2359,6 +2363,7 @@ def merge(  # type: ignore[invalid-annotation]
   /,
   *states: tp.Any,
   copy: bool = False,
+  auto_create_variables: bool = True,
 ) -> A:
   """The inverse of :func:`flax.nnx.split`.
 
@@ -2410,7 +2415,7 @@ def merge(  # type: ignore[invalid-annotation]
     _state = state
   else:
     _state = _merge_to_flat_state((state, *states))
-  node = unflatten(graphdef, _state, copy_variables=copy)
+  node = unflatten(graphdef, _state, copy_variables=copy, auto_create_variables=auto_create_variables)
   return node
 
 
@@ -2534,6 +2539,7 @@ def map(
   /,
   *,
   graph: bool | None = None,
+  auto_create_variables: bool = True,
 ) -> A:
   """Map a function over the state of a graph node.
 
@@ -2567,7 +2573,7 @@ def map(
   """
   graphdef, state = split(node, graph=graph)
   state = statelib.map_state(f, state)
-  return merge(graphdef, state)
+  return merge(graphdef, state, auto_create_variables=auto_create_variables)
 
 
 def graphdef(

flax/nnx/pytreelib.py

@@ -1061,4 +1061,4 @@ def _maybe_int(x):
     return x
 
 def _get_str(x):
-  return x if isinstance(x, str) else str(x)
+  return x if isinstance(x, str) else str(x)

tests/nnx/graph_utils_test.py

@@ -1699,6 +1699,14 @@ def test_map_replace(self):
     np.testing.assert_array_equal(new_model.kernel[...], jnp.zeros((2, 3)))
     np.testing.assert_array_equal(new_model.bias[...], jnp.zeros((3,)))
 
+  def test_map_auto_create_variables_false(self):
+    rngs = nnx.Rngs(0)
+    new_rngs = nnx.map(
+      lambda path, x: 0, rngs, auto_create_variables=False
+    )
+    self.assertNotIsInstance(new_rngs.default.count, nnx.Variable)
+    self.assertEqual(new_rngs.default.count, 0)
+
 
 if __name__ == '__main__':
   absltest.main()