Add _graph_node_set_key method for List class

Author: samanklesaria

flax/nnx/helpers.py

@@ -27,6 +27,7 @@
 from flax.nnx.rnglib import Rngs
 from flax.nnx.statelib import State
 from flax.training.train_state import struct
+from flax.nnx.variablelib import Variable
 
 A = tp.TypeVar('A')
 M = tp.TypeVar('M', bound=Module)
@@ -185,6 +186,17 @@ def __setitem__(self, index: int | slice, value: A | tp.Iterable[A]) -> None:
     else:
       raise TypeError('Invalid index type')
 
+  def _graph_node_set_key(self, key: str, value: tp.Any):
+    if not isinstance(key, int):
+      raise KeyError(f'Invalid key: {key}')
+    elif key < len(self):
+      if isinstance(variable := self[key], Variable) and isinstance(value, Variable):
+        variable.update_from_state(value)
+      else:
+        self[key] = value
+    else:
+      self.insert(key, value)
+
   def __delitem__(self, index: int | slice) -> None:
     if isinstance(index, int):
       if index < 0:

tests/nnx/graph_utils_test.py

@@ -1180,6 +1180,17 @@ def swap(path, node):
     self.assertEqual(bar2[1].d, -20)
     self.assertEqual(n, 2)
 
+  def test_recursive_map_with_list(self):
+    rngs = nnx.Rngs(0)
+    model = nnx.Sequential(nnx.Linear(2, 3, rngs=rngs), nnx.relu, nnx.Linear(3, 4, rngs=rngs))
+
+    def add_rank2_lora(_, node):
+      if isinstance(node, nnx.Linear):
+        return nnx.LoRA(node.in_features, 2, node.out_features, base_module=node, rngs=rngs)
+      return node
+
+    self.assertEqual(len(nnx.recursive_map(add_rank2_lora, model).layers), 3)
+
   def test_graphdef_hash_with_sequential(self):
     rngs = nnx.Rngs(0)
     net = nnx.Sequential(