import numpy as np


class LoRALinear:
    def __init__(self, weight: np.ndarray, rank: int, alpha: float = 1.0, seed: int = 0):
        if weight.ndim != 2 or rank <= 0 or rank > min(weight.shape):
            raise ValueError("rank must fit a rank-2 weight matrix")
        rng = np.random.default_rng(seed)
        self.weight = weight.copy()
        self.a = rng.normal(0.0, 0.02, size=(rank, weight.shape[1]))
        self.b = np.zeros((weight.shape[0], rank))
        self.scale = alpha / rank

    @property
    def trainable_parameters(self) -> int:
        return self.a.size + self.b.size

    def forward(self, inputs: np.ndarray) -> np.ndarray:
        return inputs @ self.weight.T + self.scale * (inputs @ self.a.T @ self.b.T)