diff --git a/build/cpu.o b/build/cpu.o index 7792721..6228ac2 100644 Binary files a/build/cpu.o and b/build/cpu.o differ diff --git a/build/tensor.o b/build/tensor.o index f5f937d..877b650 100644 Binary files a/build/tensor.o and b/build/tensor.o differ diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index d2c4826..5a33b4a 100644 Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc index 54c5754..1365c91 100644 Binary files a/norch/autograd/__pycache__/functions.cpython-38.pyc and b/norch/autograd/__pycache__/functions.cpython-38.pyc differ diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py index 46b5b36..8ce8625 100644 --- a/norch/autograd/functions.py +++ b/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] diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp index 5a0cb29..bde7b0c 100644 --- a/norch/csrc/cpu.cpp +++ b/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++) { diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h index 5e64293..fad467c 100644 --- a/norch/csrc/cpu.h +++ b/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); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 28501d2..ad56c94 100644 --- a/norch/csrc/tensor.cpp +++ b/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); diff --git a/norch/libtensor.so b/norch/libtensor.so index aa08bb8..dd54141 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/tensor.py b/norch/tensor.py index f50edf7..03b37ee 100644 --- a/norch/tensor.py +++ b/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 diff --git a/tests/test_autograd.py b/tests/test_autograd.py index 4dc44eb..ead6e1a 100644 --- a/tests/test_autograd.py +++ b/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 diff --git a/tests/test_operations.py b/tests/test_operations.py index ed4ee46..a48f7c0 100644 --- a/tests/test_operations.py +++ b/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):