From 3e9f48369576b2381eb47285ad0e2f329eccfbb7 Mon Sep 17 00:00:00 2001 From: eeshsaxena Date: Sat, 8 Aug 2026 14:39:58 +0530 Subject: [PATCH] Store _prev as a tuple in sum() and mean() Every op builds its _prev through _create_child, which normalises it with tuple(prev). sum() and mean() instead set out._prev = {self} directly, so those two nodes carried a set. _eval_forward reads node._prev[0] to pull an operand, and indexing a set raises "TypeError: 'set' object is not subscriptable". So evaluating any graph that reduces with sum or mean, which is almost every loss, crashed. Store (self,) in both methods so _prev is a tuple everywhere. Added a test that evaluates a sum and a mean graph through _eval_forward. --- leanpass/tensor.py | 4 ++-- tests/test_gradcheck.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/leanpass/tensor.py b/leanpass/tensor.py index 1f05a24..6da99b0 100644 --- a/leanpass/tensor.py +++ b/leanpass/tensor.py @@ -241,7 +241,7 @@ def _backward(): 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") - out._prev = {self} + out._prev = (self,) out._op = "sum" out._meta = {"axis": axis, "keepdims": keepdims, "shape": self.data.shape} @@ -260,7 +260,7 @@ def _backward(): 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") - out._prev = {self} + out._prev = (self,) out._op = "mean" out._meta = {"axis": axis, "keepdims": keepdims, "shape": self.data.shape} diff --git a/tests/test_gradcheck.py b/tests/test_gradcheck.py index d30001a..7a9fc63 100644 --- a/tests/test_gradcheck.py +++ b/tests/test_gradcheck.py @@ -79,3 +79,15 @@ def test_visualize_dot(): assert dot.startswith("digraph") assert "->" in dot assert "input" in dot + + +def test_eval_forward_handles_reductions(): + # sum() and mean() stored their _prev as a set instead of a tuple, so + # _eval_forward's `node._prev[0]` indexing raised + # "TypeError: 'set' object is not subscriptable" for any graph that reduces + # with sum or mean, which is almost every loss. + x = Tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True) + assert np.isclose((x * x).sum()._eval_forward(), 30.0) + + y = Tensor([1.0, 2.0, 3.0], requires_grad=True) + assert np.isclose((y * y).mean()._eval_forward(), 14.0 / 3.0)