diff --git a/build/cpu.o b/build/cpu.o new file mode 100644 index 0000000..7792721 Binary files /dev/null and b/build/cpu.o differ diff --git a/build/cuda.cu.o b/build/cuda.cu.o new file mode 100644 index 0000000..8734644 Binary files /dev/null and b/build/cuda.cu.o differ diff --git a/build/tensor.o b/build/tensor.o new file mode 100644 index 0000000..f5f937d Binary files /dev/null and b/build/tensor.o differ diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index 01f4c49..d2c4826 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 87e1cfd..1b23177 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 d5a2684..b114534 100644 --- a/norch/autograd/functions.py +++ b/norch/autograd/functions.py @@ -7,6 +7,12 @@ class AddBackward: def backward(self, gradient): return [gradient, gradient] +class AddBroadcastedBackward: + pass + + + + class SubBackward: def __init__(self, x, y): self.input = [x, y] diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp index ee46dca..5a0cb29 100644 --- a/norch/csrc/cpu.cpp +++ b/norch/csrc/cpu.cpp @@ -11,6 +11,46 @@ 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) { + 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 sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { for (int i = 0; i < tensor1->size; i++) { @@ -126,16 +166,52 @@ void log_tensor_cpu(Tensor* tensor, float* result_data) { } } -void sum_tensor_cpu(Tensor* tensor, float* result_data) { - float sum = 0.0; +void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis) { + if (axis == -1) { + // Sum over all elements + float sum = 0.0; + for (int i = 0; i < tensor->size; i++) { + sum += tensor->data[i]; + } + *result_data = sum; + } else { + if (axis < 0 || axis >= tensor->ndim) { + printf("Invalid axis"); + return; + } + + int result_shape[tensor->ndim - 1]; + int result_size = 1; + int axis_stride = tensor->strides[axis]; - for (int i = 0; i < tensor->size; i++) { - sum += tensor->data[i]; + int idx = 0; + for (int i = 0; i < tensor->ndim; i++) { + if (i != axis) { + result_shape[idx++] = tensor->shape[i]; + result_size *= tensor->shape[i]; + } + } + + memset(result_data, 0, result_size * sizeof(float)); + + for (int i = 0; i < tensor->shape[axis]; i++) { + for (int j = 0; j < result_size; j++) { + int index = 0; + int remainder = j; + for (int k = tensor->ndim - 2; k >= 0; k--) { + index += (remainder % result_shape[k]) * tensor->strides[k < axis ? k : k + 1]; + remainder /= result_shape[k]; + } + result_data[j] += tensor->data[index + i * axis_stride]; + } + } } - - *result_data = sum; } + + + + void ones_like_tensor_cpu(Tensor* tensor, float* result_data) { for (int i = 0; i < tensor->size; i++) { diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h index 9e5dc62..5e64293 100644 --- a/norch/csrc/cpu.h +++ b/norch/csrc/cpu.h @@ -4,7 +4,8 @@ #include "tensor.h" void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); -void sum_tensor_cpu(Tensor* tensor1, 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 elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); void scalar_div_tensor_cpu(float scalar, Tensor* tensor, float* result_data); diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index 837e720..28501d2 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -124,7 +124,44 @@ extern "C" { } } - Tensor* sum_tensor(Tensor* tensor) { + Tensor* add_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); + } + + add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape); + + return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device); + } + + Tensor* sum_tensor(Tensor* tensor, int axis) { char* device = (char*)malloc(strlen(tensor->device) + 1); if (device != NULL) { @@ -133,14 +170,25 @@ extern "C" { fprintf(stderr, "Memory allocation failed\n"); exit(-1); } - int ndim = 1; - int* shape = (int*)malloc(ndim * sizeof(int)); - if (shape == NULL) { - fprintf(stderr, "Memory allocation failed\n"); - exit(1); + int ndim; + int* shape; + if (axis == -1) { + shape = (int*) malloc(sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + shape[0] = 1; + ndim = 1; + } else { + shape = (int*) malloc((tensor->ndim - 1) * sizeof(int)); + for (int i = 0, j = 0; i < tensor->ndim; ++i) { + if (i != axis) { + shape[j++] = tensor->shape[i]; + } + } + ndim = tensor->ndim - 1; } - - shape[0] = 1; if (strcmp(tensor->device, "cuda") == 0) { @@ -155,7 +203,7 @@ extern "C" { fprintf(stderr, "Memory allocation failed\n"); exit(1); } - sum_tensor_cpu(tensor, result_data); + sum_tensor_cpu(tensor, result_data, axis); return create_tensor(result_data, shape, ndim, device); } } diff --git a/norch/csrc/tensor.h b/norch/csrc/tensor.h index d45901f..5ff62d0 100644 --- a/norch/csrc/tensor.h +++ b/norch/csrc/tensor.h @@ -16,7 +16,7 @@ extern "C" { Tensor* create_tensor(float* data, int* shape, int ndim, char* device); float get_item(Tensor* tensor, int* indices); Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2); - Tensor* sum_tensor(Tensor* tensor); + Tensor* sum_tensor(Tensor* tensor, int axis); Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2); Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2); Tensor* scalar_mul_tensor(Tensor* tensor, float scalar); diff --git a/norch/libtensor.so b/norch/libtensor.so index 5dfcec6..aa08bb8 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/tensor.py b/norch/tensor.py index 53b8416..f50edf7 100644 --- a/norch/tensor.py +++ b/norch/tensor.py @@ -223,27 +223,65 @@ class Tensor: def __add__(self, other): if isinstance(other, (int, float)): other = other * self.ones_like() + + # 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.add_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.add_broadcasted_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.add_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 = AddBroadcastedBackward(self, other) - if self.shape != other.shape: - raise ValueError("Tensors must have the same shape for addition") - - Tensor._C.add_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] - Tensor._C.add_tensor.restype = ctypes.POINTER(CTensor) + else: + # Call add_tensor if shapes are identical + Tensor._C.add_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.add_tensor.restype = ctypes.POINTER(CTensor) - result_tensor_ptr = Tensor._C.add_tensor(self.tensor, other.tensor) + result_tensor_ptr = Tensor._C.add_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 = AddBackward(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 = AddBackward(self, other) return result_data + def __radd__(self, other): if isinstance(other, (int, float)): @@ -555,24 +593,33 @@ class Tensor: return result_data - def sum(self): - Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor)] + def sum(self, axis=-1): + Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int] Tensor._C.sum_tensor.restype = ctypes.POINTER(CTensor) - result_tensor_ptr = Tensor._C.sum_tensor(self.tensor) + result_tensor_ptr = Tensor._C.sum_tensor(self.tensor, axis) result_data = Tensor() result_data.tensor = result_tensor_ptr - result_data.shape = [1] - result_data.ndim = 1 + + if axis == -1: + result_data.shape = [1] + result_data.ndim = 1 + else: + result_data.shape = self.shape[:axis] + self.shape[axis+1:] + result_data.ndim = len(result_data.shape) + result_data.device = self.device result_data.numel = 1 + for s in result_data.shape: + result_data.numel *= s result_data.requires_grad = self.requires_grad if result_data.requires_grad: result_data.grad_fn = SumBackward(self) return result_data + def sin(self): Tensor._C.sin_tensor.argtypes = [ctypes.POINTER(CTensor)] diff --git a/tests/test_autograd.py b/tests/test_autograd.py index a042ce6..4dc44eb 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -31,6 +31,26 @@ 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)) + def test_broadcasting_addition_autograd(self): + """ + Test autograd for broadcasting addition: 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_subtraction(self): """ diff --git a/tests/test_operations.py b/tests/test_operations.py index e3fac6a..3eeb8b1 100644 --- a/tests/test_operations.py +++ b/tests/test_operations.py @@ -35,6 +35,22 @@ class TestTensorOperations(unittest.TestCase): self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + def test_broadcasting_addition(self): + """ + Test addition 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 # Broadcasting addition + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_subtraction(self): """ Test subtraction of two tensors: tensor1 - tensor2 @@ -176,6 +192,21 @@ class TestTensorOperations(unittest.TestCase): self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + def test_sum_axis_no_keepdims(self): + """ + Test summation of a tensor along a specific axis without keeping the dimensions + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.sum(axis=1) + torch_result = utils.to_torch(norch_result) + + 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): """ Test transposition of a tensor: tensor.T