Add zero_grads and change the meaning of step

Author: parsiad

src/micrograd_pp/__init__.py

@@ -1,4 +1,4 @@
-from ._expr import Constant, Expr, Parameter, is_grad_enabled, maximum, no_grad, relu
+from ._expr import Constant, Expr, Parameter, is_grad_enabled, maximum, no_grad, relu, zero_grads
 from ._func import cat, cross_entropy_loss, softmax
 from ._nn import (
     BatchNorm1d,
@@ -44,4 +44,5 @@
     "no_grad",
     "relu",
     "softmax",
+    "zero_grads",
 )

src/micrograd_pp/_expr.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import contextlib
+import functools
 import itertools
 from abc import ABC, abstractmethod
 from collections import deque
-from typing import Any, Callable, Generator, Sequence
+from typing import Any, Callable, Generator, Iterable, Sequence
 
 import numpy.typing as npt
 
@@ -150,6 +151,7 @@ def __sub__(self, other: int | float | Expr) -> Expr:
     def _backward(self, grad: npt.NDArray) -> None:
         del grad
 
+    @functools.lru_cache(maxsize=1)  # Cache for when loss.params is called multiple times
     def _get_nodes(self) -> deque[Expr]:
         retval: deque[Expr] = deque()
         if not self._requires_grad:
@@ -175,6 +177,10 @@ def visit(node: Expr) -> None:
         visit(self)
         return retval
 
+    @property
+    def params(self) -> list[Expr]:
+        return [node for node in self._get_nodes() if len(node._children) == 0]
+
     def backward(
         self,
         init: np.ndarray | float = 1.0,
@@ -214,7 +220,7 @@ def backward(
             if not retain_grad:
                 node._grad = None
         if opt is not None:
-            opt.step()
+            opt.update_state()
 
     def exp(self) -> Expr:
         """Return the element-wise exponential."""
@@ -356,6 +362,9 @@ def var(self, dim: int | tuple[int, ...] | None = None, keepdim: bool = False) -
             retval = _Squeeze(retval, dim=dim)
         return retval
 
+    def zero_grad(self) -> None:
+        self._grad = None
+
     @property
     def dtype(self) -> npt.DTypeLike:
         """Data type."""
@@ -390,13 +399,23 @@ def shape(self) -> tuple[int, ...]:
         return self._value.shape
 
 
+def zero_grads(params: Iterable[Expr]) -> None:
+    for param in params:
+        param.zero_grad()
+
+
 class Opt(ABC):
+    def step(self, params: Iterable[Expr]) -> None:
+        for param in params:
+            self.update_param(param)
+        self.update_state()
+
     @abstractmethod
     def update_param(self, param: Expr) -> None:
         pass
 
     @abstractmethod
-    def step(self) -> None:
+    def update_state(self) -> None:
         pass
 
 

src/micrograd_pp/_opt.py

@@ -17,7 +17,7 @@ def __init__(self, lr: float) -> None:
     def update_param(self, param: Expr) -> None:
         param.update_value(-self._lr * param.grad)
 
-    def step(self) -> None:
+    def update_state(self) -> None:
         pass
 
 
@@ -79,7 +79,7 @@ def __init__(
         self._moments: dict[Expr, tuple[np.ndarray, np.ndarray]] = {}
 
         self._t = 0
-        self.step()
+        self.update_state()
 
     def update_param(self, param: Expr) -> None:
         if param not in self._moments:
@@ -95,7 +95,7 @@ def update_param(self, param: Expr) -> None:
         update = -self._lr * (corrected_moment_1 / denom + self._weight_decay * param.value)
         param.update_value(update)
 
-    def step(self):
+    def update_state(self):
         self._t += 1
         self._bias_correction_1 = 1.0 - self._beta_1**self._t
         self._bias_correction_2 = 1.0 - self._beta_2**self._t

tests/test_opt.py

@@ -12,14 +12,17 @@ def run_before_and_after_tests():
 
 
 @pytest.mark.parametrize(
-    ("opt_factory", "num_steps", "atol"),
+    ("opt_factory", "num_steps", "atol", "pass_opt_to_backward"),
     [
-        (lambda: mpp.SGD(lr=0.1), 150, 1e-8),
-        (lambda: mpp.AdamW(lr=0.2, weight_decay=0.0), 600, 1e-8),
+        (*cfg, pass_opt_to_backward)
+        for cfg in (
+            (lambda: mpp.SGD(lr=0.1), 150, 1e-8),
+            (lambda: mpp.AdamW(lr=0.2, weight_decay=0.0), 600, 1e-8),
+        )
+        for pass_opt_to_backward in (False, True)
     ],
-    ids=("sgd", "adamw"),
 )
-def test_mse(opt_factory, num_steps: int, atol: float):
+def test_mse(opt_factory, num_steps: int, atol: float, pass_opt_to_backward: bool):
     n = 10
     coef = np.random.randn(3, 1)
     coef_hat = np.random.randn(3, 1)
@@ -35,6 +38,11 @@ def test_mse(opt_factory, num_steps: int, atol: float):
     for _ in range(num_steps):
         y_pred_ = x_ @ coef_hat_
         mse = ((y_pred_ - y_) ** 2).sum() / n
-        mse.backward(opt=opt)
+        if pass_opt_to_backward:
+            mse.backward(opt=opt)  # Automatically handles zeroing gradients and updating the optimizer state
+        else:
+            mpp.zero_grads(mse.params)
+            mse.backward()
+            opt.step(mse.params)
 
     np.testing.assert_allclose(coef, coef_hat, rtol=0.0, atol=atol)