fix sum axis autograd and reversed add broadcasted order

norch/autograd/functions.py

@@ -84,9 +84,9 @@ class MatmulBackward:
         x, y = self.input
         
         if x.ndim != y.ndim: # broadcasted case
-            print("KKKKKKKKK", x.ndim, y.ndim)
             aux = (gradient @ y.transpose(-1,-2))
-            return [aux.sum(axis=0), x.transpose(-1,-2) @ gradient]
+            aux_sum = aux.sum(axis=0)
+            return [aux_sum, x.transpose(-1,-2) @ gradient]
         else:
             return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
         
@@ -133,14 +133,15 @@ class SumBackward:
 
     def backward(self, gradient):
         input_shape = self.input[0].shape
-        if self.axis is None:
+        if self.axis == -1:
             # 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_shape = list(input_shape)
+            grad_output_shape[self.axis] = 1
             grad_output = gradient.reshape(grad_output_shape)
-            grad_output = grad_output * self.input[0].ones_like()
+            grad_output = grad_output + self.input[0].zeros_like()
         
         return [grad_output]
 class ReshapeBackward:

norch/csrc/cpu.cpp

@@ -11,7 +11,7 @@ void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
     }
 }
 
-void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
+void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) {
     int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
 
     // Calculate strides for broadcasting
@@ -28,12 +28,12 @@ void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_
         int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
         strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
         strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
-        stride1 *= broadcasted_shape[i];
-        stride2 *= broadcasted_shape[i];
+        stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
+        stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
     }
 
     // Perform element-wise addition with broadcasting
-    for (int i = 0; i < tensor1->size; i++) {
+    for (int i = 0; i < broadcasted_size; i++) {
         int index1 = 0, index2 = 0;
         int linear_index = i;
         for (int j = max_ndim - 1; j >= 0; j--) {
@@ -219,7 +219,11 @@ void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis) {
             return;
         }
         
-        int result_shape[tensor->ndim - 1];
+        int* result_shape = (int*)malloc((tensor->ndim - 1) * sizeof(int));
+        if (result_shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
         int result_size = 1;
         int axis_stride = tensor->strides[axis];
 

norch/csrc/cpu.h

@@ -4,7 +4,7 @@
 #include "tensor.h"
 
 void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
-void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape);
+void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
 void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis);
 void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape);

norch/csrc/tensor.cpp

@@ -149,20 +149,25 @@ extern "C" {
             broadcasted_shape[max_ndim - 1 - i] = dim1 > dim2 ? dim1 : dim2;
         }
 
+        int broadcasted_size = 1;
+        for (int i = 0; i < max_ndim; i++) {
+            broadcasted_size *= broadcasted_shape[i];
+        }
+
         if (strcmp(tensor1->device, "cuda") == 0) {
             float* result_data;
-            cudaMalloc((void **)&result_data, tensor1->size * sizeof(float));
+            cudaMalloc((void **)&result_data, broadcasted_size * sizeof(float));
             add_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape);
             return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
         } 
         else {
-            float* result_data = (float*)malloc(tensor1->size * sizeof(float));
+            float* result_data = (float*)malloc(broadcasted_size * sizeof(float));
             if (result_data == NULL) {
                 fprintf(stderr, "Memory allocation failed\n");
                 exit(1);
             }
 
-            add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
+            add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size);
             return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
         }
     }

norch/tensor.py

@@ -653,7 +653,7 @@ class Tensor:
 
         result_data.requires_grad = self.requires_grad
         if result_data.requires_grad:
-            result_data.grad_fn = SumBackward(self)
+            result_data.grad_fn = SumBackward(self, axis)
 
         return result_data
 

tests/test_autograd.py

@@ -41,17 +41,19 @@ class TestTensorAutograd(unittest.TestCase):
         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))
 
@@ -76,6 +78,21 @@ 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))
+
+        ## reversed order broadcasting
+
+        norch_result = (norch_tensor2 + norch_tensor1).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_result = (torch_tensor2 + torch_tensor1).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_subtraction(self):
         """
@@ -267,7 +284,7 @@ class TestTensorAutograd(unittest.TestCase):
 
     def test_broadcasted_batched_matmul(self):
         """
-        Test autograd from batched matrix multiplication: BxMxP = NxM @ BxMxP
+        Test autograd from broadcasted batched matrix multiplication: BxMxP = NxM @ BxMxP
         """
         B = 3  # Batch size
 

tests/test_operations.py

@@ -37,7 +37,7 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
-    def test_broadcasting_addition(self):
+    def test_addition_broadcasted(self):
         """
         Test addition of two tensors with broadcasting: tensor1 + tensor2
         """
@@ -52,6 +52,18 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
+        norch_tensor1 = norch.Tensor([[0, 2]]).to(self.device) 
+        norch_tensor2 = norch.Tensor([[3, 4], [5, -1]]).to(self.device) 
+        norch_result = norch_tensor1 + norch_tensor2
+        torch_result = utils.to_torch(norch_result).to(self.device)
+
+        torch_tensor1 = torch.tensor([[0, 2]]).to(self.device)  
+        torch_tensor2 = torch.tensor([[3, 4], [5, -1]]).to(self.device) 
+        torch_expected = torch_tensor1 + torch_tensor2
+
+        print(torch_result, torch_expected)
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+
         norch_result = norch_tensor2 + norch_tensor1
         torch_expected = torch_tensor2 + torch_tensor1
 
@@ -244,7 +256,7 @@ class TestTensorOperations(unittest.TestCase):
 
     def test_matmul(self):
         """
-        Test batched matrix multiplication: MxP = NxM @ MxP
+        Test matrix multiplication: MxP = NxM @ MxP
         """
         # Creating batched tensors for Norch
         norch_tensor1 = norch.Tensor([[1, 2], [3, -4], [5, 6], [7, 8]]).to(self.device)