import numpy as np
from .core import Tensor

class Module:
    def parameters(self):
        return [getattr(self, a) for a in dir(self) if isinstance(getattr(self, a), Tensor)]

class Linear(Module):
    def __init__(self, in_dim, out_dim):
        # Инициализация весов
        scale = np.sqrt(2.0 / in_dim)
        self.w = Tensor(np.random.randn(in_dim, out_dim) * scale)
        self.b = Tensor(np.zeros(out_dim))

    def __call__(self, x):
        return x @ self.w + self.b

def relu(x):
    out = Tensor(np.maximum(0, x.data))
    out._prev = {x}
    def _backward():
        x.grad = (x.grad if x.grad is not None else 0) + (x.data > 0) * out.grad
    out._backward = _backward
    return out