diff --git a/build/cpu.o b/build/cpu.o index 6228ac2..054ce2e 100644 Binary files a/build/cpu.o and b/build/cpu.o differ diff --git a/build/cuda.cu.o b/build/cuda.cu.o index 376fab2..9809e05 100644 Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ diff --git a/build/tensor.o b/build/tensor.o index d330a48..5df7819 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 eeeaf80..d3ad7c7 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 b0f6eab..b9738eb 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 e4f91f5..0023139 100644 --- a/norch/autograd/functions.py +++ b/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: diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp index bde7b0c..a2f107b 100644 --- a/norch/csrc/cpu.cpp +++ b/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]; diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h index fad467c..2a97fd2 100644 --- a/norch/csrc/cpu.h +++ b/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); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 104994e..730d204 100644 --- a/norch/csrc/tensor.cpp +++ b/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); } } diff --git a/norch/libtensor.so b/norch/libtensor.so index 06f1149..9f15e45 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/tensor.py b/norch/tensor.py index 03b37ee..18192c3 100644 --- a/norch/tensor.py +++ b/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 diff --git a/tests/test_autograd.py b/tests/test_autograd.py index d40318a..7635e58 100644 --- a/tests/test_autograd.py +++ b/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 diff --git a/tests/test_operations.py b/tests/test_operations.py index 3a6f414..e7dd833 100644 --- a/tests/test_operations.py +++ b/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)