sub broadcasted cpu autograd

norch/autograd/functions.py

@@ -28,10 +28,30 @@ class AddBroadcastedBackward:
                 gradient = gradient.sum(axis=i)
 
         return gradient
-
-
     
+
+class SubBroadcastedBackward:
+    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
+
 class SubBackward:
     def __init__(self, x, y):
         self.input = [x, y]

norch/csrc/cpu.cpp

@@ -58,6 +58,45 @@ 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) {
+    int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
+
+    // Calculate strides for broadcasting
+    int* strides1 = (int*)malloc(max_ndim * sizeof(int));
+    int* strides2 = (int*)malloc(max_ndim * sizeof(int));
+    if (strides1 == NULL || strides2 == NULL) {
+        fprintf(stderr, "Memory allocation failed\n");
+        exit(1);
+    }
+
+    int stride1 = 1, stride2 = 1;
+    for (int i = max_ndim - 1; i >= 0; i--) {
+        int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - max_ndim + i] : 1;
+        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];
+    }
+
+    // Perform element-wise addition with broadcasting
+    for (int i = 0; i < tensor1->size; i++) {
+        int index1 = 0, index2 = 0;
+        int linear_index = i;
+        for (int j = max_ndim - 1; j >= 0; j--) {
+            int pos = linear_index % broadcasted_shape[j];
+            linear_index /= broadcasted_shape[j];
+            if (strides1[j] != 0) index1 += pos * strides1[j];
+            if (strides2[j] != 0) index2 += pos * strides2[j];
+        }
+        result_data[i] = tensor1->data[index1] - tensor2->data[index2];
+    }
+
+    // Free strides
+    free(strides1);
+    free(strides2);
+}
+
 void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
     
     for (int i = 0; i < tensor1->size; i++) {

norch/csrc/cpu.h

@@ -7,6 +7,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 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);
 void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void scalar_div_tensor_cpu(float scalar, Tensor* tensor, float* result_data);
 void tensor_div_scalar_cpu(Tensor* tensor, float scalar, float* result_data);

norch/csrc/tensor.cpp

@@ -259,6 +259,43 @@ extern "C" {
         }
     }
 
