sum axis autograd

norch/autograd/functions.py

@@ -82,7 +82,11 @@ class MatmulBackward:
 
     def backward(self, gradient):
         x, y = self.input
-        return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
+        
+        if x.shape != y.shape: # broadcasted case
+            return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
+        else:
+            return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
         
 """class PowBackward:
     def __init__(self, x, power):
@@ -121,13 +125,22 @@ class LogBackward:
         return [grad_input]
        
 class SumBackward:
-    def __init__(self, x):
+    def __init__(self, x, axis=None):
         self.input = [x]
+        self.axis = axis
 
     def backward(self, gradient):
-        # Since sum reduces a tensor to a scalar, gradient is broadcasted to match the original shape.
-        return [float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()]
-    
+        input_shape = self.input[0].shape
+        if self.axis is None:
+            # If axis is None, sum reduces the tensor to a scalar.
+            grad_output = float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()
+        else:
+            # Broadcast the gradient to the input shape along the specified axis.
+            grad_output_shape = [1 if i in self.axis else dim for i, dim in enumerate(input_shape)]
+            grad_output = gradient.reshape(grad_output_shape)
+            grad_output = grad_output * self.input[0].ones_like()
+        
+        return [grad_output]
 class ReshapeBackward:
     def __init__(self, x):
         self.input = [x]

tests/test_autograd.py

@@ -34,6 +34,28 @@ class TestTensorAutograd(unittest.TestCase):
         self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
         self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
 
+    def test_sum_axis(self):
+        """
+        Test autograd from sum specifying axis
+        """
+        norch_tensor1 = norch.Tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        norch_tensor2 = norch.Tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        norch_result = (norch_tensor1 + norch_tensor2).sum(axis=0).sum(axis=0).sum()
+        norch_result.backward()
+        norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
+        norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
+
+        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        torch_result = (torch_tensor1 + torch_tensor2).sum(axis=0).sum(axis=0).sum()
+        torch_result.backward()
+        torch_tensor1_grad = torch_tensor1.grad
+        torch_tensor2_grad = torch_tensor2.grad
+
+        self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
+        self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
+
+
     def test_broadcasting_addition_autograd(self):
         """
         Test autograd for broadcasting addition: tensor1 + tensor2