feat: Add lora_alpha tests to Dense, Embedding, and EinsumDense layers

keras/src/layers/convolutional/conv_test.py

@@ -735,6 +735,51 @@ def call(self, x):
             model.conv2d.lora_kernel_a.path, "mymodel/conv2d/lora_kernel_a"
         )
 
+    @pytest.mark.requires_trainable_backend
+    def test_enable_lora_with_alpha(self):
+        # Create a Conv2D layer with a small kernel for simplicity.
+        layer = layers.Conv2D(filters=3, kernel_size=(2, 2), padding="valid")
+        # Use a fixed input shape: batch size 1, height=4, width=4, channels=3.
+        input_shape = (1, 4, 4, 3)
+        layer.build(input_shape)
+
+        # Set the base kernel to known, deterministic values.
+        base_kernel = np.linspace(
+            0, 1, num=np.prod(layer.kernel.shape), dtype=np.float32
+        )
+        base_kernel = base_kernel.reshape(layer.kernel.shape)
+        layer.kernel.assign(base_kernel)
+
+        # Enable LoRA with rank=2 and a custom lora_alpha value (e.g. 3.0).
+        layer.enable_lora(rank=2, lora_alpha=3.0)
+        self.assertEqual(layer.lora_rank, 2)
+        self.assertEqual(layer.lora_alpha, 3.0)
+
+        # For Conv2D, assume the LoRA weights have shapes:
+        #   lora_kernel_a: (kernel_height, kernel_width, in_channels, rank)
+        #   lora_kernel_b: (rank, out_channels)
+        lora_a_shape = layer.lora_kernel_a.shape
+        lora_b_shape = layer.lora_kernel_b.shape
+
+        # Assign known constant values to LoRA weights.
+        lora_a = np.full(lora_a_shape, 0.1, dtype=np.float32)
+        lora_b = np.full(lora_b_shape, 0.2, dtype=np.float32)
+        layer.lora_kernel_a.assign(lora_a)
+        layer.lora_kernel_b.assign(lora_b)
+
+        # Compute the expected delta.
+        # Flatten lora_kernel_a to shape (-1, rank),
+        # multiply with lora_kernel_b,
+        # then reshape to the kernel's shape.
+        scaling = 3.0 / 2  # lora_alpha / lora_rank
+        delta = np.matmul(lora_a.reshape(-1, 2), lora_b)
+        delta = delta.reshape(base_kernel.shape)
+        expected_effective_kernel = base_kernel + scaling * delta
+
+        # Compare the effective kernel computed via the property.
+        actual_effective_kernel = layer.kernel.numpy()
+        self.assertAllClose(actual_effective_kernel, expected_effective_kernel)
+
     @pytest.mark.requires_trainable_backend
     def test_lora_rank_argument(self):
         self.run_layer_test(

keras/src/layers/core/dense_test.py

@@ -271,6 +271,29 @@ def test_enable_lora(self):
         model.load_weights(temp_filepath)
         self.assertAllClose(model.predict(x), new_model.predict(x))
 
+    @pytest.mark.requires_trainable_backend
+    def test_enable_lora_with_alpha(self):
+        # Create a Dense layer and build it.
+        layer = layers.Dense(units=8)
+        layer.build((None, 4))
+
+        # Enable LoRA with rank=2 and lora_alpha=3.0.
+        layer.enable_lora(2, lora_alpha=3.0)
+        self.assertEqual(layer.lora_rank, 2)
+        self.assertEqual(layer.lora_alpha, 3.0)
+
+        # Manually compute the expected effective kernel:
+        # effective_kernel = base_kernel +
+        # (lora_alpha / lora_rank) * (lora_kernel_a @ lora_kernel_b)
+        base_kernel = layer._kernel.numpy()
+        lora_update = np.matmul(
+            layer.lora_kernel_a.numpy(), layer.lora_kernel_b.numpy()
+        )
+        effective_kernel_expected = base_kernel + (3.0 / 2) * lora_update
+
+        # Verify that the effective kernel matches expectation.
+        self.assertAllClose(layer.kernel.numpy(), effective_kernel_expected)
+
     @pytest.mark.requires_trainable_backend
     def test_lora_weight_name(self):
         class MyModel(models.Model):

keras/src/layers/core/einsum_dense_test.py

@@ -361,6 +361,48 @@ def test_enable_lora(self):
         model.load_weights(temp_filepath)
         self.assertAllClose(model.predict(x), new_model.predict(x))
 
+    @pytest.mark.requires_trainable_backend
+    def test_enable_lora_with_alpha(self):
+        # Use a simple equation that mimics a Dense layer behavior.
+        equation = "ab,bc->ac"
+        output_shape = 3  # This means the kernel shape will be (input_dim, 3)
+        bias_axes = None
+
+        # Create and build the EinsumDense layer with an input shape (None, 2)
+        layer = layers.EinsumDense(
+            equation=equation, output_shape=output_shape, bias_axes=bias_axes
+        )
+        # Build the layer with an input shape of (batch, 2)
+        layer.build((None, 2))
+
+        # Set the base kernel weights to a known value.
+        base_kernel = np.array(
+            [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32
+        )
+        layer._kernel.assign(base_kernel)
+
+        # Enable LoRA with rank=2 and a custom lora_alpha=3.0.
+        layer.enable_lora(rank=2, lora_alpha=3.0)
+        self.assertEqual(layer.lora_rank, 2)
+        self.assertEqual(layer.lora_alpha, 3.0)
+
+        # The expected shapes are:
+        #   base_kernel: (2, 3)
+        #   lora_kernel_a: (2, 2) and lora_kernel_b: (2, 3)
+        a_val = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
+        b_val = np.array([[0.5, 0.6, 0.7], [0.8, 0.9, 1.0]], dtype=np.float32)
+        layer.lora_kernel_a.assign(a_val)
+        layer.lora_kernel_b.assign(b_val)
+
+        # Compute expected effective kernel.
+        # Scaling factor is lora_alpha / lora_rank = 3.0 / 2 = 1.5
+        expected_delta = 1.5 * np.matmul(a_val, b_val)
+        expected_kernel = base_kernel + expected_delta
+
+        # Verify that the effective kernel property returns the expected value.
+        actual_kernel = layer.kernel.numpy()
+        self.assertAllClose(actual_kernel, expected_kernel)
+
     @pytest.mark.requires_trainable_backend
     def test_lora_rank_argument(self):
         self.run_layer_test(

keras/src/layers/core/embedding_test.py

@@ -181,6 +181,38 @@ def test_enable_lora(self):
         model.load_weights(temp_filepath)
         self.assertAllClose(model.predict(x), new_model.predict(x))
 
+    @pytest.mark.requires_trainable_backend
+    def test_enable_lora_with_alpha(self):
+        # Create an Embedding layer without specifying lora_rank
+        layer = layers.Embedding(input_dim=3, output_dim=2)
+        layer.build((None,))  # Build the layer
+
+        # Set the base embeddings to known values.
+        base_emb = np.array(
+            [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]], dtype=np.float32
+        )
+        layer.embeddings.assign(base_emb)
+
+        # Enable LoRA with a custom alpha: rank=2, lora_alpha=3.0.
+        layer.enable_lora(2, lora_alpha=3.0)
+        self.assertEqual(layer.lora_rank, 2)
+        self.assertEqual(layer.lora_alpha, 3.0)
+
+        # Manually assign known values to lora weights.
+        a_val = np.array([[0.1, 0.1], [0.2, 0.2], [0.3, 0.3]], dtype=np.float32)
+        b_val = np.array([[0.5, 0.5], [0.6, 0.6]], dtype=np.float32)
+        layer.lora_embeddings_a.assign(a_val)
+        layer.lora_embeddings_b.assign(b_val)
+
+        # Compute the expected delta.
+        # Scaling factor: (3.0 / 2) = 1.5
+        effective_delta = 1.5 * np.matmul(a_val, b_val)
+        expected_embeddings = base_emb + effective_delta
+
+        # Verify that the effective embeddings match expectation.
+        actual_embeddings = layer.embeddings.numpy()
+        self.assertAllClose(actual_embeddings, expected_embeddings)
+
     @pytest.mark.requires_trainable_backend
     def test_lora_rank_argument(self):
         self.run_layer_test(