Add ReLU

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
 

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
 

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",

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

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"]

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

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)

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 ───────────────────────────────────────────────────────────────────