diff --git a/AGENTS.md b/AGENTS.md index 7d259d9..ee5c9eb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -15,9 +15,6 @@ inspired by PyTorch. It currently provides: - Tests under `tests/`, grouped roughly by subsystem. - Example notebooks under `examples/`. -The project is still evolving, so prefer focused, understandable changes over -large framework-style abstractions. - ## Environment and Tooling This project uses `uv`. @@ -39,28 +36,9 @@ Notes: ## Testing Expectations -- Add or update pytest tests for behavioral changes. +- Only use the public API in tests if possible; do not test internals unless necessary - Prefer `np.testing.assert_array_equal` or `np.testing.assert_allclose` for tensor data comparisons. -- Run the narrowest relevant pytest target first, then broaden if the change - touches shared behavior. - -## Code Style and Design Preferences - -- Follow the existing PyTorch-inspired public API where it is already - established. -- Keep implementation simple and readable; this is an educational library. -- Prefer NumPy operations and structured Tensor/autograd helpers over ad hoc - special cases. -- Preserve existing module boundaries: - - Core tensor behavior belongs in `src/motorch/tensor.py`. - - Gradient graph construction and propagation belongs in - `src/motorch/autograd/`. - - Layers, activations, losses, and parameters belong in `src/motorch/nn/`. - - Optimizers belong in `src/motorch/optim/`. -- Avoid broad refactors unless they are directly needed for the task. -- Keep comments useful and sparse; explain non-obvious math or graph behavior, - not routine assignments. ## Repository Hygiene diff --git a/README.md b/README.md index eb6917f..568bab9 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ x = mo.tensor([[0.1, 0.2], [0.3, 0.4]], requires_grad=True) y = mo.tensor([[1.0], [-1.0]]) model = nn.Linear(in_features=2, out_features=1) -activation = nn.AltSigmoid() +activation = nn.ReLU() criterion = nn.LogisticLoss() optimizer = optim.SGD(model.parameters(), lr=0.01) @@ -57,7 +57,7 @@ class AltSigmoidMLP(nn.Module): def __init__(self): super().__init__() self.layer1 = nn.Linear(2, 3) - self.act1 = nn.AltSigmoid() + self.act1 = nn.ReLU() self.output = nn.Linear(3, 1) def forward(self, x): @@ -71,7 +71,7 @@ model = AltSigmoidMLP() Currently available modules (more to come): - Layers: nn.Linear -- Activations: nn.Sigmoid, nn.AltSigmoid, nn.Sgn +- Activations: nn.Sigmoid, nn.AltSigmoid, nn.Sgn, nn.ReLU - Loss: nn.LogisticLoss - Optimizers: optim.SGD diff --git a/src/motorch/nn/__init__.py b/src/motorch/nn/__init__.py index 853abff..e9159e1 100644 --- a/src/motorch/nn/__init__.py +++ b/src/motorch/nn/__init__.py @@ -6,13 +6,14 @@ from . import functional from . import init from .parameter import Parameter -from .modules import Linear, Module, LogisticLoss, Sgn, Sigmoid, AltSigmoid +from .modules import Linear, Module, LogisticLoss, Sgn, Sigmoid, AltSigmoid, ReLU __all__ = [ "Linear", "Module", "LogisticLoss", "Sgn", + "ReLU", "Sigmoid", "AltSigmoid", "functional", diff --git a/src/motorch/nn/functional.py b/src/motorch/nn/functional.py index 10eed7a..20a66f7 100644 --- a/src/motorch/nn/functional.py +++ b/src/motorch/nn/functional.py @@ -10,23 +10,31 @@ # --- Activations --- # -def sigmoid(x): +def sigmoid(z): """Compute the sigmoid activation for a tensor input.""" - x_clipped = mo.clip(x, -700, 700) # to avoid numerical instability - return 1 / (1 + mo.exp(-x_clipped)) + z_clipped = mo.clip(z, -700, 700) # to avoid numerical instability + return 1 / (1 + mo.exp(-z_clipped)) -def sigmoid_grad(x, precomputed=False): +def sigmoid_grad(z, precomputed=False): """ Compute the gradient of the sigmoid activation. args: - - x (Tensor): the input data - - precomputed (boolean): whether x is the output of sigmoid + - z (Tensor): the input data + - precomputed (boolean): whether z is the output of sigmoid """ if not precomputed: - x = sigmoid(x) - return x * (1 - x) + z = sigmoid(z) + return z * (1 - z) + + +def relu(z): + return mo.where(z > 0, z, 0) + + +def relu_grad(z): + return mo.where(z > 0, 1, 0) # --- Loss Functions --- # TODO: Change name to explicit +1/-1 loss. @@ -45,19 +53,19 @@ def logloss_grad(logits, labels): # --- Layers --- # -def linear(x, weight, bias): +def linear(z, weight, bias): """Compute a linear transformation with optional bias. Parameters: - x: Input tensor with last dimension matching ``weight.shape[0]``. + z: Input tensor with last dimension matching ``weight.shape[0]``. weight: Weight tensor of shape ``(in_features, out_features)``. bias: Bias tensor added to each row of the result. """ - if not isinstance(x, Tensor): + if not isinstance(z, Tensor): raise ValueError("Input must be of type motorch.Tensor") - assert x.shape[-1] == weight.shape[0], ( - f"Expected x.shape = ({len(x)}, {weight.shape[0]}), got {x.shape}." + assert z.shape[-1] == weight.shape[0], ( + f"Expected z.shape = ({len(z)}, {weight.shape[0]}), got {z.shape}." ) - return x @ weight + bias + return z @ weight + bias diff --git a/src/motorch/nn/modules/__init__.py b/src/motorch/nn/modules/__init__.py index dfc080d..88015a9 100644 --- a/src/motorch/nn/modules/__init__.py +++ b/src/motorch/nn/modules/__init__.py @@ -6,6 +6,6 @@ from .module import Module from .linear import Linear from .loss import LogisticLoss -from .activations import Sgn, Sigmoid, AltSigmoid +from .activations import Sgn, Sigmoid, AltSigmoid, ReLU -__all__ = ["Module", "Linear", "LogisticLoss", "Sgn", "Sigmoid", "AltSigmoid"] +__all__ = ["Module", "Linear", "LogisticLoss", "Sgn", "Sigmoid", "AltSigmoid", "ReLU"] diff --git a/src/motorch/nn/modules/activations.py b/src/motorch/nn/modules/activations.py index f403d6e..933ecf0 100644 --- a/src/motorch/nn/modules/activations.py +++ b/src/motorch/nn/modules/activations.py @@ -11,26 +11,26 @@ class Sgn(Module): """Sign activation module that returns +1 for non-negative inputs.""" - def forward(self, x): - return mo.where(x >= 0, 1, -1) + def forward(self, z): + return mo.where(z >= 0, 1, -1) class Sigmoid(Module): """Sigmoid activation module with stored gradient support.""" - def forward(self, x): + def forward(self, z): """Compute the sigmoid activation and cache its gradient.""" with no_grad(): - out = F.sigmoid(x) + out = F.sigmoid(z) result = tensor(out.data) local_grad = [self._grad_fn(result)] - apply_forward_pass(result, [x], local_grad) + apply_forward_pass(result, [z], local_grad) return result def _grad_fn(self, z): - """Compute the gradient of the sigmoid activation w.r.t the input x.""" + """Compute the gradient of the sigmoid activation w.r.t the input z.""" with no_grad(): out = F.sigmoid_grad(z, precomputed=True) return out @@ -39,17 +39,36 @@ def _grad_fn(self, z): class AltSigmoid(Module): """Sigmoid function rescaled to output values between -1 and 1.""" - def forward(self, x): + def forward(self, z): with no_grad(): - out = 2 * F.sigmoid(x) - 1 + out = 2 * F.sigmoid(z) - 1 result = tensor(out.data) - local_grad = [self._grad_fn(x)] - apply_forward_pass(result, [x], local_grad) + local_grad = [self._grad_fn(z)] + apply_forward_pass(result, [z], local_grad) return result - def _grad_fn(self, x): + def _grad_fn(self, z): + with no_grad(): + out = 2 * F.sigmoid_grad(z) + return out + + +class ReLU(Module): + """ReLU activation function.""" + + def forward(self, z): + with no_grad(): + out = F.relu(z) + + result = tensor(out.data) + local_grad = [self._grad_fn(result)] + apply_forward_pass(result, [z], local_grad) + + return result + + def _grad_fn(self, z): with no_grad(): - out = 2 * F.sigmoid_grad(x) + out = F.relu_grad(z) return out diff --git a/tests/nn/modules/test_activations.py b/tests/nn/modules/test_activations.py index 4ca00a4..f0e3649 100644 --- a/tests/nn/modules/test_activations.py +++ b/tests/nn/modules/test_activations.py @@ -2,6 +2,7 @@ import numpy as np import pytest +import motorch.nn as nn from motorch.tensor import tensor, Tensor from motorch.nn.modules.activations import Sgn, Sigmoid, AltSigmoid @@ -270,3 +271,65 @@ def test_grad_does_not_accumulate_across_calls(self): self.altsig(x).backward() grad_after_second = float(x.grad.item()) if x.grad is not None else 0.0 assert grad_after_first == pytest.approx(grad_after_second, rel=1e-6) + + +# ── ReLU ────────────────────────────────────────────────────────────────────── + + +class TestReLU: + def setup_method(self): + self.relu = nn.ReLU() + + def test_zero_input(self): + assert_close(self.relu(tensor(0.0)), 0.0) + + def test_positive_input(self): + assert_close(self.relu(tensor(2.5)), 2.5) + + def test_negative_input(self): + assert_close(self.relu(tensor(-2.5)), 0.0) + + def test_vector_input(self): + out = self.relu(tensor([-2.0, -0.5, 0.0, 0.5, 2.0])) + assert_close(out, [0.0, 0.0, 0.0, 0.5, 2.0]) + + def test_matrix_input(self): + out = self.relu(tensor([[-1.0, 2.0], [0.0, -3.0]])) + assert_close(out, [[0.0, 2.0], [0.0, 0.0]]) + + def test_returns_tensor(self): + assert isinstance(self.relu(tensor(1.0)), Tensor) + + def test_grad_positive_input(self): + z = tensor(2.5, requires_grad=True) + self.relu(z).backward() + assert_close(z.grad, 1.0) + + def test_grad_negative_input(self): + z = tensor(-2.5, requires_grad=True) + self.relu(z).backward() + assert_close(z.grad, 0.0) + + def test_grad_zero_input(self): + z = tensor(0.0, requires_grad=True) + self.relu(z).backward() + assert_close(z.grad, 0.0) + + def test_grad_vector_input(self): + z = tensor([-2.0, -0.5, 0.0, 0.5, 2.0], requires_grad=True) + self.relu(z).backward() + assert_close(z.grad, [0.0, 0.0, 0.0, 1.0, 1.0]) + + def test_grad_matrix_input(self): + z = tensor([[-1.0, 2.0], [0.0, -3.0]], requires_grad=True) + self.relu(z).backward() + assert_close(z.grad, [[0.0, 1.0], [0.0, 0.0]]) + + def test_grad_does_not_accumulate_across_calls(self): + z = tensor(2.5, requires_grad=True) + self.relu(z).backward() + grad_after_first = float(z.grad.item()) if z.grad is not None else 0.0 + z.grad = None + self.relu(z).backward() + grad_after_second = float(z.grad.item()) if z.grad is not None else 0.0 + assert grad_after_first == pytest.approx(grad_after_second, rel=1e-6) diff --git a/tests/nn/test_functional.py b/tests/nn/test_functional.py index 010a1a3..d14d928 100644 --- a/tests/nn/test_functional.py +++ b/tests/nn/test_functional.py @@ -114,6 +114,52 @@ def test_returns_tensor(self): assert isinstance(F.sigmoid_grad(tensor(0.0)), Tensor) +# ── relu ────────────────────────────────────────────────────────────────────── + + +class TestRelu: + def test_zero_input(self): + assert_close(F.relu(tensor(0.0)), 0.0) + + def test_positive_input(self): + assert_close(F.relu(tensor(2.5)), 2.5) + + def test_negative_input(self): + assert_close(F.relu(tensor(-2.5)), 0.0) + + def test_vector_input(self): + out = F.relu(tensor([-2.0, -0.5, 0.0, 0.5, 2.0])) + assert_close(out, [0.0, 0.0, 0.0, 0.5, 2.0]) + + def test_matrix_input(self): + out = F.relu(tensor([[-1.0, 2.0], [0.0, -3.0]])) + assert_close(out, [[0.0, 2.0], [0.0, 0.0]]) + + def test_returns_tensor(self): + assert isinstance(F.relu(tensor(1.0)), Tensor) + + +# ── relu_grad ───────────────────────────────────────────────────────────────── + + +class TestReluGrad: + def test_zero_input(self): + assert_close(F.relu_grad(tensor(0.0)), 0.0) + + def test_positive_input(self): + assert_close(F.relu_grad(tensor(2.5)), 1.0) + + def test_negative_input(self): + assert_close(F.relu_grad(tensor(-2.5)), 0.0) + + def test_vector_input(self): + out = F.relu_grad(tensor([-2.0, -0.5, 0.0, 0.5, 2.0])) + assert_close(out, [0.0, 0.0, 0.0, 1.0, 1.0]) + + def test_returns_tensor(self): + assert isinstance(F.relu_grad(tensor(1.0)), Tensor) + + # ── logloss ───────────────────────────────────────────────────────────────────