import unittest
import numpy as np
from lora import LoRALinear


class LoRATests(unittest.TestCase):
    def test_zero_initialized_b_preserves_base_output(self):
        weight = np.arange(12, dtype=float).reshape(3, 4)
        inputs = np.ones((2, 4))
        layer = LoRALinear(weight, rank=2)
        np.testing.assert_allclose(layer.forward(inputs), inputs @ weight.T)

    def test_trainable_parameter_count(self):
        layer = LoRALinear(np.zeros((4096, 4096)), rank=8)
        self.assertEqual(layer.trainable_parameters, 8 * (4096 + 4096))


if __name__ == "__main__":
    unittest.main()