Add SGD momentum & weight decay; Dropout; Tensor.clip/clamp & indexing

README.md

@@ -30,6 +30,11 @@ LeanPass is a lightweight, transparent NumPy-based autodiff library for small ne
 - Linear layers and multilayer perceptrons
 - Loss utilities: `mse_loss`, `cross_entropy_loss`, and `binary_cross_entropy_loss`
 - SGD and Adam optimizers
+- SGD now supports momentum and L2 weight decay: `optim.SGD(params, lr=..., momentum=0.9, weight_decay=1e-4)`
+- Adam supports an optional L2 weight decay: `optim.Adam(params, lr=..., weight_decay=1e-4)`
+- Dropout layer: `nn.Dropout(p=0.5)` for inverted dropout during training
+- Tensor utilities: `Tensor.clip(a_min, a_max)` / `Tensor.clamp(min, max)` for value clamping
+- Basic indexing: `Tensor.__getitem__` supports slicing/indexing that participates in autodiff
 - A simple demo script that trains a toy model
 
 ## Installation

leanpass/README.md

@@ -37,6 +37,8 @@ Supported operations:
 - `Tensor.gelu()`
 - `Tensor.softmax()`
 - `Tensor.backward()`
+ - `Tensor.clip(a_min, a_max)` / `Tensor.clamp(min, max)`
+ - Basic indexing via `Tensor[...]` which participates in autodiff
 
 ### Neural network layers
 
@@ -51,6 +53,7 @@ Available components:
 
 - `nn.Linear(in_features, out_features)` creates a linear layer with weights and bias.
 - `nn.MLP(layer_sizes)` creates a multilayer perceptron with ReLU activations between layers.
+ - `nn.Dropout(p=0.5)` creates an inverted-dropout layer for training-time regularization.
 - `nn.mse_loss(predictions, targets)` computes mean squared error.
 - `nn.cross_entropy_loss(predictions, targets)` computes categorical cross-entropy for multi-class targets.
 - `nn.binary_cross_entropy_loss(predictions, targets)` computes binary cross-entropy for binary classification.
