from .module import Module
from .parameter import Parameter
from .functional import randn, zeros

class Linear(Module):
    def __init__(self, d_in, d_out):
        self.weights = Parameter(randn((d_in, d_out)))
        self.bias = Parameter(zeros((d_out,)))
    
    def __call__(self, x):
        return self.forward(x)

    def forward(self, x):
        # y = x @ W + b
        # Autograd will automatically track this computation
        return x.matmul(self.weights) + self.bias