Repository: google/flax

Issue:
nnx.nn.LoRA does not respect dtype argument
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
Reproduced on Google Colab (CPU runtime)

flax: 0.10.5
jax: 0.5.2
jaxlib: 0.5.1
python: 3.11.11 (main, Dec  4 2024, 08:55:07) [GCC 11.4.0]
platform: Linux-6.1.85+-x86_64-with-glibc2.35

### Problem you have encountered:

Unlike nnx.nn.Linear, nnx.nn.LoRA does perform the computation in the passed in dtype.

```
import platform
import sys

import flax
import jax
import jax.numpy as jnp
import jaxlib
from flax.nnx import Linear, LoRA, Rngs

print(f"flax: {flax.__version__}")
print(f"jax: {jax.__version__}")
print(f"jaxlib: {jaxlib.__version__}")

print(f"python: {sys.version}")
print(f"platform: {platform.platform()}")

linear = Linear(
    in_features=1,
    out_features=1,
    dtype=jnp.float16,
    param_dtype=jnp.float16,
    rngs=Rngs(0),
)
lora = LoRA(
    in_features=1,
    lora_rank=1,
    out_features=1,
    dtype=jnp.float16,
    param_dtype=jnp.float16,
    rngs=Rngs(0),
)

one = jnp.ones((1), dtype=jnp.float32)
print(linear(one).dtype)
print(lora(one).dtype)
```

The output is:

```
float16
float32
```

### What you expected to happen:

Both outputs should be float16. Currently, the dtype argument is ignored.

### Steps to reproduce:

https://colab.research.google.com/drive/1CBYAGHL_A3HevKmJHhBSJ88yp4Gipi69?usp=sharing

Discussion:

Fix (patch):
diff --git a/flax/nnx/nn/lora.py b/flax/nnx/nn/lora.py
index 76725f0a..36f19d8b 100644
--- a/flax/nnx/nn/lora.py
+++ b/flax/nnx/nn/lora.py
@@ -19,6 +19,7 @@ from flax.nnx import rnglib, variablelib
 from flax.nnx.module import Module
 from flax.nnx.nn import initializers
 from flax.nnx.nn.linear import Linear
+from flax.nnx.nn.dtypes import promote_dtype
 from flax.typing import Dtype, Initializer
 import jax
 import jax.numpy as jnp
@@ -105,7 +106,10 @@ class LoRA(Module):
     )
 
   def __call__(self, x: jax.Array):
-    out = x @ self.lora_a @ self.lora_b
+    x, lora_a, lora_b = promote_dtype(
+      (x, self.lora_a.value, self.lora_b.value), dtype=self.dtype
+    )
+    out = x @ lora_a @ lora_b
     if self.base_module is not None:
       if not callable(self.base_module):
         raise ValueError('`self.base_module` must be callable.')


Test changes:
diff --git a/tests/nnx/nn/lora_test.py b/tests/nnx/nn/lora_test.py
index db501e5e..1e89ca58 100644
--- a/tests/nnx/nn/lora_test.py
+++ b/tests/nnx/nn/lora_test.py
@@ -17,6 +17,7 @@ from absl.testing import absltest
 import numpy as np
 
 from flax import nnx
+from jax import numpy as jnp
 
 
 class TestLora(absltest.TestCase):
@@ -115,6 +116,14 @@ class TestLora(absltest.TestCase):
     np.testing.assert_allclose(params['lora_a'].value, model.lora_a.value)
     assert lora_params == {}
 
+  def test_dtype(self):
+    rngs = nnx.Rngs(0)
+    model = nnx.LoRA(3, 4, 2, dtype=jnp.float16, param_dtype=jnp.float32,
+                     rngs=rngs)
+    assert model.lora_a.value.dtype == jnp.float32
+    y = model(jnp.ones((1, 3)).astype(jnp.float32))
+    assert y.dtype == jnp.float16
+
 
 if __name__ == '__main__':
   absltest.main()