add broadcasting autograd

norch/autograd/functions.py

@@ -8,7 +8,26 @@ class AddBackward:
         return [gradient, gradient]
     
 class AddBroadcastedBackward:
-    pass
+    def __init__(self, x, y):
+        self.input = [x, y]
+
+    def backward(self, gradient):
+        x, y = self.input
+        grad_x = self._reshape_gradient(gradient, x.shape)
+        grad_y = self._reshape_gradient(gradient, y.shape)
+        return [grad_x, grad_y]
+    
+    def _reshape_gradient(self, gradient, shape):
+        # Reduce gradient dimensions to match the target shape dimensions
+        while len(gradient.shape) > len(shape):
+            gradient = gradient.sum(axis=0)
+
+        # Sum along axes where the target shape dimension is 1
+        for i in range(len(shape)):
+            if shape[i] == 1:
+                gradient = gradient.sum(axis=i)
+
+        return gradient
 
 
     

tests/test_operations.py

@@ -192,7 +192,7 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
-    def test_sum_axis_no_keepdims(self):
+    def test_sum_axis(self):
         """
         Test summation of a tensor along a specific axis without keeping the dimensions
         """
@@ -202,9 +202,7 @@ class TestTensorOperations(unittest.TestCase):
 
         torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
         torch_expected = torch.sum(torch_tensor, dim=1)
-
-        print(torch_result)
-        print("\n\n", torch_expected)
+        
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
     def test_transpose_T(self):