def test_axis_sum(self) -> None: t1 = Tensor([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]], requires_grad=True) # First axis t2 = t1.sum(0) assert t2.data.tolist() == [2, 4, 6] t2.backward(np.asarray([1.0, 2.0, 3.0])) assert t1.grad.tolist() == [[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]] # Second axis t1.zero_grad() t3 = t1.sum(1) assert t3.data.tolist() == [6, 6] t3.backward(np.asarray([1.0, 2.0])) assert t1.grad.tolist() == [[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]
def test_simple_sum(self) -> None: t1 = Tensor([1, 2, 3], requires_grad=True) t2 = t1.sum() assert t2.data.tolist() == 6 assert t2.requires_grad t2.backward() assert t1.grad.tolist() == [1.0, 1.0, 1.0] # Also try with a specified grad t1.zero_grad() t3 = t1.sum() assert t3.data.tolist() == 6 t2.backward(2) assert t1.grad.tolist() == [2.0, 2.0, 2.0]
def test_simple_sum_no_grad(self) -> None: t1 = Tensor([1, 2, 3]) t2 = t1.sum() assert t2.data.tolist() == 6 assert not t2.requires_grad