@@ -65,10 +68,10 @@ optimizer = optim.SGD(model.parameters(), lr=0.01)
 optimizer = optim.Adam(model.parameters(), lr=0.001)
 ```
 
-Available optimizers:
+Available optimizers (options):
 
-- `optim.SGD(parameters, lr=...)` performs simple gradient descent.
-- `optim.Adam(parameters, lr=...)` performs Adam optimization with bias correction.
+- `optim.SGD(parameters, lr=..., momentum=0.0, weight_decay=0.0)` supports momentum and L2 weight decay.
+- `optim.Adam(parameters, lr=..., weight_decay=0.0)` supports an optional L2 weight decay term.
 
 ### Example
 

leanpass/nn.py

@@ -38,6 +38,7 @@ def forward(self, x: Tensor) -> Tensor:
 
     def __call__(self, x: Tensor) -> Tensor:
         return self.forward(x)
+
 
 
 class MLP(Module):
@@ -67,6 +68,31 @@ def parameters(self):
         return params
 
 
+class Dropout(Module):
+    """Simple dropout layer.
+
+    Usage: `Dropout(p=0.5)(x, training=True)` or call with `training=False` to
+    disable dropout at evaluation time.
+    """
+
+    def __init__(self, p: float = 0.5):
+        if not 0 <= p < 1:
+            raise ValueError("p must be in the interval [0, 1)")
+        self.p = float(p)
+
+    def forward(self, x: Tensor, training: bool = True) -> Tensor:
+        if not training or self.p == 0.0:
+            return x
+        # create a binary mask and scale to preserve expectation (inverted dropout)
+        mask = (np.random.rand(*x.data.shape) > self.p).astype(np.float64) / (1.0 - self.p)
+        return x * Tensor(mask)
+
+    def __call__(self, x: Tensor, training: bool = True) -> Tensor:
+        return self.forward(x, training=training)
+
+
+
+
 def mse_loss(prediction: Tensor, target: Tensor) -> Tensor:
     """Mean squared error loss used for regression training."""
     return ((prediction - target) ** 2).mean()

leanpass/optim.py

@@ -2,18 +2,40 @@
 
 
 class SGD:
-    """Simple stochastic gradient descent optimizer."""
+    """Stochastic gradient descent with optional momentum and L2 weight decay.
 
-    def __init__(self, params, lr=1e-2):
+    Args:
+        params: iterable of `Tensor` parameters
+        lr: learning rate
+        momentum: momentum factor (0.0 means no momentum)
+        weight_decay: L2 penalty factor applied to parameters (aka weight decay)
+    """
+
+    def __init__(self, params, lr=1e-2, momentum=0.0, weight_decay=0.0):
         self.params = list(params)
         self.lr = lr
+        self.momentum = float(momentum)
+        self.weight_decay = float(weight_decay)
+        # velocity buffers for momentum (kept even if momentum==0 for simplicity)
+        self.v = [np.zeros_like(p.data) for p in self.params]
 
     def step(self):
-        """Apply a plain gradient descent update to each parameter."""
-        for param in self.params:
+        """Perform a parameter update step."""
+        for i, param in enumerate(self.params):
             if param.grad is None:
                 continue
-            param.data = param.data - self.lr * param.grad
+            # apply L2 weight decay directly to the gradient (common choice)
+            g = param.grad
+            if self.weight_decay:
+                g = g + self.weight_decay * param.data
+
+            if self.momentum:
+                self.v[i] = self.momentum * self.v[i] + g
+                update = self.v[i]
+            else:
+                update = g
+
+            param.data = param.data - self.lr * update
 
     def zero_grad(self):
         """Zero out gradients so the next backward pass starts clean."""
@@ -22,13 +44,14 @@ def zero_grad(self):
 
 
 class Adam:
-    """Adam optimizer with bias-corrected moment estimates."""
+    """Adam optimizer with bias-corrected moment estimates and optional L2 weight decay."""
 
-    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
         self.params = list(params)
         self.lr = lr
         self.b1, self.b2 = betas
         self.eps = eps
+        self.weight_decay = float(weight_decay)
         self.m = [np.zeros_like(p.data) for p in self.params]
         self.v = [np.zeros_like(p.data) for p in self.params]
         self.t = 0
@@ -39,7 +62,11 @@ def step(self):
         for i, param in enumerate(self.params):
             if param.grad is None:
                 continue
+            # apply L2 weight decay to the gradient
             g = param.grad
+            if self.weight_decay:
+                g = g + self.weight_decay * param.data
+
             self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
             self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (g ** 2)
 

leanpass/tensor.py

@@ -238,6 +238,31 @@ def _backward():
         out._backward = _backward
         return out
 
+    def clip(self, a_min=None, a_max=None):
+        """Clamp tensor values to the interval [a_min, a_max].
+
+        Both bounds are optional. The backward pass only propagates gradient
+        for elements that were not clipped (standard straight-through behavior).
+        """
+        out_data = np.clip(self.data, a_min, a_max)
+        out = self._create_child(out_data, "clip", (self,), meta={"min": a_min, "max": a_max})
+
+        def _backward():
+            if self.requires_grad:
+                mask = np.ones_like(self.data, dtype=np.float64)
+                if a_min is not None:
+                    mask = mask * (self.data > a_min)
+                if a_max is not None:
+                    mask = mask * (self.data < a_max)
+                self.grad += _sum_to_shape(out.grad * mask, self.data.shape)
+
+        out._backward = _backward
+        return out
+
+    # alias common name
+    def clamp(self, min=None, max=None):
+        return self.clip(min, max)
+
     def sum(self, axis=None, keepdims=False):
         out_data = self.data.sum(axis=axis, keepdims=keepdims)
         out = Tensor(out_data, requires_grad=self.requires_grad, name="sum")
@@ -257,6 +282,26 @@ def _backward():
         out._backward = _backward
         return out
 
+    def __getitem__(self, idx):
+        """Basic indexing / slicing returning a new Tensor view (not a view in-place).
+
+        The returned tensor participates in autodiff; gradients are placed back
+        into the source tensor at the same indices during the backward pass.
+        """
+        out_data = self.data[idx]
+        out = self._create_child(out_data, "getitem", (self,), meta={"index": idx})
+
+        def _backward():
+            if self.requires_grad:
+                if out.grad is None:
+                    return
+                grad_buf = np.zeros_like(self.data)
+                grad_buf[idx] = out.grad
+                self.grad += _sum_to_shape(grad_buf, self.data.shape)
+
+        out._backward = _backward
+        return out
+
     def mean(self, axis=None, keepdims=False):
         out_data = self.data.mean(axis=axis, keepdims=keepdims)
         out = Tensor(out_data, requires_grad=self.requires_grad, name="mean")

scripts/debug_inspect.py

@@ -0,0 +1,36 @@
+import sys
+sys.path.insert(0, '.')
+from leanpass import Tensor, nn, optim
+import numpy as np
+
+model = nn.MLP([2,4,1])
+opt = optim.SGD(model.parameters(), lr=0.1)
+
+print('model __dict__ keys and types:')
+for k, v in model.__dict__.items():
+    print(k, type(v))
+
+params = model.parameters()
+print('model.parameters() ->', params)
+print('len=', len(params))
+
+x = Tensor([[1.0, 1.0]], requires_grad=False)
+y_true = Tensor([[2.0]], requires_grad=False)
+
+pred = model(x)
+loss = nn.mse_loss(pred, y_true)
+
+print('param data before:')
+for p in model.parameters():
+    print(repr(p), p.data)
+
+model.zero_grad()
+loss.backward()
+print('\nparam grads after backward:')
+for p in model.parameters():
+    print(repr(p), p.grad)
+
+opt.step()
+print('\nparam data after step:')
+for p in model.parameters():
+    print(repr(p), p.data)