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)