Store _prev as a tuple in sum() and mean()

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}
 

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)