+    Tensor* sub_broadcasted_tensor(Tensor* tensor1, Tensor* tensor2) {
+
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
+
+        // Determine the broadcasted shape
+        int* broadcasted_shape = (int*)malloc(max_ndim * sizeof(int));
+        if (broadcasted_shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        for (int i = 0; i < max_ndim; i++) {
+            int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - 1 - i] : 1;
+            int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - 1 - i] : 1;
+            if (dim1 != dim2 && dim1 != 1 && dim2 != 1) {
+                fprintf(stderr, "Shapes are not compatible for broadcasting\n");
+                exit(1);
+            }
+            broadcasted_shape[max_ndim - 1 - i] = dim1 > dim2 ? dim1 : dim2;
+        }
+
+        // Allocate memory for result tensor
+        float* result_data = (float*)malloc(tensor1->size * sizeof(float));
+        if (result_data == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        sub_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
+
+        return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
+    }
+
     Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2) {
         if (tensor1->ndim != tensor2->ndim) {
             fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for element-wise multiplication\n", tensor1->ndim, tensor2->ndim);

norch/tensor.py

@@ -312,24 +312,61 @@ class Tensor:
         if isinstance(other, (int, float)):
             other = other * self.ones_like()
 
-        if self.shape != other.shape:
-            raise ValueError("Tensors must have the same shape for subtraction")
+        # Function to determine if broadcasting is needed and get the broadcasted shape
+        def broadcast_shape(shape1, shape2):
+            if shape1 == shape2:
+                return shape1, False
+            
+            max_len = max(len(shape1), len(shape2))
+            shape1 = [1] * (max_len - len(shape1)) + shape1
+            shape2 = [1] * (max_len - len(shape2)) + shape2
+
+            broadcasted_shape = []
+            for dim1, dim2 in zip(shape1, shape2):
+                if dim1 != dim2 and dim1 != 1 and dim2 != 1:
+                    raise ValueError("Shapes are not compatible for broadcasting")
+                broadcasted_shape.append(max(dim1, dim2))
+            return broadcasted_shape, True
+
+        broadcasted_shape, needs_broadcasting = broadcast_shape(self.shape, other.shape)
+
+        if needs_broadcasting:
+            # Call add_broadcasted_tensor if broadcasting is needed
+            Tensor._C.sub_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.sub_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
+
+            result_tensor_ptr = Tensor._C.sub_broadcasted_tensor(self.tensor, other.tensor)
+
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = broadcasted_shape
+            result_data.ndim = len(broadcasted_shape)
+
+            result_data.device = self.device
+            result_data.numel = self.numel  # Update this to calculate the correct number of elements if broadcasting
+
+            result_data.requires_grad = self.requires_grad or other.requires_grad
+            if result_data.requires_grad:
+                result_data.grad_fn = SubBroadcastedBackward(self, other)
         
-        Tensor._C.sub_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
-        Tensor._C.sub_tensor.restype = ctypes.POINTER(CTensor)
+        else:
+            # Call add_tensor if shapes are identical
+            Tensor._C.sub_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.sub_tensor.restype = ctypes.POINTER(CTensor)
 
-        result_tensor_ptr = Tensor._C.sub_tensor(self.tensor, other.tensor)
+            result_tensor_ptr = Tensor._C.sub_tensor(self.tensor, other.tensor)
 
-        result_data = Tensor()
-        result_data.tensor = result_tensor_ptr
-        result_data.shape = self.shape.copy()
-        result_data.ndim = self.ndim
-        result_data.device = self.device
-        result_data.numel = self.numel
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = self.shape.copy()
+            result_data.ndim = self.ndim
 
-        result_data.requires_grad = self.requires_grad or other.requires_grad
-        if result_data.requires_grad:
-            result_data.grad_fn = SubBackward(self, other)
+            result_data.device = self.device
+            result_data.numel = self.numel  # Update this to calculate the correct number of elements if broadcasting
+
+            result_data.requires_grad = self.requires_grad or other.requires_grad
+            if result_data.requires_grad:
+                result_data.grad_fn = SubBackward(self, other)
 
         return result_data
     

tests/test_autograd.py

@@ -73,6 +73,28 @@ class TestTensorAutograd(unittest.TestCase):
         self.assertTrue(utils.compare_torch(norch_tensor1_grad_sub, torch_tensor1_grad_sub))
         self.assertTrue(utils.compare_torch(norch_tensor2_grad_sub, torch_tensor2_grad_sub))
         
+    def test_broadcasting_subtraction_autograd(self):
+        """
+        Test autograd for broadcasting subtraction: tensor1 - tensor2
+        """
+        norch_tensor1 = norch.Tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device)  # Shape (1, 2, 3)
+        norch_tensor2 = norch.Tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+        norch_result = (norch_tensor1 - norch_tensor2).sum()
+        norch_result.backward()
+        norch_tensor1_grad = utils.to_torch(norch_tensor1.grad)
+        norch_tensor2_grad = utils.to_torch(norch_tensor2.grad)
+
+        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+        torch_result = (torch_tensor1 - torch_tensor2).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_division(self):
         """
         Test autograd from dividing two tensors: tensor1 / tensor2

tests/test_operations.py

@@ -46,7 +46,7 @@ class TestTensorOperations(unittest.TestCase):
 
         torch_tensor1 = torch.tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device)  # Shape (1, 2, 3)
         torch_tensor2 = torch.tensor([1, 1, 1]).to(self.device)  # Shape (3)
-        torch_expected = torch_tensor1 + torch_tensor2  # Broadcasting addition
+        torch_expected = torch_tensor1 + torch_tensor2
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
@@ -66,6 +66,22 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
+    def test_broadcasting_subtraction(self):
+        """
+        Test subtraction of two tensors with broadcasting: tensor1 - tensor2
+        """
+        norch_tensor1 = norch.Tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device)  # Shape (1, 2, 3)
+        norch_tensor2 = norch.Tensor([1, 1, 1]).to(self.device)  # Shape (3)
+        norch_result = norch_tensor1 - norch_tensor2
+        torch_result = utils.to_torch(norch_result)
+
+        torch_tensor1 = torch.tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1, 1, 1]).to(self.device)  # Shape (3)
+        torch_expected = torch_tensor1 - torch_tensor2 
+
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+
+
     def test_division_by_scalar(self):
         """
         Test division of a tensor by a scalar: tensor / scalar
@@ -202,7 +218,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)
-        
+
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
     def test_transpose_T(self):