Merge pull request #3988 from google:nnx-lora-param

PiperOrigin-RevId: 648503216
google · Jul 1, 2024 · b5503da · b5503da
2 parents afaa721 + 24c44f6
commit b5503da
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/flax/nnx/nnx/nn/lora.py b/flax/nnx/nnx/nn/lora.py
@@ -46,10 +46,11 @@
 default_kernel_init = initializers.lecun_normal()
 
 
-class LoRAParam(variables.Variable[A]):
+class LoRAParam(variables.Param[A]):
   pass
 
 
+
 class LoRA(Module):
   """A standalone LoRA layer.
 

diff --git a/flax/nnx/tests/nn/lora_test.py b/flax/nnx/tests/nn/lora_test.py
@@ -105,7 +105,7 @@ def __call__(self, x):
   def test_lora_param_type(self):
     rngs = nnx.Rngs(0)
     model = nnx.LoRA(3, 4, 2, lora_param_type=nnx.LoRAParam, rngs=rngs)
-    _, params, lora_params = nnx.split(model, nnx.Param, nnx.LoRAParam)
+    _, lora_params, params = nnx.split(model, nnx.LoRAParam, nnx.Param)
     assert params == {}
     assert ('lora_a' in lora_params) and ('lora_b' in lora_params)
     np.testing.assert_allclose(lora_params.lora_a.value, model.lora_a.value)