diff --git a/build/lib/norch/__init__.py b/build/lib/norch/__init__.py new file mode 100644 index 0000000..c2a3122 --- /dev/null +++ b/build/lib/norch/__init__.py @@ -0,0 +1,9 @@ +from norch.tensor import * +from .nn import * +from .optim import * +from .utils import * +from .norchvision import * + +__version__ = "0.0.4" +__author__ = 'Lucas de Lima Nogueira' +__credits__ = 'Lucas de Lima Nogueira' \ No newline at end of file diff --git a/build/lib/norch/autograd/__init__.py b/build/lib/norch/autograd/__init__.py new file mode 100644 index 0000000..e7d4dc0 --- /dev/null +++ b/build/lib/norch/autograd/__init__.py @@ -0,0 +1 @@ +from .functions import * \ No newline at end of file diff --git a/build/lib/norch/autograd/functions.py b/build/lib/norch/autograd/functions.py new file mode 100644 index 0000000..477837d --- /dev/null +++ b/build/lib/norch/autograd/functions.py @@ -0,0 +1,304 @@ +import math +import norch + +class AddBackward: + def __init__(self, x, y): + self.input = [x, y] + + def backward(self, gradient): + return [gradient, gradient] + +class AddBroadcastedBackward: + 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, keepdim=True) + + return gradient + +class SubBackward: + def __init__(self, x, y): + self.input = [x, y] + + def backward(self, gradient): + return [gradient, -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 ScalarMulBackward: + def __init__(self, x, scalar): + self.input = [x] + self.scalar = scalar + + def backward(self, gradient): + return [gradient * self.scalar] + +class ElementwiseMulBackward: + def __init__(self, x, y): + self.input = [x, y] + + def backward(self, gradient): + x = self.input[0] + y = self.input[1] + return [y * gradient, x * gradient] + +class MatmulBackward: + def __init__(self, x, y): + self.input = [x, y] + + def backward(self, gradient): + x, y = self.input + + if x.ndim != y.ndim: # broadcasted case + aux = (gradient @ y.transpose(-1,-2)) + 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] + +"""class PowBackward: + def __init__(self, x, power): + self.input = [x] + self.power = power + + def backward(self, gradient): + return [(gradient * self.power) * (self.input[0]) ** (self.power - 1)]""" + +class PowBackward: + def __init__(self, base, exponent): + self.input = [base, exponent] + + def backward(self, gradient): + base, exponent = self.input[0], self.input[1] + + if isinstance(base, (int, float)): + grad_base = gradient * (base ** (exponent - 1)) + grad_exponent = (gradient * base ** exponent) * math.log(base) + + else: + grad_base = gradient * exponent * (base ** (exponent - 1)) + grad_exponent = (gradient * base ** exponent) * (base.log()) + + return [grad_base, grad_exponent] + + +class LogBackward: + def __init__(self, x): + self.input = [x] + + def backward(self, gradient): + + grad_input = gradient / self.input[0] + + return [grad_input] + +class SumBackward: + def __init__(self, x, axis=None, keepdim=False): + self.input = [x] + self.axis = axis + self.keepdim = keepdim + + def backward(self, gradient): + input_shape = self.input[0].shape.copy() + if self.axis == -1: + # If axis is None, sum reduces the tensor to a scalar. + grad_output = float(gradient[[0] * len(gradient.shape)]) * self.input[0].ones_like() + else: + + if self.keepdim: + input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:] + else: + input_shape = input_shape[:self.axis] + input_shape[self.axis+1:] + + # Broadcast the gradient to the input shape along the specified axis. + grad_output_shape = list(input_shape) + grad_output = gradient.reshape(grad_output_shape) + grad_output = grad_output + self.input[0].zeros_like() + + return [grad_output] + +class ReshapeBackward: + def __init__(self, x): + self.input = [x] + + def backward(self, gradient): + return [gradient.reshape(self.input[0].shape)] + +class TransposeBackward: + def __init__(self, x, axis1, axis2): + self.input = [x] + self.axis1 = axis1 + self.axis2 = axis2 + + def backward(self, gradient): + return [gradient.transpose(self.axis2, self.axis1)] + +class TBackward: + def __init__(self, x): + self.input = [x] + + def backward(self, gradient): + return [gradient.T] + + +class DivisionBackward: + def __init__(self, x, y): + self.input = [x, y] + + def backward(self, gradient): + x, y = self.input + grad_x = gradient / y + grad_y = -1 * gradient * (x / (y * y)) + + return [grad_x, grad_y] + +class SinBackward: + def __init__(self, x): + self.input = [x] + + def backward(self, gradient): + x = self.input[0] + return [gradient * x.cos()] + +class CosBackward: + def __init__(self, x): + self.input = [x] + + def backward(self, gradient): + x = self.input[0] + return [-gradient * x.sin()] + +class MaxBackward: + def __init__(self, x, axis=None, keepdim=False): + self.input = [x] + self.axis = axis + self.keepdim = keepdim + + def backward(self, gradient): + input_shape = self.input[0].shape.copy() + if self.axis == -1: + max_value = self.input[0].max() + mask = self.input[0].equal(max_value) + + grad_output = float(gradient[[0] * len(gradient.shape)]) * self.input[0].ones_like() + + grad_output = (grad_output * mask) / mask.sum()[0] + + else: + + if self.keepdim: + input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:] + else: + input_shape = input_shape[:self.axis] + input_shape[self.axis+1:] + + # Broadcast the gradient to the input shape along the specified axis. + grad_output_shape = list(input_shape) + + grad_output = gradient.reshape(grad_output_shape) + grad_output = grad_output + self.input[0].zeros_like() + max_values = self.input[0].max(axis=self.axis, keepdim=True) + mask = self.input[0].equal(max_values) + + grad_output = (grad_output * mask) + + return [grad_output] + + +class MinBackward: + def __init__(self, x, axis=None, keepdim=False): + self.input = [x] + self.axis = axis + self.keepdim = keepdim + + def backward(self, gradient): + input_shape = self.input[0].shape.copy() + if self.axis == -1: + min_value = self.input[0].min() + mask = self.input[0].equal(min_value) + + grad_output = float(gradient[[0] * len(gradient.shape)]) * self.input[0].ones_like() + + grad_output = (grad_output * mask) / mask.sum()[0] + + else: + if self.keepdim: + input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:] + else: + input_shape = input_shape[:self.axis] + input_shape[self.axis+1:] + + # Broadcast the gradient to the input shape along the specified axis. + grad_output_shape = list(input_shape) + grad_output = gradient.reshape(grad_output_shape) + grad_output = grad_output + self.input[0].zeros_like() + max_values = self.input[0].min(axis=self.axis, keepdim=True) + mask = self.input[0].equal(max_values) + + grad_output = (grad_output * mask) + + return [grad_output] + + +class CrossEntropyLossBackward: + def __init__(self, logits, targets): + self.input = [logits, targets] + + def backward(self, gradient): + logits, targets = self.input + + if logits.ndim == 1: + softmax = norch.softmax(logits, dim=0) + grad_logits = (softmax - targets) + + elif logits.ndim == 2: + # batched + batch_size = logits.shape[0] + softmax = norch.softmax(logits, dim=1) + + grad_logits = (softmax - targets) / batch_size + + return [grad_logits, None] # targets do not have a gradient + +class SigmoidBackward: + def __init__(self, input): + self.input = [input] + + def backward(self, gradient): + sigmoid_x = self.input[0].sigmoid() + grad_input = gradient * sigmoid_x * (1 - sigmoid_x) + + return [grad_input] + diff --git a/build/lib/norch/csrc/cpu.cpp b/build/lib/norch/csrc/cpu.cpp new file mode 100644 index 0000000..63e1615 --- /dev/null +++ b/build/lib/norch/csrc/cpu.cpp @@ -0,0 +1,443 @@ +#include "tensor.h" +#include +#include +#include +#include + +void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + for (int i = 0; i < tensor1->size; i++) { + result_data[i] = tensor1->data[i] + tensor2->data[i]; + } +} + +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 + 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 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1; + stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1; + } + + // Perform element-wise addition with broadcasting + 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--) { + 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++) { + result_data[i] = tensor1->data[i] - tensor2->data[i]; + } +} + +void sub_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 + 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 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1; + stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1; + } + + // Perform element-wise addition with broadcasting + 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--) { + 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++) { + result_data[i] = tensor1->data[i] * tensor2->data[i]; + } +} + +void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data) { + + for (int i = 0; i < tensor->size; i++) { + result_data[i] = scalar * tensor->data[i]; + } +} + +void scalar_div_tensor_cpu(float scalar, Tensor* tensor, float* result_data) { + + for (int i = 0; i < tensor->size; i++) { + result_data[i] = scalar / tensor->data[i]; + } +} + +void tensor_div_scalar_cpu(Tensor* tensor, float scalar, float* result_data) { + + for (int i = 0; i < tensor->size; i++) { + result_data[i] = tensor->data[i] / scalar; + } +} + +void tensor_div_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + for (int i = 0; i < tensor1->size; i++) { + result_data[i] = tensor1->data[i] / tensor2->data[i]; + } +} + +void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + for (int i = 0; i < tensor1->shape[0]; i++) { + for (int j = 0; j < tensor2->shape[1]; j++) { + float sum = 0.0; + for (int k = 0; k < tensor1->shape[1]; k++) { + sum += tensor1->data[i * tensor1->shape[1] + k] * tensor2->data[k * tensor2->shape[1] + j]; + } + result_data[i * tensor2->shape[1] + j] = sum; + } + } +} + + +void broadcasted_batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int result_data_stride = tensor1->shape[0] * tensor2->shape[2]; + + for (int batch = 0; batch < tensor2->shape[0]; batch++) { + + for (int i = 0; i < tensor1->shape[0]; i++) { + for (int j = 0; j < tensor2->shape[2]; j++) { + float sum = 0.0; + for (int k = 0; k < tensor1->shape[1]; k++) { + sum += tensor1->data[i * tensor1->shape[1] + k] * tensor2->data[batch*tensor2->strides[0] + (k * tensor2->shape[2] + j)]; + } + result_data[(batch * result_data_stride) + (i * tensor2->shape[2] + j)] = sum; + } + } + } +} + +void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { + int result_data_stride = tensor1->shape[1] * tensor2->shape[2]; + + for (int batch = 0; batch < tensor2->shape[0]; batch++) { + + for (int i = 0; i < tensor1->shape[1]; i++) { + for (int j = 0; j < tensor2->shape[2]; j++) { + float sum = 0.0; + for (int k = 0; k < tensor1->shape[2]; k++) { + sum += tensor1->data[(batch * tensor1->strides[0]) + i * tensor1->shape[2] + k] * tensor2->data[batch*tensor2->strides[0] + (k * tensor2->shape[2] + j)]; + } + result_data[(batch * result_data_stride) + (i * tensor2->shape[2] + j)] = sum; + } + } + } +} + +void scalar_pow_tensor_cpu(float base, Tensor* tensor, float* result_data) { + + for (int i = 0; i < tensor->size; i++) { + result_data[i] = powf(base, tensor->data[i]); + } +} + +void tensor_pow_scalar_cpu(Tensor* tensor, float exponent, float* result_data) { + + for (int i = 0; i < tensor->size; i++) { + result_data[i] = powf(tensor->data[i], exponent); + } +} + +void log_tensor_cpu(Tensor* tensor, float* result_data) { + for (int i = 0; i < tensor->size; i++) { + result_data[i] = logf(tensor->data[i]); + } +} + +void sum_tensor_cpu(Tensor* tensor, float* result_data, int size, int* result_shape, 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 axis_stride = tensor->strides[axis]; + + for (int i = 0; i < tensor->shape[axis]; i++) { + for (int j = 0; j < 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]; + } + } + } +} + +void max_tensor_cpu(Tensor* tensor, float* result_data, int size, int* result_shape, int axis) { + if (axis == -1) { + float max_value = -INFINITY; + for (int i = 0; i < tensor->size; i++) { + max_value = fmax(max_value, tensor->data[i]); + } + *result_data = max_value; + } else { + for (int i = 0; i < size; i++) { + result_data[i] = -INFINITY; + } + if (axis < 0 || axis >= tensor->ndim) { + printf("Invalid axis"); + return; + } + + int axis_stride = tensor->strides[axis]; + + for (int i = 0; i < tensor->shape[axis]; i++) { + for (int j = 0; j < 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] = fmax(result_data[j], tensor->data[index + i * axis_stride]); + } + } + } +} + +void min_tensor_cpu(Tensor* tensor, float* result_data, int size, int* result_shape, int axis) { + if (axis == -1) { + float min_value = INFINITY; + for (int i = 0; i < tensor->size; i++) { + min_value = fmin(min_value, tensor->data[i]); + } + *result_data = min_value; + } else { + for (int i = 0; i < size; i++) { + result_data[i] = INFINITY; + } + if (axis < 0 || axis >= tensor->ndim) { + printf("Invalid axis"); + return; + } + + int axis_stride = tensor->strides[axis]; + + for (int i = 0; i < tensor->shape[axis]; i++) { + for (int j = 0; j < 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] = fmin(result_data[j], tensor->data[index + i * axis_stride]); + } + } + } +} + +void equal_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + for (int i = 0; i < tensor1->size; i++) { + result_data[i] = (tensor1->data[i] == tensor2->data[i]) ? 1.0f : 0.0f; + } +} + +void equal_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 + 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 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1; + stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1; + } + + // Perform element-wise equal with broadcasting + 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--) { + 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]) ? 1.0f : 0.0f; + } + + // Free strides + free(strides1); + free(strides2); +} + + +void ones_like_tensor_cpu(Tensor* tensor, float* result_data) { + + for (int i = 0; i < tensor->size; i++) { + result_data[i] = 1.0; + } +} + +void zeros_like_tensor_cpu(Tensor* tensor, float* result_data) { + + for (int i = 0; i < tensor->size; i++) { + result_data[i] = 0.0; + } +} + +void transpose_1D_tensor_cpu(Tensor* tensor, float* result_data) { + + for (int i = 0; i < tensor->shape[0]; i++) { + result_data[i] = tensor->data[i]; + } +} + +void transpose_2D_tensor_cpu(Tensor* tensor, float* result_data) { + int rows = tensor->shape[0]; + int cols = tensor->shape[1]; + + for (int i = 0; i < rows; i++) { + for (int j = 0; j < cols; j++) { + result_data[j * rows + i] = tensor->data[i * cols + j]; + } + } +} + +void transpose_3D_tensor_cpu(Tensor* tensor, float* result_data) { + int batch = tensor->shape[0]; + int rows = tensor->shape[1]; + int cols = tensor->shape[2]; + + for (int i = 0; i < batch; i++) { + for (int j = 0; j < rows; j++) { + for (int k = 0; k < cols; k++) { + result_data[k * rows * batch + j * batch + i] = tensor->data[i * rows * cols + j * cols + k]; + } + } + } +} + +void assign_tensor_cpu(Tensor* tensor, float* result_data) { + + for (int i = 0; i < tensor->size; i++) { + result_data[i] = tensor->data[i]; + } +} + +void make_contiguous_tensor_cpu(Tensor* tensor, float* result_data, int* new_strides) { + + for (int i = 0; i < tensor->size; i++) { + int index = 0; + int offset = i; + for (int j = 0; j < tensor->ndim; j++) { + index += (offset / new_strides[j]) * tensor->strides[j]; + offset %= new_strides[j]; + } + result_data[i] = tensor->data[index]; + } + + // Free old data and update tensor properties + free(tensor->data); + free(tensor->strides); + tensor->data = result_data; + tensor->strides = new_strides; +} + +void sin_tensor_cpu(Tensor* tensor, float* result_data) { + for (int i = 0; i < tensor->size; i++) { + result_data[i] = sinf(tensor->data[i]); + } +} + +void sigmoid_tensor_cpu(Tensor* tensor, float* result_data) { + for (int i = 0; i < tensor->size; i++) { + // avoid overflow + if (tensor->data[i] >= 0) { + + float z = expf(-tensor->data[i]); + result_data[i] = 1 / (1 + z); + + } else { + + float z = expf(tensor->data[i]); + result_data[i] = z / (1 + z); + } + } +} + +void cos_tensor_cpu(Tensor* tensor, float* result_data) { + for (int i = 0; i < tensor->size; i++) { + result_data[i] = cosf(tensor->data[i]); + } +} diff --git a/build/lib/norch/csrc/cpu.h b/build/lib/norch/csrc/cpu.h new file mode 100644 index 0000000..9a27a26 --- /dev/null +++ b/build/lib/norch/csrc/cpu.h @@ -0,0 +1,38 @@ +#ifndef CPU_H +#define CPU_H + +#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, int broadcasted_size); +void sum_tensor_cpu(Tensor* tensor, float* result_data, int size, int* shape, int axis); +void max_tensor_cpu(Tensor* tensor, float* result_data, int size, int* result_shape, int axis); +void min_tensor_cpu(Tensor* tensor, float* result_data, int size, int* result_shape, 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, int broadcasted_size); +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); +void tensor_div_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); +void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); +void broadcasted_batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); +void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); +void scalar_pow_tensor_cpu(float base, Tensor* tensor, float* result_data); +void tensor_pow_scalar_cpu(Tensor* tensor, float exponent, float* result_data); +void log_tensor_cpu(Tensor* tensor, float* result_data); +void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data); +void equal_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data); +void equal_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size); +void ones_like_tensor_cpu(Tensor* tensor, float* result_data); +void zeros_like_tensor_cpu(Tensor* tensor, float* result_data); +void transpose_1D_tensor_cpu(Tensor* tensor, float* result_data); +void transpose_2D_tensor_cpu(Tensor* tensor, float* result_data); +void transpose_3D_tensor_cpu(Tensor* tensor, float* result_data); +void transpose_axes_cpu(Tensor* tensor, float* result_data, int axis1, int axis2, int* new_shape); +void assign_tensor_cpu(Tensor* tensor, float* result_data); +void make_contiguous_tensor_cpu(Tensor* tensor, float* result_data, int* new_strides); +void sin_tensor_cpu(Tensor* tensor, float* result_data); +void cos_tensor_cpu(Tensor* tensor, float* result_data); +void sigmoid_tensor_cpu(Tensor* tensor, float* result_data); + +#endif /* CPU_H */ diff --git a/build/lib/norch/csrc/cuda.cu b/build/lib/norch/csrc/cuda.cu new file mode 100644 index 0000000..e39c20c --- /dev/null +++ b/build/lib/norch/csrc/cuda.cu @@ -0,0 +1,1248 @@ +#include "tensor.h" +#include +#include +#include +#include + +#define THREADS_PER_BLOCK 128 +#define TILE_SIZE 32 + +__host__ void cpu_to_cuda(Tensor* tensor, int device_id) { + + int deviceCount; + cudaGetDeviceCount(&deviceCount); + if (device_id >= deviceCount) { + fprintf(stderr, "Could not send tensor to device %d, only %d devices available\n", device_id, deviceCount); + exit(1); + } + + cudaSetDevice(device_id); + + float* data_tmp; + cudaMalloc((void **)&data_tmp, tensor->size * sizeof(float)); + cudaMemcpy(data_tmp, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice); + + tensor->data = data_tmp; + + const char* device_str = "cuda"; + tensor->device = (char*)malloc(strlen(device_str) + 1); + strcpy(tensor->device, device_str); +} + +__host__ void cuda_to_cpu(Tensor* tensor) { + float* data_tmp = (float*)malloc(tensor->size * sizeof(float)); + + cudaMemcpy(data_tmp, tensor->data, tensor->size * sizeof(float), cudaMemcpyDeviceToHost); + cudaFree(tensor->data); + + tensor->data = data_tmp; + + const char* device_str = "cpu"; + tensor->device = (char*)malloc(strlen(device_str) + 1); + strcpy(tensor->device, device_str); + + printf("Successfully sent tensor to: %s\n", tensor->device); +} + +__global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = data1[i] + data2[i]; + } +} + +__host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + add_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, tensor1->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= size) return; + + 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] = data1[index1] + data2[index2]; +} + +__host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) { + int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim; + + 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 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1; + stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1; + } + + int* d_broadcasted_shape; + int* d_strides1; + int* d_strides2; + + cudaMalloc((void**)&d_broadcasted_shape, max_ndim * sizeof(int)); + cudaMemcpy(d_broadcasted_shape, broadcasted_shape, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + cudaMalloc((void**)&d_strides1, max_ndim * sizeof(int)); + cudaMemcpy(d_strides1, strides1, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + cudaMalloc((void**)&d_strides2, max_ndim * sizeof(int)); + cudaMemcpy(d_strides2, strides2, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + int number_of_blocks = (broadcasted_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + add_broadcasted_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, d_broadcasted_shape, d_strides1, d_strides2, max_ndim, broadcasted_size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + cudaFree(d_broadcasted_shape); +} + +__global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size) { + __shared__ float partial_sum[THREADS_PER_BLOCK * sizeof(float)]; + + int tid = threadIdx.x; + int i = blockIdx.x * blockDim.x + threadIdx.x; + + partial_sum[tid] = (i < size) ? data[i] : 0; + + __syncthreads(); + + // Perform block-wise reduction + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + partial_sum[tid] += partial_sum[tid + s]; + } + __syncthreads(); + } + + // Write block sum to global memory + if (tid == 0) { + result_data[blockIdx.x] = partial_sum[0]; + } +} + +__global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + if (tid < result_size) { + for (int i = 0; i < shape[axis]; i++) { + int index = 0; + int remainder = tid; + for (int k = ndim - 2; k >= 0; k--) { + index += (remainder % shape[k < axis ? k : k + 1]) * strides[k < axis ? k : k + 1]; + remainder /= shape[k < axis ? k : k + 1]; + } + index += i * axis_stride; + + atomicAdd(&result_data[tid], data[index]); + } + } +} + + +__host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) { + + if (axis == -1) { + cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice); + + int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + + // First-level reduction + sum_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + // If necessary, perform multiple levels of reduction + while (num_blocks > 1) { + int num_blocks_next = (num_blocks + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + sum_tensor_cuda_kernel<<>>(result_data, result_data, num_blocks); + num_blocks = num_blocks_next; + } + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + } else { + int axis_stride = tensor->strides[axis]; + + // Calculate the size of the resulting tensor + int result_size = 1; + for (int i = 0; i < tensor->ndim; i++) { + if (i != axis) { + result_size *= tensor->shape[i]; + } + } + + // Allocate memory for strides and shape on the device + int* d_strides; + int* d_shape; + cudaMalloc(&d_strides, tensor->ndim * sizeof(int)); + cudaMalloc(&d_shape, tensor->ndim * sizeof(int)); + cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + + // Initialize result_data to 0 + cudaMemset(result_data, 0, result_size * sizeof(float)); + + int num_threads = result_size; + int num_blocks = (num_threads + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + sum_tensor_cuda_kernel_axis<<>>(tensor->data, result_data, d_strides, d_shape, axis, tensor->ndim, axis_stride, tensor->size, result_size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + // Free allocated memory + cudaFree(d_strides); + cudaFree(d_shape); + } +} + +__global__ void max_tensor_cuda_kernel(float* data, float* result_data, int size) { + __shared__ float partial_max[THREADS_PER_BLOCK * sizeof(float)]; + + int tid = threadIdx.x; + int i = blockIdx.x * blockDim.x + threadIdx.x; + + partial_max[tid] = (i < size) ? data[i] : -FLT_MAX; + + __syncthreads(); + + // Perform block-wise reduction + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + partial_max[tid] = fmax(partial_max[tid], partial_max[tid + s]); + } + __syncthreads(); + } + + // Write block sum to global memory + if (tid == 0) { + result_data[blockIdx.x] = partial_max[0]; + } +} + +__device__ float atomicMaxFloat(float* address, float val) { + int* address_as_int = (int*)address; + int old = *address_as_int, assumed; + + do { + assumed = old; + old = atomicCAS(address_as_int, assumed, + __float_as_int(fmaxf(val, __int_as_float(assumed)))); + } while (assumed != old); + + return __int_as_float(old); +} + +__global__ void max_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + if (tid < result_size) { + for (int i = 0; i < shape[axis]; i++) { + int index = 0; + int remainder = tid; + for (int k = ndim - 2; k >= 0; k--) { + index += (remainder % shape[k < axis ? k : k + 1]) * strides[k < axis ? k : k + 1]; + remainder /= shape[k < axis ? k : k + 1]; + } + index += i * axis_stride; + + atomicMaxFloat(&result_data[tid], data[index]); + } + } +} + +__host__ void max_tensor_cuda(Tensor* tensor, float* result_data, int axis) { + + if (axis == -1) { + cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice); + + int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + + // First-level reduction + max_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + // If necessary, perform multiple levels of reduction + while (num_blocks > 1) { + int num_blocks_next = (num_blocks + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + max_tensor_cuda_kernel<<>>(result_data, result_data, num_blocks); + num_blocks = num_blocks_next; + } + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + } else { + int axis_stride = tensor->strides[axis]; + + // Calculate the size of the resulting tensor + int result_size = 1; + for (int i = 0; i < tensor->ndim; i++) { + if (i != axis) { + result_size *= tensor->shape[i]; + } + } + + // Allocate memory for strides and shape on the device + int* d_strides; + int* d_shape; + cudaMalloc(&d_strides, tensor->ndim * sizeof(int)); + cudaMalloc(&d_shape, tensor->ndim * sizeof(int)); + cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + + // Initialize result_data to 0 + float neg_inf = -FLT_MAX; + cudaMemset(result_data, *reinterpret_cast(&neg_inf), result_size * sizeof(float)); + + int num_threads = result_size; + int num_blocks = (num_threads + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + max_tensor_cuda_kernel_axis<<>>(tensor->data, result_data, d_strides, d_shape, axis, tensor->ndim, axis_stride, tensor->size, result_size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + // Free allocated memory + cudaFree(d_strides); + cudaFree(d_shape); + } +} + +__global__ void min_tensor_cuda_kernel(float* data, float* result_data, int size) { + __shared__ float partial_min[THREADS_PER_BLOCK * sizeof(float)]; + + int tid = threadIdx.x; + int i = blockIdx.x * blockDim.x + threadIdx.x; + + partial_min[tid] = (i < size) ? data[i] : FLT_MAX; + + __syncthreads(); + + // Perform block-wise reduction + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + partial_min[tid] = fmin(partial_min[tid], partial_min[tid + s]); + } + __syncthreads(); + } + + // Write block sum to global memory + if (tid == 0) { + result_data[blockIdx.x] = partial_min[0]; + } +} + +__device__ float atomicMinFloat(float* address, float val) { + int* address_as_int = (int*)address; + int old = *address_as_int, assumed; + + do { + assumed = old; + old = atomicCAS(address_as_int, assumed, + __float_as_int(fminf(val, __int_as_float(assumed)))); + } while (assumed != old); + + return __int_as_float(old); +} + +__global__ void min_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + if (tid < result_size) { + for (int i = 0; i < shape[axis]; i++) { + int index = 0; + int remainder = tid; + for (int k = ndim - 2; k >= 0; k--) { + index += (remainder % shape[k < axis ? k : k + 1]) * strides[k < axis ? k : k + 1]; + remainder /= shape[k < axis ? k : k + 1]; + } + index += i * axis_stride; + + atomicMinFloat(&result_data[tid], data[index]); + } + } +} + +__host__ void min_tensor_cuda(Tensor* tensor, float* result_data, int axis) { + + if (axis == -1) { + cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice); + + int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + + // First-level reduction + min_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + // If necessary, perform multiple levels of reduction + while (num_blocks > 1) { + int num_blocks_next = (num_blocks + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + min_tensor_cuda_kernel<<>>(result_data, result_data, num_blocks); + num_blocks = num_blocks_next; + } + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + } else { + int axis_stride = tensor->strides[axis]; + + // Calculate the size of the resulting tensor + int result_size = 1; + for (int i = 0; i < tensor->ndim; i++) { + if (i != axis) { + result_size *= tensor->shape[i]; + } + } + + // Allocate memory for strides and shape on the device + int* d_strides; + int* d_shape; + cudaMalloc(&d_strides, tensor->ndim * sizeof(int)); + cudaMalloc(&d_shape, tensor->ndim * sizeof(int)); + cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + + // Initialize result_data to 0 + float inf = FLT_MAX; + cudaMemset(result_data, *reinterpret_cast(&inf), result_size * sizeof(float)); + + int num_threads = result_size; + int num_blocks = (num_threads + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + min_tensor_cuda_kernel_axis<<>>(tensor->data, result_data, d_strides, d_shape, axis, tensor->ndim, axis_stride, tensor->size, result_size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + // Free allocated memory + cudaFree(d_strides); + cudaFree(d_shape); + } +} + + + +__global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = data1[i] - data2[i]; + } +} + +__host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + sub_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, tensor1->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void sub_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= size) return; + + 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] = data1[index1] - data2[index2]; +} + +__host__ void sub_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) { + int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim; + + 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 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1; + stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1; + } + + int* d_broadcasted_shape; + int* d_strides1; + int* d_strides2; + + cudaMalloc((void**)&d_broadcasted_shape, max_ndim * sizeof(int)); + cudaMemcpy(d_broadcasted_shape, broadcasted_shape, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + cudaMalloc((void**)&d_strides1, max_ndim * sizeof(int)); + cudaMemcpy(d_strides1, strides1, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + cudaMalloc((void**)&d_strides2, max_ndim * sizeof(int)); + cudaMemcpy(d_strides2, strides2, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + int number_of_blocks = (broadcasted_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + sub_broadcasted_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, d_broadcasted_shape, d_strides1, d_strides2, max_ndim, broadcasted_size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + cudaFree(d_broadcasted_shape); +} + +__global__ void elementwise_mul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = data1[i] * data2[i]; + } +} + +__host__ void elementwise_mul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + elementwise_mul_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, tensor1->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void scalar_mul_tensor_cuda_kernel(float* data, float scalar, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = scalar * data[i]; + } +} + +__host__ void scalar_mul_tensor_cuda(Tensor* tensor, float scalar, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + scalar_mul_tensor_cuda_kernel<<>>(tensor->data, scalar, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void scalar_div_tensor_cuda_kernel(float scalar, float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = scalar / data[i]; + } +} + +__host__ void scalar_div_tensor_cuda(float scalar, Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + scalar_div_tensor_cuda_kernel<<>>(scalar, tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void tensor_div_scalar_cuda_kernel(float* data, float scalar, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = data[i] / scalar; + } +} + +__host__ void tensor_div_scalar_cuda(Tensor* tensor, float scalar, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + tensor_div_scalar_cuda_kernel<<>>(tensor->data, scalar, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void tensor_div_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = data1[i] / data2[i]; + } +} + +__host__ void tensor_div_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + tensor_div_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, tensor1->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +/*__global__ void matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int rows1, int cols1, int cols2) { + + int row = blockIdx.y * blockDim.y + threadIdx.y; + int col = blockIdx.x * blockDim.x + threadIdx.x; + + if (row < rows1 && col < cols2) { + float sum = 0.0; + for (int k = 0; k < cols1; k++) { + sum += data1[row * cols1 + k] * data2[k * cols2 + col]; + } + result_data[row * cols2 + col] = sum; + } + +}*/ + +__global__ void matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int rows1, int cols1, int cols2) { + + // Shared memory for tiles + __shared__ float tile1[TILE_SIZE][TILE_SIZE]; + __shared__ float tile2[TILE_SIZE][TILE_SIZE]; + + // Thread indices + int tx = threadIdx.x; + int ty = threadIdx.y; + + // Output position + int row = blockIdx.y * blockDim.y + threadIdx.y; + int col = blockIdx.x * blockDim.x + threadIdx.x; + + float sum = 0.0; + + // Iterate over tiles + for (int i = 0; i < (cols1 + TILE_SIZE - 1) / TILE_SIZE; ++i) { + + // Load tiles into shared memory + if (row < rows1 && i * TILE_SIZE + tx < cols1) + tile1[ty][tx] = data1[row * cols1 + i * TILE_SIZE + tx]; + else + tile1[ty][tx] = 0.0; + + if (col < cols2 && i * TILE_SIZE + ty < cols1) + tile2[ty][tx] = data2[(i * TILE_SIZE + ty) * cols2 + col]; + else + tile2[ty][tx] = 0.0; + + // Synchronize threads + __syncthreads(); + + // Accumulate sum + for (int k = 0; k < TILE_SIZE; ++k) + sum += tile1[ty][k] * tile2[k][tx]; + + // Synchronize threads + __syncthreads(); + } + + // Write result to global memory + if (row < rows1 && col < cols2) + result_data[row * cols2 + col] = sum; +} + + +__host__ void matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int rows1 = tensor1->shape[0]; + int cols1 = tensor1->shape[1]; + int cols2 = tensor2->shape[1]; + + dim3 threadsPerBlock(16, 16); + dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y); + matmul_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, rows1, cols1, cols2); + + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void batched_matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int batch_size, int rows1, int cols1, int cols2) { + int batch = blockIdx.z; + int row = blockIdx.y * blockDim.y + threadIdx.y; + int col = blockIdx.x * blockDim.x + threadIdx.x; + + if (row < rows1 && col < cols2) { + float sum = 0.0f; + for (int k = 0; k < cols1; ++k) { + sum += data1[batch * rows1 * cols1 + row * cols1 + k] * + data2[batch * cols1 * cols2 + k * cols2 + col]; + } + result_data[batch * rows1 * cols2 + row * cols2 + col] = sum; + } +} + +__host__ void batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int batch_size = tensor2->shape[0]; + int rows1 = tensor1->shape[1]; + int cols1 = tensor1->shape[2]; + int cols2 = tensor2->shape[2]; + + dim3 threadsPerBlock(16, 16); + dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y, batch_size); + batched_matmul_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, batch_size, rows1, cols1, cols2); + + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void broadcasted_batched_matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int batch_size, int rows1, int cols1, int cols2) { + int batch = blockIdx.z; + int row = blockIdx.y * blockDim.y + threadIdx.y; + int col = blockIdx.x * blockDim.x + threadIdx.x; + + if (row < rows1 && col < cols2) { + float sum = 0.0f; + for (int k = 0; k < cols1; ++k) { + sum += data1[row * cols1 + k] * + data2[batch * cols1 * cols2 + k * cols2 + col]; + } + result_data[batch * rows1 * cols2 + row * cols2 + col] = sum; + } +} + +__host__ void broadcasted_batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int batch_size = tensor2->shape[0]; + int rows1 = tensor1->shape[0]; + int cols1 = tensor1->shape[1]; + int cols2 = tensor2->shape[2]; + + dim3 threadsPerBlock(16, 16); + dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y, batch_size); + broadcasted_batched_matmul_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, batch_size, rows1, cols1, cols2); + + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void tensor_pow_scalar_cuda_kernel(float* data, float exponent, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = powf(data[i], exponent); + } +} + +__host__ void tensor_pow_scalar_cuda(Tensor* tensor, float exponent, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + tensor_pow_scalar_cuda_kernel<<>>(tensor->data, exponent, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void scalar_pow_tensor_cuda_kernel(float base, float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = powf(base, data[i]); + } +} + +__host__ void scalar_pow_tensor_cuda(float base, Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + scalar_pow_tensor_cuda_kernel<<>>(base, tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void log_tensor_cuda_kernel(float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = logf(data[i]); + } +} + +__host__ void log_tensor_cuda(Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + log_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void equal_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = (data1[i] == data2[i]) ? 1.0f : 0.0f; + } +} + +__host__ void equal_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data) { + + int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + equal_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, tensor1->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void equal_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= size) return; + + 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] = (data1[index1] == data2[index2]) ? 1.0f : 0.0f; +} + +__host__ void equal_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) { + int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim; + + 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 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1; + stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1; + } + + int* d_broadcasted_shape; + int* d_strides1; + int* d_strides2; + + cudaMalloc((void**)&d_broadcasted_shape, max_ndim * sizeof(int)); + cudaMemcpy(d_broadcasted_shape, broadcasted_shape, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + cudaMalloc((void**)&d_strides1, max_ndim * sizeof(int)); + cudaMemcpy(d_strides1, strides1, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + cudaMalloc((void**)&d_strides2, max_ndim * sizeof(int)); + cudaMemcpy(d_strides2, strides2, max_ndim * sizeof(int), cudaMemcpyHostToDevice); + + int number_of_blocks = (broadcasted_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + equal_broadcasted_tensor_cuda_kernel<<>>(tensor1->data, tensor2->data, result_data, d_broadcasted_shape, d_strides1, d_strides2, max_ndim, broadcasted_size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + cudaFree(d_broadcasted_shape); +} + + +__global__ void ones_like_tensor_cuda_kernel(float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = 1.0; + } +} + +__host__ void ones_like_tensor_cuda(Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + ones_like_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void zeros_like_tensor_cuda_kernel(float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = 0.0; + } +} + +__host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + zeros_like_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + + +__global__ void transpose_1D_tensor_cuda_kernel(float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = data[i]; + } +} + +__host__ void transpose_1D_tensor_cuda(Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + transpose_1D_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void transpose_2D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + int j = blockIdx.y * blockDim.y + threadIdx.y; + + if (i < rows && j < cols) { + result_data[j * rows + i] = data[i * cols + j]; + } +} + +__host__ void transpose_2D_tensor_cuda(Tensor* tensor, float* result_data) { + + int rows = tensor->shape[0]; + int cols = tensor->shape[1]; + + dim3 threadsPerBlock(16, 16); + dim3 number_of_blocks((rows + threadsPerBlock.x - 1) / threadsPerBlock.x, (cols + threadsPerBlock.y - 1) / threadsPerBlock.y); + transpose_2D_tensor_cuda_kernel<<>>(tensor->data, result_data, rows, cols); + + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void transpose_3D_tensor_cuda_kernel(float* data, float* result_data, int batch, int rows, int cols) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + int j = blockIdx.y * blockDim.y + threadIdx.y; + int k = blockIdx.z * blockDim.z + threadIdx.z; + + if (i < batch && j < rows && k < cols) { + result_data[k * rows * batch + j * batch + i] = data[i * rows * cols + j * cols + k]; + } +} + +__host__ void transpose_3D_tensor_cuda(Tensor* tensor, float* result_data) { + + int batch = tensor->shape[0]; + int rows = tensor->shape[1]; + int cols = tensor->shape[2]; + + dim3 threadsPerBlock(8, 8, 8); + dim3 number_of_blocks((batch + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows + threadsPerBlock.y - 1) / threadsPerBlock.y, (cols + threadsPerBlock.z - 1) / threadsPerBlock.z); + transpose_3D_tensor_cuda_kernel<<>>(tensor->data, result_data, batch, rows, cols); + + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + + +__global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = data[i]; + } +} + +__host__ void assign_tensor_cuda(Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + assign_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void make_contiguous_tensor_cuda_kernel(float* data, float* result_data, int ndim, int size, int* strides, int* new_strides) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + int index = 0; + int offset = i; + for (int j = 0; j < ndim; j++) { + index += (offset / new_strides[j]) * strides[j]; + offset %= new_strides[j]; + } + result_data[i] = data[index]; + } +} + +__host__ void make_contiguous_tensor_cuda(Tensor* tensor, float* result_data, int* new_strides) { + + int* d_strides; + cudaMalloc((void **)&d_strides, tensor->ndim * sizeof(int)); + cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + + int* d_new_strides; + cudaMalloc((void **)&d_new_strides, tensor->ndim * sizeof(int)); + cudaMemcpy(d_new_strides, new_strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice); + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + make_contiguous_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->ndim, tensor->size, d_strides, d_new_strides); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); + + // Free old data and update tensor properties + cudaFree(tensor->data); + free(tensor->strides); + tensor->data = result_data; + tensor->strides = new_strides; +} + +__global__ void sin_tensor_cuda_kernel(float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = sinf(data[i]); + } +} + +__host__ void sin_tensor_cuda(Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + sin_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void cos_tensor_cuda_kernel(float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) { + result_data[i] = cosf(data[i]); + } +} + +__host__ void cos_tensor_cuda(Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + cos_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + +__global__ void sigmoid_tensor_cuda_kernel(float* data, float* result_data, int size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + + if (i < size) { + // avoid overflow + if (data[i] >= 0) { + + float z = expf(-data[i]); + result_data[i] = 1 / (1 + z); + + } else { + + float z = expf(data[i]); + result_data[i] = z / (1 + z); + } + } +} + +__host__ void sigmoid_tensor_cuda(Tensor* tensor, float* result_data) { + + int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + sigmoid_tensor_cuda_kernel<<>>(tensor->data, result_data, tensor->size); + + cudaError_t error = cudaGetLastError(); + if (error != cudaSuccess) { + printf("CUDA error: %s\n", cudaGetErrorString(error)); + exit(-1); + } + + cudaDeviceSynchronize(); +} + + + + diff --git a/build/lib/norch/csrc/cuda.h b/build/lib/norch/csrc/cuda.h new file mode 100644 index 0000000..702dea1 --- /dev/null +++ b/build/lib/norch/csrc/cuda.h @@ -0,0 +1,104 @@ +#ifndef CUDA_KERNEL_H_ +#define CUDA_KERNEL_H_ + + + __host__ void cpu_to_cuda(Tensor* tensor, int device_id); + __host__ void cuda_to_cpu(Tensor* tensor); + + __global__ void add_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size); + __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + + __global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size); + __host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size); + + __global__ void sub_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size); + __host__ void sub_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size); + + __global__ void sum_tensor_cuda_kernel(float* data, float* result_data); + __global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size); + __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis); + + __global__ void max_tensor_cuda_kernel(float* data, float* result_data); + __global__ void max_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size); + __host__ void max_tensor_cuda(Tensor* tensor, float* result_data, int axis); + + __global__ void min_tensor_cuda_kernel(float* data, float* result_data); + __global__ void min_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size); + __host__ void min_tensor_cuda(Tensor* tensor, float* result_data, int axis); + + __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size); + __host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + + __global__ void elementwise_mul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size); + __host__ void elementwise_mul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + + __global__ void scalar_mul_tensor_cuda_kernel(float* data, float scalar, float* result_data, int size); + __host__ void scalar_mul_tensor_cuda(Tensor* tensor, float scalar, float* result_data); + + __global__ void scalar_div_tensor_cuda_kernel(float scalar, float* data, float* result_data, int size); + __host__ void scalar_div_tensor_cuda(float scalar, Tensor* tensor, float* result_data); + + __global__ void tensor_div_scalar_cuda_kernel(float* data, float scalar, float* result_data, int size); + __host__ void tensor_div_scalar_cuda(Tensor* tensor, float scalar, float* result_data); + + __global__ void tensor_div_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size); + __host__ void tensor_div_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + + __global__ void matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int rows1, int cols1, int cols2); + __host__ void matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + + __global__ void batched_matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int batch_size, int rows1, int cols1, int cols2); + __host__ void batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + + __global__ void broadcasted_batched_matmul_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int batch_size, int rows1, int cols1, int cols2); + __host__ void broadcasted_batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + + __global__ void tensor_pow_scalar_cuda_kernel(float* data, float exponent, float* result_data, int size); + __host__ void tensor_pow_scalar_cuda(Tensor* tensor, float exponent, float* result_data); + + __global__ void scalar_pow_tensor_cuda_kernel(float base, float* data, float* result_data, int size); + __host__ void scalar_pow_tensor_cuda(float base, Tensor* tensor, float* result_data); + + __global__ void log_tensor_cuda_kernel(float* data, float* result_data, int size); + __host__ void log_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void equal_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size); + __host__ void equal_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data); + + __global__ void equal_broadcasted_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int* broadcasted_shape, int* strides1, int*strides2, int max_ndim, int size); + __host__ void equal_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size); + + __global__ void ones_like_tensor_cuda_kernel(float* data, float* result_data, int size); + __host__ void ones_like_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void zeros_like_tensor_cuda_kernel(float* data, float* result_data, int size); + __host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void transpose_1D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols); + __host__ void transpose_1D_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void transpose_2D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols); + __host__ void transpose_2D_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void transpose_3D_tensor_cuda_kernel(float* data, float* result_data, int batch, int rows, int cols); + __host__ void transpose_3D_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size); + __host__ void assign_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void make_contiguous_tensor_cuda_kernel(float* data, float* result_data, int ndim, int size, int* strides, int* new_strides); + __host__ void make_contiguous_tensor_cuda(Tensor* tensor, float* result_data, int* new_strides); + + __global__ void sin_tensor_cuda_kernel(float* data, float* result_data, int size); + __host__ void sin_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void cos_tensor_cuda_kernel(float* data, float* result_data, int size); + __host__ void cos_tensor_cuda(Tensor* tensor, float* result_data); + + __global__ void sigmoid_tensor_cuda_kernel(float* data, float* result_data, int size); + __host__ void sigmoid_tensor_cuda(Tensor* tensor, float* result_data); + + + + +#endif /* CUDA_KERNEL_H_ */ diff --git a/build/lib/norch/csrc/distributed.cpp b/build/lib/norch/csrc/distributed.cpp new file mode 100644 index 0000000..8124ec0 --- /dev/null +++ b/build/lib/norch/csrc/distributed.cpp @@ -0,0 +1,65 @@ +#include +#include +#include +#include "distributed.h" +#include "tensor.h" +#include "cuda.h" +#include +#include + +int rank; +int world_size; +ncclComm_t nccl_comm; + +extern "C" { + +void init_process_group(int env_rank, int env_world_size) { + + rank = env_rank; + world_size = env_world_size; + + // init MPI + MPI_CHECK(MPI_Init(NULL, NULL)); + + + ncclUniqueId nccl_id; + + if (rank == 0) { + ncclGetUniqueId(&nccl_id); + } + + // send nccl unique id to all processes + MPI_CHECK(MPI_Bcast((void *)&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD)); + + // init NCCL communication group + NCCL_CHECK(ncclCommInitRank(&nccl_comm, world_size, nccl_id, rank)); + +} + +void broadcast_tensor(Tensor* tensor) { + + cudaStream_t stream; + cudaStreamCreate(&stream); + + NCCL_CHECK(ncclBroadcast(tensor->data, tensor->data, tensor->size * sizeof(float), ncclFloat, 0, nccl_comm, stream)); + cudaStreamSynchronize(stream); + cudaStreamDestroy(stream); +} + +void allreduce_sum_tensor(Tensor* tensor) { + cudaStream_t stream; + cudaStreamCreate(&stream); + + // Perform NCCL AllReduce operation to calculate the sum of all tensors across all processes + NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream)); + + cudaStreamSynchronize(stream); + cudaStreamDestroy(stream); +} + +void end_process_group() { + MPI_CHECK(MPI_Finalize()); + NCCL_CHECK(ncclCommDestroy(nccl_comm)); +} + +} diff --git a/build/lib/norch/csrc/distributed.h b/build/lib/norch/csrc/distributed.h new file mode 100644 index 0000000..1040fc0 --- /dev/null +++ b/build/lib/norch/csrc/distributed.h @@ -0,0 +1,28 @@ +#ifndef DISTRIBUTED_H +#define DISTRIBUTED_H + +#define MPI_CHECK(cmd) do { \ + int e = cmd; \ + if( e != MPI_SUCCESS ) { \ + printf("Failed: MPI error %s:%d '%d'\n", \ + __FILE__,__LINE__, e); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + +#define NCCL_CHECK(cmd) do { \ + ncclResult_t r = cmd; \ + if (r!= ncclSuccess) { \ + printf("Failed, NCCL error %s:%d '%s'\n", \ + __FILE__,__LINE__,ncclGetErrorString(r)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + +extern "C" { + void init_process_group(int rank, int world_size); + void broadcast_tensor(Tensor* tensor); + void allreduce_sum_tensor(Tensor* tensor); +} + +#endif /* DISTRIBUTED_H */ diff --git a/build/lib/norch/csrc/tensor.cpp b/build/lib/norch/csrc/tensor.cpp new file mode 100644 index 0000000..725b9f5 --- /dev/null +++ b/build/lib/norch/csrc/tensor.cpp @@ -0,0 +1,1558 @@ +#include +#include +#include +#include +#include +#include "tensor.h" +#include "cuda.h" +#include "cpu.h" + +extern "C" { + + Tensor* create_tensor(float* data, int* shape, int ndim, char* device) { + + Tensor* tensor = (Tensor*)malloc(sizeof(Tensor)); + if (tensor == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + tensor->data = data; + tensor->shape = shape; + tensor->ndim = ndim; + + tensor->device = (char*)malloc(strlen(device) + 1); + if (device != NULL) { + strcpy(tensor->device, device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + tensor->size = 1; + for (int i = 0; i < ndim; i++) { + tensor->size *= shape[i]; + } + + tensor->strides = (int*)malloc(ndim * sizeof(int)); + if (tensor->strides == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + int stride = 1; + for (int i = ndim - 1; i >= 0; i--) { + tensor->strides[i] = stride; + stride *= shape[i]; + } + + return tensor; + } + + float get_item(Tensor* tensor, int* indices) { + int index = 0; + for (int i = 0; i < tensor->ndim; i++) { + index += indices[i] * tensor->strides[i]; + } + + float result; + if (strcmp(tensor->device, "cuda") == 0) { + cudaMemcpy(&result, tensor->data + index, sizeof(float), cudaMemcpyDeviceToHost); + } else { + result = tensor->data[index]; + } + + return result; + } + + void to_device(Tensor* tensor, char* target_device) { + int device_id = 0; + char* endptr; + long num = strtol(target_device, &endptr, 10); + if (*endptr == '\0') { + device_id = (int)num; + } + + if ((strcmp(target_device, "cuda") == 0) && (strcmp(tensor->device, "cpu") == 0)) { + cpu_to_cuda(tensor, device_id); + } + + else if ((strcmp(target_device, "cpu") == 0) && (strcmp(tensor->device, "cuda") == 0)) { + cuda_to_cpu(tensor); + } + } + + Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2) { + if (tensor1->ndim != tensor2->ndim) { + fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for addition\n", tensor1->ndim, tensor2->ndim); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor1->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + if (tensor1->shape[i] != tensor2->shape[i]) { + fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for addition\n", tensor1->shape[i], tensor2->shape[i], i); + exit(1); + } + shape[i] = tensor1->shape[i]; + } + + if (strcmp(tensor1->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor1->size * sizeof(float)); + add_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor1->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + add_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + 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; + } + + 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, broadcasted_size * sizeof(float)); + add_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size); + return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device); + } + else { + 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, broadcasted_size); + return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device); + } + } + + Tensor* sum_tensor(Tensor* tensor, int axis, bool keepdim) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim; + int* shape; + + if (axis > tensor->ndim - 1) { + fprintf(stderr, "Error: axis argument %d must be smaller than tensor dimension %d", axis, tensor->ndim); + } + + if (axis == -1) { + + shape = (int*) malloc(sizeof(int)); + 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; + } + + int axis_size = 1; + for (int i = 0; i < ndim; i++) { + axis_size *= shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + if (axis == -1) { + cudaMalloc((void**)&result_data, tensor->size * sizeof(float)); + } else { + cudaMalloc((void**)&result_data, axis_size * sizeof(float)); + } + sum_tensor_cuda(tensor, result_data, axis); + + if (keepdim) { + if (axis == -1){ + ndim = tensor->ndim; + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = 1; + } + } else { + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = tensor->shape[i]; + } + shape[axis] = 1; + ndim = tensor->ndim; + } + + } + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)calloc(axis_size, sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + sum_tensor_cpu(tensor, result_data, axis_size, shape, axis); + + if (keepdim) { + if (axis == -1){ + ndim = tensor->ndim; + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = 1; + } + } else { + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = tensor->shape[i]; + } + shape[axis] = 1; + ndim = tensor->ndim; + } + + } + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* max_tensor(Tensor* tensor, int axis, bool keepdim) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim; + int* shape; + if (axis == -1) { + shape = (int*) malloc(sizeof(int)); + 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; + } + + int axis_size = 1; + for (int i = 0; i < ndim; i++) { + axis_size *= shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + if (axis == -1) { + cudaMalloc((void**)&result_data, tensor->size * sizeof(float)); + } else { + cudaMalloc((void**)&result_data, axis_size * sizeof(float)); + } + max_tensor_cuda(tensor, result_data, axis); + + if (keepdim) { + if (axis == -1){ + ndim = tensor->ndim; + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = 1; + } + } else { + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = tensor->shape[i]; + } + shape[axis] = 1; + ndim = tensor->ndim; + } + + } + return create_tensor(result_data, shape, ndim, device); + + } + else { + float* result_data = (float*)malloc(axis_size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + max_tensor_cpu(tensor, result_data, axis_size, shape, axis); + + if (keepdim) { + if (axis == -1){ + ndim = tensor->ndim; + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = 1; + } + } else { + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = tensor->shape[i]; + } + shape[axis] = 1; + ndim = tensor->ndim; + } + + } + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* min_tensor(Tensor* tensor, int axis, bool keepdim) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim; + int* shape; + if (axis == -1) { + shape = (int*) malloc(sizeof(int)); + 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; + } + + int axis_size = 1; + for (int i = 0; i < ndim; i++) { + axis_size *= shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + if (axis == -1) { + cudaMalloc((void**)&result_data, tensor->size * sizeof(float)); + } else { + cudaMalloc((void**)&result_data, axis_size * sizeof(float)); + } + min_tensor_cuda(tensor, result_data, axis); + + if (keepdim) { + if (axis == -1){ + ndim = tensor->ndim; + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = 1; + } + } else { + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = tensor->shape[i]; + } + shape[axis] = 1; + ndim = tensor->ndim; + } + + } + return create_tensor(result_data, shape, ndim, device); + + } + else { + float* result_data = (float*)malloc(axis_size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + min_tensor_cpu(tensor, result_data, axis_size, shape, axis); + + if (keepdim) { + if (axis == -1){ + ndim = tensor->ndim; + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = 1; + } + } else { + shape = (int*) malloc((tensor->ndim) * sizeof(int)); + for (int i = 0; i < tensor->ndim; i++) { + shape[i] = tensor->shape[i]; + } + shape[axis] = 1; + ndim = tensor->ndim; + } + + } + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2) { + if (tensor1->ndim != tensor2->ndim) { + fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for subtraction\n", tensor1->ndim, tensor2->ndim); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor1->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + if (tensor1->shape[i] != tensor2->shape[i]) { + fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for subtraction\n", tensor1->shape[i], tensor2->shape[i], i); + exit(1); + } + shape[i] = tensor1->shape[i]; + } + + if (strcmp(tensor1->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor1->size * sizeof(float)); + sub_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor1->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + sub_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + 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; + } + + 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, broadcasted_size * sizeof(float)); + sub_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size); + return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device); + } + else { + float* result_data = (float*)malloc(broadcasted_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, broadcasted_size); + 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); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor1->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + if (tensor1->shape[i] != tensor2->shape[i]) { + fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for element-wise multiplication\n", tensor1->shape[i], tensor2->shape[i], i); + exit(1); + } + shape[i] = tensor1->shape[i]; + } + + if (strcmp(tensor1->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor1->size * sizeof(float)); + elementwise_mul_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor1->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + elementwise_mul_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* scalar_mul_tensor(Tensor* tensor, float scalar) { + + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + scalar_mul_tensor_cuda(tensor, scalar, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + scalar_mul_tensor_cpu(tensor, scalar, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* scalar_div_tensor(float scalar, Tensor* tensor) { + + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + scalar_div_tensor_cuda(scalar, tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + scalar_div_tensor_cpu(scalar, tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* tensor_div_scalar(Tensor* tensor, float scalar) { + + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + tensor_div_scalar_cuda(tensor, scalar, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + tensor_div_scalar_cpu(tensor, scalar, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* tensor_div_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 division\n", tensor1->ndim, tensor2->ndim); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor1->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + if (tensor1->shape[i] != tensor2->shape[i]) { + fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for division\n", tensor1->shape[i], tensor2->shape[i], i); + exit(1); + } + shape[i] = tensor1->shape[i]; + } + + if (strcmp(tensor1->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor1->size * sizeof(float)); + tensor_div_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor1->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + tensor_div_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* matmul_tensor(Tensor* tensor1, Tensor* tensor2) { + //MxN @ NxP = MxP + // Check if tensors have compatible shapes for matrix multiplication + if (tensor1->shape[1] != tensor2->shape[0]) { + fprintf(stderr, "Incompatible shapes for matrix multiplication %dx%d and %dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1]); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor1->ndim + tensor2->ndim - 2; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + for (int i = 0; i < tensor1->ndim - 1; i++) { + shape[i] = tensor1->shape[i]; + } + for (int i = tensor1->ndim - 1; i < ndim; i++) { + shape[i] = tensor2->shape[i - tensor1->ndim + 2]; + } + + int size = 1; + for (int i = 0; i < ndim; i++) { + size *= shape[i]; + } + + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + if (strcmp(tensor1->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, size * sizeof(float)); + matmul_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + matmul_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* broadcasted_batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) { + //BATCHxMxP = MxN @ BATCHxNxP + // Check if tensors have compatible shapes for matrix multiplication + if (tensor1->shape[1] != tensor2->shape[1]) { + fprintf(stderr, "Incompatible shapes for broadcasted batched matrix multiplication %dx%d and %dx%dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1], tensor2->shape[2]); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + int ndim = 3; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + shape[0] = tensor2->shape[0];; + shape[1] = tensor1->shape[0]; + shape[2] = tensor2->shape[2]; + + int size = 1; + for (int i = 0; i < ndim; i++) { + size *= shape[i]; + } + + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + if (strcmp(tensor1->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, size * sizeof(float)); + broadcasted_batched_matmul_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + broadcasted_batched_matmul_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + + } + + Tensor* batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) { + //BATCHxMxP = BATCHxMxN @ BATCHxNxP + // Check if tensors have compatible shapes for matrix multiplication + + if (tensor1->shape[0] != tensor2->shape[0]) { + fprintf(stderr, "Tensors must have same batch dimension for batch matmul %d and %d\n", tensor1->shape[0], tensor2->shape[0]); + exit(1); + } + + if (tensor1->shape[2] != tensor2->shape[1]) { + fprintf(stderr, "Incompatible shapes for matrix multiplication %dx%d and %dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1]); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + int ndim = 3; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + shape[0] = tensor2->shape[0];; + shape[1] = tensor1->shape[1]; + shape[2] = tensor2->shape[2]; + + int size = 1; + for (int i = 0; i < ndim; i++) { + size *= shape[i]; + } + + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + if (strcmp(tensor1->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, size * sizeof(float)); + batched_matmul_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + batched_matmul_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + + } + + + Tensor* tensor_pow_scalar(Tensor* tensor, float exponent) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + tensor_pow_scalar_cuda(tensor, exponent, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + tensor_pow_scalar_cpu(tensor, exponent, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* scalar_pow_tensor(float base, Tensor* tensor) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + scalar_pow_tensor_cuda(base, tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + scalar_pow_tensor_cpu(base, tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* log_tensor(Tensor* tensor) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + log_tensor_cuda(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + log_tensor_cpu(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + int ndim = new_ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = new_shape[i]; + } + + // Calculate the total number of elements in the new shape + int size = 1; + for (int i = 0; i < new_ndim; i++) { + size *= shape[i]; + } + + // Check if the total number of elements matches the current tensor's size + if (size != tensor->size) { + fprintf(stderr, "Cannot reshape tensor. Total number of elements in new shape does not match the current size of the tensor.\n"); + exit(1); + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + assign_tensor_cuda(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + assign_tensor_cpu(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* equal_tensor(Tensor* tensor1, Tensor* tensor2) { + if (tensor1->ndim != tensor2->ndim) { + fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for equal\n", tensor1->ndim, tensor2->ndim); + exit(1); + } + + 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); + } + + char* device = (char*)malloc(strlen(tensor1->device) + 1); + if (device != NULL) { + strcpy(device, tensor1->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor1->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + if (tensor1->shape[i] != tensor2->shape[i]) { + fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for equal\n", tensor1->shape[i], tensor2->shape[i], i); + exit(1); + } + shape[i] = tensor1->shape[i]; + } + + if (strcmp(tensor1->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor1->size * sizeof(float)); + equal_tensor_cuda(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor1->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + equal_tensor_cpu(tensor1, tensor2, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* equal_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; + } + + 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, broadcasted_size * sizeof(float)); + equal_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size); + return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device); + } + else { + float* result_data = (float*)malloc(broadcasted_size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + equal_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size); + return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device); + } + } + + + Tensor* ones_like_tensor(Tensor* tensor) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + ones_like_tensor_cuda(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + ones_like_tensor_cpu(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* zeros_like_tensor(Tensor* tensor) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + zeros_like_tensor_cuda(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + zeros_like_tensor_cpu(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* sin_tensor(Tensor* tensor) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + sin_tensor_cuda(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + sin_tensor_cpu(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* cos_tensor(Tensor* tensor) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + cos_tensor_cuda(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + cos_tensor_cpu(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* sigmoid_tensor(Tensor* tensor) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + sigmoid_tensor_cuda(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + sigmoid_tensor_cpu(tensor, result_data); + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* transpose_tensor(Tensor* tensor) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[ndim - 1 - i]; + } + + int size = tensor->size; + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, size * sizeof(float)); + switch (ndim) { + case 1: + transpose_1D_tensor_cuda(tensor, result_data); + break; + case 2: + transpose_2D_tensor_cuda(tensor, result_data); + break; + case 3: + transpose_3D_tensor_cuda(tensor, result_data); + break; + default: + fprintf(stderr, "Transpose only supports tensors up to 3 dimensions.\n"); + exit(-1); + } + return create_tensor(result_data, shape, ndim, device); + } + else { + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + switch (ndim) { + case 1: + transpose_1D_tensor_cpu(tensor, result_data); + break; + case 2: + transpose_2D_tensor_cpu(tensor, result_data); + break; + case 3: + transpose_3D_tensor_cpu(tensor, result_data); + break; + default: + fprintf(stderr, "Transpose only supports tensors up to 3 dimensions.\n"); + exit(-1); + } + return create_tensor(result_data, shape, ndim, device); + } + } + + Tensor* transpose_axes_tensor(Tensor* tensor, int axis1, int axis2) { + char* device = (char*)malloc(strlen(tensor->device) + 1); + if (device != NULL) { + strcpy(device, tensor->device); + } else { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + int ndim = tensor->ndim; + int* shape = (int*)malloc(ndim * sizeof(int)); + if (shape == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(-1); + } + + for (int i = 0; i < ndim; i++) { + shape[i] = tensor->shape[i]; + } + + shape[axis1] = tensor->shape[axis2]; + shape[axis2] = tensor->shape[axis1]; + + int size = tensor->size; + + if (strcmp(tensor->device, "cuda") == 0) { + + float* result_data; + cudaMalloc((void **)&result_data, size * sizeof(float)); + assign_tensor_cuda(tensor, result_data); + + Tensor* new_tensor = create_tensor(result_data, shape, ndim, device); + for (int i = 0; i < ndim; i++) { + new_tensor->strides[i] = tensor->strides[i]; + } + new_tensor->strides[axis1] = tensor->strides[axis2]; + new_tensor->strides[axis2] = tensor->strides[axis1]; + make_contiguous(new_tensor); + return new_tensor; + } + else { + float* result_data = (float*)malloc(size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + exit(1); + } + assign_tensor_cpu(tensor, result_data); + + Tensor* new_tensor = create_tensor(result_data, shape, ndim, device); + for (int i = 0; i < ndim; i++) { + new_tensor->strides[i] = tensor->strides[i]; + } + new_tensor->strides[axis1] = tensor->strides[axis2]; + new_tensor->strides[axis2] = tensor->strides[axis1]; + make_contiguous(new_tensor); + return new_tensor; + } + } + + void make_contiguous(Tensor* tensor) { + + int* new_strides = (int*)malloc(tensor->ndim * sizeof(int)); + if (new_strides == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + } + + // Calculate new strides assuming C-contiguous order + int stride = 1; + for (int i = tensor->ndim - 1; i >= 0; i--) { + new_strides[i] = stride; + stride *= tensor->shape[i]; + } + + if (strcmp(tensor->device, "cuda") == 0) { + float* result_data; + cudaMalloc((void **)&result_data, tensor->size * sizeof(float)); + make_contiguous_tensor_cuda(tensor, result_data, new_strides); + + } else { + float* result_data = (float*)malloc(tensor->size * sizeof(float)); + if (result_data == NULL) { + fprintf(stderr, "Memory allocation failed\n"); + } + make_contiguous_tensor_cpu(tensor, result_data, new_strides); + } + } +} diff --git a/build/lib/norch/csrc/tensor.h b/build/lib/norch/csrc/tensor.h new file mode 100644 index 0000000..5669d4f --- /dev/null +++ b/build/lib/norch/csrc/tensor.h @@ -0,0 +1,45 @@ +#ifndef TENSOR_H +#define TENSOR_H + +typedef struct { + float* data; + int* strides; + int* shape; + int* strides_cuda; + int* shape_cuda; + int ndim; + int size; + char* device; +} Tensor; + +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, int axis, bool keepdims); + Tensor* max_tensor(Tensor* tensor, int axis, bool keepdim); + Tensor* min_tensor(Tensor* tensor, int axis, bool keepdim); + Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2); + Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2); + Tensor* scalar_mul_tensor(Tensor* tensor, float scalar); + Tensor* scalar_div_tensor(float scalar, Tensor* tensor); + Tensor* tensor_div_scalar(Tensor* tensor, float scalar); + Tensor* tensor_div_tensor(Tensor* tensor1, Tensor* tensor2); + Tensor* reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim); + Tensor* matmul_tensor(Tensor* tensor1, Tensor* tensor2); + Tensor* tensor_pow_scalar(Tensor* tensor, float exponent); + Tensor* scalar_pow_tensor(float base, Tensor* tensor); + Tensor* log_tensor(Tensor* tensor); + Tensor* equal_tensor(Tensor* tensor1, Tensor* tensor2); + Tensor* equal_broadcasted_tensor(Tensor* tensor1, Tensor* tensor2); + void to_device(Tensor* tensor, char* device); + Tensor* ones_like_tensor(Tensor* tensor); + Tensor* zeros_like_tensor(Tensor* tensor); + Tensor* sin_tensor(Tensor* tensor); + Tensor* cos_tensor(Tensor* tensor); + Tensor* transpose_tensor(Tensor* tensor); + Tensor* transpose_axes_tensor(Tensor* tensor, int axis1, int axis2); + void make_contiguous(Tensor* tensor); +} + +#endif /* TENSOR_H */ diff --git a/build/lib/norch/distributed/__init__.py b/build/lib/norch/distributed/__init__.py new file mode 100644 index 0000000..a979451 --- /dev/null +++ b/build/lib/norch/distributed/__init__.py @@ -0,0 +1 @@ +from .distributed import * \ No newline at end of file diff --git a/build/lib/norch/distributed/distributed.py b/build/lib/norch/distributed/distributed.py new file mode 100644 index 0000000..d033ede --- /dev/null +++ b/build/lib/norch/distributed/distributed.py @@ -0,0 +1,24 @@ +import os +import ctypes +from norch import Tensor, CTensor + +def init_process_group(rank, world_size, backend='nccl'): + + Tensor._C.init_process_group.argtypes = [ctypes.c_int, ctypes.c_int] + Tensor._C.init_process_group.restype = None + + Tensor._C.init_process_group(rank, world_size) + +def broadcast_tensor(tensor): + + Tensor._C.broadcast_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.broadcast_tensor.restype = None + + Tensor._C.broadcast_tensor(tensor) + +def allreduce_sum_tensor(tensor): + + Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.allreduce_mean_tensor.restype = None + + Tensor._C.allreduce_mean_tensor(tensor) diff --git a/build/lib/norch/libtensor.so b/build/lib/norch/libtensor.so new file mode 100755 index 0000000..b70187e Binary files /dev/null and b/build/lib/norch/libtensor.so differ diff --git a/build/lib/norch/nn/__init__.py b/build/lib/norch/nn/__init__.py new file mode 100644 index 0000000..8f7997b --- /dev/null +++ b/build/lib/norch/nn/__init__.py @@ -0,0 +1,4 @@ +from .modules import * +from .activation import * +from .loss import * +from .functional import * \ No newline at end of file diff --git a/build/lib/norch/nn/activation.py b/build/lib/norch/nn/activation.py new file mode 100644 index 0000000..a1361f6 --- /dev/null +++ b/build/lib/norch/nn/activation.py @@ -0,0 +1,30 @@ +from .module import Module +from . import functional as F +import math + +class Activation(Module): + """ + Abstract classes for activations + """ + def __init__(self): + super(Activation, self).__init__() + + def forward(self, x): + raise NotImplementedError + + +class Sigmoid(Activation): + def __init__(self): + super().__init__() + + def forward(self, x): + return F.sigmoid(x) + +class Softmax(Activation): + def __init__(self, dim): + super(Softmax, self).__init__() + + self.dim = dim + + def forward(self, x): + return F.softmax(x, self.dim) \ No newline at end of file diff --git a/build/lib/norch/nn/functional.py b/build/lib/norch/nn/functional.py new file mode 100644 index 0000000..29fb68d --- /dev/null +++ b/build/lib/norch/nn/functional.py @@ -0,0 +1,32 @@ +import math +import norch +import numpy as np +from norch.autograd.functions import * + +def sigmoid(x): + z = x.sigmoid() + return z + +def softmax(x, dim=None): + if dim is not None and dim < 0: + dim = x.ndim + dim + + x_max = x.max(axis=dim, keepdim=True) + exp_x = math.e ** (x - x_max) + + if dim is not None: + sum_exp_x = exp_x.sum(axis=dim, keepdim=True) + exp_x.zeros_like() + return exp_x / sum_exp_x + else: + sum_exp_x = exp_x.sum() + return exp_x / sum_exp_x + +def one_hot_encode(x, num_classes): + one_hot = [[0] * num_classes for _ in range(x.numel)] + + # Set the appropriate elements to 1 + for i in range(x.numel): + target_idx = int(x[i]) + one_hot[i][target_idx] = 1 + + return norch.Tensor(one_hot) \ No newline at end of file diff --git a/build/lib/norch/nn/loss.py b/build/lib/norch/nn/loss.py new file mode 100644 index 0000000..6c0abea --- /dev/null +++ b/build/lib/norch/nn/loss.py @@ -0,0 +1,93 @@ +from .module import Module +from norch.autograd.functions import * +import norch +from abc import ABC + +class Loss(Module, ABC): + "Abstract class for loss functions" + + def __init__(self): + super().__init__() + + def forward(self, predictions, labels): + raise NotImplementedError + + def __call__(self, *inputs): + return self.forward(*inputs) + + +class MSELoss(Loss): + def __init__(self): + super().__init__() + + def forward(self, predictions, target): + assert target.shape == predictions.shape, \ + "Labels and predictions shape does not match: {} and {}".format(target.shape, predictions.shape) + + cost = ((predictions - target) ** 2).sum() / predictions.numel + return cost + + +class CrossEntropyLoss(Loss): + def __init__(self): + super().__init__() + + def forward(self, input, target): + assert isinstance(input, norch.Tensor), \ + "Cross entropy argument 'input' must be Tensor, not {}".format(type(input)) + + assert isinstance(target, norch.Tensor), \ + "Cross entropy argument 'target' must be Tensor, not {}".format(type(target)) + if input.ndim > 2: + input = input.squeeze(-1) + + if input.ndim == 1: + if target.numel == 1: + num_classes = input.shape[0] + target = norch.one_hot_encode(target, num_classes).to(target.device) + + logits = norch.softmax(input, dim=0) + target = target.reshape(logits.shape) + cost = -(logits.log() * target).sum() + + else: + # target -> class probabilities (one-hot encoded) + assert target.shape == input.shape, \ + "Input and target shape does not match: {} and {}".format(input.shape, target.shape) + logits = norch.softmax(input, dim=0) + target = target.reshape(logits.shape) + cost = -(logits.log() * target).sum() + + + elif input.ndim == 2: + if target.ndim > 1: + target = target.squeeze(-1) + # batched + if target.ndim == 1: + # target -> Ground truth class indices: + num_classes = input.shape[1] + + target = norch.one_hot_encode(target, num_classes).to(target.device) + + batch_size = input.shape[0] + logits = norch.softmax(input, dim=1) + target = target.reshape(logits.shape) + cost = -(logits.log() * target).sum() / batch_size + + else: + # target -> class probabilities (one-hot encoded) + assert target.shape == input.shape, \ + "Input and target shape does not match: {} and {}".format(input.shape, target.shape) + + batch_size = input.shape[0] + logits = norch.softmax(input, dim=1) + target = target.reshape(logits.shape) + cost = -(logits.log() * target).sum() / batch_size + + if input.requires_grad: + cost.grad_fn = CrossEntropyLossBackward(input, target) + + return cost + + + diff --git a/build/lib/norch/nn/module.py b/build/lib/norch/nn/module.py new file mode 100644 index 0000000..0135bb6 --- /dev/null +++ b/build/lib/norch/nn/module.py @@ -0,0 +1,107 @@ +from .parameter import Parameter +from collections import OrderedDict +from abc import ABC +import pickle +import json +import inspect +import warnings + +class Module(ABC): + """ + Abstract class for modules + """ + def __init__(self): + self._modules = OrderedDict() + self._params = OrderedDict() + self._grads = OrderedDict() + self.training = True + + def forward(self, *inputs, **kwargs): + raise NotImplementedError + + def __call__(self, *inputs, **kwargs): + return self.forward(*inputs, **kwargs) + + def train(self): + self.training = True + for param in self.parameters(): + param.requires_grad = True + + def eval(self): + self.training = False + for param in self.parameters(): + param.requires_grad = False + + def parameters(self): + for name, value in inspect.getmembers(self): + if isinstance(value, Parameter): + yield self, name, value + elif isinstance(value, Module): + yield from value.parameters() + + def modules(self): + yield from self._modules.values() + + def gradients(self): + for module in self.modules(): + yield module._grads + + def zero_grad(self): + for _, _, parameter in self.parameters(): + parameter.zero_grad() + + def to(self, device): + for module, name, _ in self.parameters(): + parameter = getattr(module, name) + parameter = parameter.to(device) + setattr(module, name, parameter) + + return self + + def state_dict(self): + state = OrderedDict() + for i, param in enumerate(self.parameters()): + state[f'param{i}'] = param.tolist() + return state + + def load_state(self, state_dict): + for i, param in self.parameters(): + data = state_dict[f'param{i}'] + if param.shape != data.shape: + warnings.warn(f"The 'state_dict' shape does not match model's parameter shape. " + f"Got {data.shape}, expected {param.shape}.") + param.data = Parameter(data=data) + + def save(self, filename='model.pickle'): + with open(filename, 'wb') as f: + pickle.dump(self, f) + + def save_dict(self, filename='state_dict.json'): + state = self.state_dict() + with open(filename, 'w') as f: + json.dump(state, f) + + def inner_repr(self): + return "" + + def __repr__(self): + string = f"{self.get_name()}(" + tab = " " + modules = self._modules + if modules == {}: + string += f'\n{tab}(parameters): {self.inner_repr()}' + else: + for key, module in modules.items(): + string += f"\n{tab}({key}): {module.get_name()}({module.inner_repr()})" + return f'{string}\n)' + + def get_name(self): + return self.__class__.__name__ + + def __setattr__(self, key, value): + self.__dict__[key] = value + + if isinstance(value, Module): + self._modules[key] = value + elif isinstance(value, Parameter): + self._params[key] = value \ No newline at end of file diff --git a/build/lib/norch/nn/modules/__init__.py b/build/lib/norch/nn/modules/__init__.py new file mode 100644 index 0000000..5c73a2e --- /dev/null +++ b/build/lib/norch/nn/modules/__init__.py @@ -0,0 +1 @@ +from .linear import * \ No newline at end of file diff --git a/build/lib/norch/nn/modules/linear.py b/build/lib/norch/nn/modules/linear.py new file mode 100644 index 0000000..b90d845 --- /dev/null +++ b/build/lib/norch/nn/modules/linear.py @@ -0,0 +1,26 @@ +from ..module import Module +from ..parameter import Parameter + +class Linear(Module): + def __init__(self, input_dim, output_dim, bias=True): + super().__init__() + self.input_dim = input_dim + self.output_dim = output_dim + self.weight = Parameter(shape=[self.output_dim, self.input_dim]) + + if bias: + self.bias = Parameter(shape=[self.output_dim, 1]) + else: + self.bias = None + + def forward(self, x): + if self.bias: + z = self.weight @ x + self.bias + else: + z = self.weight @ x + + return z + + def inner_repr(self): + return f"input_dim={self.input_dim}, output_dim={self.output_dim}, " \ + f"bias={True if self.bias is not None else False}" \ No newline at end of file diff --git a/build/lib/norch/nn/parameter.py b/build/lib/norch/nn/parameter.py new file mode 100644 index 0000000..2324a1a --- /dev/null +++ b/build/lib/norch/nn/parameter.py @@ -0,0 +1,11 @@ +from norch.tensor import Tensor +from norch.utils import utils +import random + +class Parameter(Tensor): + """ + A parameter is a trainable tensor. + """ + def __init__(self, shape): + data = utils.generate_random_list(shape=shape) + super().__init__(data, requires_grad=True) \ No newline at end of file diff --git a/build/lib/norch/norchvision/__init__.py b/build/lib/norch/norchvision/__init__.py new file mode 100644 index 0000000..a6ccfac --- /dev/null +++ b/build/lib/norch/norchvision/__init__.py @@ -0,0 +1 @@ +from .datasets import * \ No newline at end of file diff --git a/build/lib/norch/norchvision/datasets/__init__.py b/build/lib/norch/norchvision/datasets/__init__.py new file mode 100644 index 0000000..fe34b11 --- /dev/null +++ b/build/lib/norch/norchvision/datasets/__init__.py @@ -0,0 +1 @@ +from .mnist import * \ No newline at end of file diff --git a/build/lib/norch/norchvision/datasets/mnist.py b/build/lib/norch/norchvision/datasets/mnist.py new file mode 100644 index 0000000..ff11edf --- /dev/null +++ b/build/lib/norch/norchvision/datasets/mnist.py @@ -0,0 +1,90 @@ +import gzip +import os +import norch +from norch.utils.data import Dataset +import numpy as np + +class MNIST(Dataset): + """ + Loads training, validation, and test partitions of the mnist dataset + (http://yann.lecun.com/exdb/mnist/). If the data is not already contained in data_dir, it will + try to download it. + + This dataset contains 60000 training examples, and 10000 test examples of handwritten digits + in {0, ..., 9} and corresponding labels. Each handwritten image has an "original" dimension of + 28x28x1, and is stored row-wise as a string of 784x1 bytes. Pixel values are in range 0 to 255 + (inclusive). + + Args: + data_dir: String. Relative or absolute path of the dataset. + devel_size: Integer. Size of the development (validation) dataset partition. + + Returns: + X_train: float64 numpy array with shape [784, 60000-devel_size] with values in [0, 1]. + Y_train: uint8 numpy array with shape [60000-devel_size]. Labels. + X_devel: float64 numpy array with shape [784, devel_size] with values in [0, 1]. + Y_devel: uint8 numpy array with shape [devel_size]. Labels. + X_test: float64 numpy array with shape [784, 10000] with values in [0, 1]. + Y_test: uint8 numpy array with shape [10000]. Labels. + """ + + urls = ['https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz', + 'https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz', + 'https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz', + 'https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz',] + name = 'mnist-data-py' + dirname = 'mnist' + + def __init__(self, path_data, path_label, transform=None, target_transform=None): + self.data = self._load_mnist(path_data, header_size=16).reshape((-1, 28, 28)) + self.labels = self._load_mnist(path_label, header_size=8) + self.transform = transform + self.target_transform = target_transform + + def _load_mnist(self, path, header_size): + with gzip.open(path, 'rb') as f: + data = np.frombuffer(f.read(), np.uint8, offset=header_size) + return np.asarray(data, dtype=np.uint8) + + + @classmethod + def splits(cls, root='.data', train_data='train-images-idx3-ubyte.gz', train_label='train-labels-idx1-ubyte.gz', + test_data='t10k-images-idx3-ubyte.gz', test_label='t10k-labels-idx1-ubyte.gz', **kwargs): + r""" + Loads training and test partitions of the [mnist dataset](https://www.cs.toronto.edu/~kriz/cifar.html). If + the data is not already contained in the ``root`` folder, it will download it. + + Args: + root (str): relative or absolute path of the dataset. + + Returns: + tuple(Dataset): training and testing datasets + """ + path = os.path.join(root, cls.dirname, cls.name) + if not os.path.isdir(path): + path = cls.download(root) + train_data = os.path.join(path, train_data) + train_label = os.path.join(path, train_label) + test_data = os.path.join(path, test_data) + test_label = os.path.join(path, test_label) + return MNIST(train_data, train_label, **kwargs), MNIST(test_data, test_label, **kwargs) + + def __getitem__(self, item): + data = self.data[item].tolist() + label = self.labels[item].tolist() + + if self.transform is not None: + data = self.transform(data) + + if self.target_transform is not None: + label = self.target_transform(label) + + + + return data, label + + def __setitem__(self, key, value): + self.data[key], self.labels[key] = value + + def __len__(self): + return len(self.data) diff --git a/build/lib/norch/norchvision/transforms.py b/build/lib/norch/norchvision/transforms.py new file mode 100644 index 0000000..03dafad --- /dev/null +++ b/build/lib/norch/norchvision/transforms.py @@ -0,0 +1,22 @@ +import norch + +class ToTensor: + def __call__(self, x): + return norch.Tensor(x) + +class Reshape: + def __init__(self, shape): + self.shape = shape + + def __call__(self, x): + return x.reshape(self.shape) + +class Sequential: + def __init__(self, transforms): + self.transforms = transforms + + def __call__(self, x): + for transform in self.transforms: + x = transform(x) + return x + diff --git a/build/lib/norch/optim/__init__.py b/build/lib/norch/optim/__init__.py new file mode 100644 index 0000000..c5baa22 --- /dev/null +++ b/build/lib/norch/optim/__init__.py @@ -0,0 +1 @@ +from .optimizers import * \ No newline at end of file diff --git a/build/lib/norch/optim/optimizer.py b/build/lib/norch/optim/optimizer.py new file mode 100644 index 0000000..fbbd537 --- /dev/null +++ b/build/lib/norch/optim/optimizer.py @@ -0,0 +1,22 @@ +from abc import ABC +from norch.tensor import Tensor + +class Optimizer(ABC): + """ + Abstract class for optimizers + """ + + def __init__(self, parameters): + if isinstance(parameters, Tensor): + raise TypeError("parameters should be an iterable but got {}".format(type(parameters))) + elif isinstance(parameters, dict): + parameters = parameters.values() + + self.parameters = list(parameters) + + def step(self): + raise NotImplementedError + + def zero_grad(self): + for module, name, parameter in self.parameters: + parameter.zero_grad() \ No newline at end of file diff --git a/build/lib/norch/optim/optimizers/__init__.py b/build/lib/norch/optim/optimizers/__init__.py new file mode 100644 index 0000000..bd076ac --- /dev/null +++ b/build/lib/norch/optim/optimizers/__init__.py @@ -0,0 +1 @@ +from .sgd import * \ No newline at end of file diff --git a/build/lib/norch/optim/optimizers/sgd.py b/build/lib/norch/optim/optimizers/sgd.py new file mode 100644 index 0000000..1d785c7 --- /dev/null +++ b/build/lib/norch/optim/optimizers/sgd.py @@ -0,0 +1,27 @@ +from ..optimizer import Optimizer +from norch.tensor import Tensor +import time + +class SGD(Optimizer): + def __init__(self, parameters, lr=1e-1, momentum=0): + super().__init__(parameters) + self.lr = lr + self.momentum = momentum + self._cache = {'velocity': [p.zeros_like() for (_, _, p) in self.parameters]} + + def step(self): + for i, (module, name, _) in enumerate(self.parameters): + parameter = getattr(module, name) + + velocity = self._cache['velocity'][i] + + velocity = self.momentum * velocity - self.lr * parameter.grad + + updated_parameter = parameter + velocity + + setattr(module, name, updated_parameter) + + self._cache['velocity'][i] = velocity + + parameter.detach() + velocity.detach() diff --git a/build/lib/norch/tensor.py b/build/lib/norch/tensor.py new file mode 100644 index 0000000..2e80e72 --- /dev/null +++ b/build/lib/norch/tensor.py @@ -0,0 +1,1033 @@ +import ctypes +import os +from .autograd.functions import * + +class CTensor(ctypes.Structure): + _fields_ = [ + ('data', ctypes.POINTER(ctypes.c_float)), + ('strides', ctypes.POINTER(ctypes.c_int)), + ('shape', ctypes.POINTER(ctypes.c_int)), + ('ndim', ctypes.c_int), + ('size', ctypes.c_int), + ('device', ctypes.c_char_p) + ] + +class Tensor: + module_dir = os.path.dirname(os.path.abspath(__file__)) + _C = ctypes.CDLL(os.path.join(module_dir, "libtensor.so")) + + def __init__(self, data=None, device="cpu", requires_grad=False): + + if data != None: + if isinstance(data, (float, int)): + data = [data] + + data, shape = self.flatten(data) + + self.shape = shape.copy() + + self.data_ctype = (ctypes.c_float * len(data))(*data.copy()) + self.shape_ctype = (ctypes.c_int * len(shape))(*shape.copy()) + self.ndim_ctype = ctypes.c_int(len(shape)) + self.device_ctype = device.encode('utf-8') + + self.ndim = len(shape) + self.device = device + + self.numel = 1 + for s in self.shape: + self.numel *= s + + self.requires_grad = requires_grad + self.grad = None + self.grad_fn = None + + Tensor._C.create_tensor.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_char_p] + Tensor._C.create_tensor.restype = ctypes.POINTER(CTensor) + + + self.tensor = Tensor._C.create_tensor( + self.data_ctype, + self.shape_ctype, + self.ndim_ctype, + self.device_ctype + ) + + else: + self.tensor = None, + self.shape = None, + self.ndim = None, + self.device = device + self.requires_grad = None + self.grad = None + self.grad_fn = None + + def flatten(self, nested_list): + def flatten_recursively(nested_list): + flat_data = [] + shape = [] + if isinstance(nested_list, list): + for sublist in nested_list: + inner_data, inner_shape = flatten_recursively(sublist) + flat_data.extend(inner_data) + shape.append(len(nested_list)) + shape.extend(inner_shape) + else: + flat_data.append(nested_list) + return flat_data, shape + + flat_data, shape = flatten_recursively(nested_list) + return flat_data, shape + + def ones_like(self): + + Tensor._C.ones_like_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.ones_like_tensor.restype = ctypes.POINTER(CTensor) + Tensor._C.ones_like_tensor(self.tensor) + + result_tensor_ptr = Tensor._C.ones_like_tensor(self.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 + + return result_data + + def zeros_like(self): + + Tensor._C.zeros_like_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.zeros_like_tensor.restype = ctypes.POINTER(CTensor) + Tensor._C.zeros_like_tensor(self.tensor) + + result_tensor_ptr = Tensor._C.zeros_like_tensor(self.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 + + return result_data + + def reshape(self, new_shape): + # Calculate the total number of elements in the tensor + total_elements = self.numel + + # Check for the presence of -1 in new_shape + if new_shape.count(-1) > 1: + raise ValueError("Only one dimension can be inferred (set to -1).") + + inferred_dim = None + known_dims_product = 1 + for dim in new_shape: + if dim == -1: + inferred_dim = dim + else: + known_dims_product *= dim + + # Calculate the inferred dimension if -1 is present + if inferred_dim == -1: + inferred_dim_size = total_elements // known_dims_product + new_shape = [inferred_dim_size if dim == -1 else dim for dim in new_shape] + + new_shape_ctype = (ctypes.c_int * len(new_shape))(*new_shape) + new_ndim_ctype = ctypes.c_int(len(new_shape)) + + Tensor._C.reshape_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(ctypes.c_int), ctypes.c_int] + Tensor._C.reshape_tensor.restype = ctypes.POINTER(CTensor) + result_tensor_ptr = Tensor._C.reshape_tensor(self.tensor, new_shape_ctype, new_ndim_ctype) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = new_shape.copy() + result_data.ndim = len(new_shape) + result_data.device = self.device + result_data.numel = self.numel + + result_data.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = ReshapeBackward(self) + + return result_data + + def unsqueeze(self, dim): + if dim < 0: + dim = self.ndim + dim + 1 + + # Ensure the dimension is valid + if dim > self.ndim: + raise ValueError("Dimension out of range (expected to be in range of [0, {0}], but got {1})".format(self.ndim, dim)) + + # Create the new shape with an extra dimension of size 1 + new_shape = self.shape[:dim] + [1] + self.shape[dim:] + + return self.reshape(new_shape) + + def squeeze(self, dim=None): + if dim is not None: + if dim < 0: + dim = self.ndim + dim + + # Ensure the dimension is valid + if dim >= self.ndim or dim < 0: + raise ValueError("Dimension out of range (expected to be in range of [0, {0}), but got {1})".format(self.ndim, dim)) + + # Only squeeze the specified dimension if its size is 1 + if self.shape[dim] != 1: + return self + #raise ValueError("Dimension {0} does not have size 1 and cannot be squeezed".format(dim)) + + # Create the new shape without the specified dimension + new_shape = self.shape[:dim] + self.shape[dim+1:] + else: + # Create the new shape by removing all dimensions of size 1 + new_shape = [s for s in self.shape if s != 1] + + return self.reshape(new_shape) + + def to(self, device): + device = str(device) + + self.device = device + self.device_ctype = self.device.encode('utf-8') + + Tensor._C.to_device.argtypes = [ctypes.POINTER(CTensor), ctypes.c_char_p] + Tensor._C.to_device.restype = None + Tensor._C.to_device(self.tensor, self.device_ctype) + + return self + + def backward(self, gradient=None): + if not self.requires_grad: + return + + if gradient is None: + if self.shape == [1]: + gradient = Tensor([1]).to(self.device) + else: + raise RuntimeError("Gradient argument must be specified for non-scalar tensors.") + + + stack = [(self, gradient)] + visited = set() + + while stack: + tensor, grad = stack.pop() + + if tensor.grad is None: + tensor.grad = grad + else: + tensor.grad += grad + + # Propagate gradients to inputs if not a leaf tensor + if tensor.grad_fn is not None: + grads = tensor.grad_fn.backward(grad) + for tensor, grad in zip(tensor.grad_fn.input, grads): + if isinstance(tensor, Tensor) and tensor not in visited: + stack.append((tensor, grad)) + visited.add(tensor) + + def zero_grad(self): + self.grad = None + + def detach(self): + self.grad = None + self.grad_fn = None + + def __getitem__(self, indices): + if isinstance(indices, int): + indices = [indices] + if len(indices) != self.ndim: + raise ValueError("Number of indices must match the number of dimensions") + + Tensor._C.get_item.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(ctypes.c_int)] + Tensor._C.get_item.restype = ctypes.c_float + + indices = (ctypes.c_int * len(indices))(*indices) + value = Tensor._C.get_item(self.tensor, indices) + + return value + + def __str__(self): + def print_recursively(tensor, depth, index): + if depth == tensor.ndim - 1: + result = "" + for i in range(tensor.shape[-1]): + index[-1] = i + result += str(tensor[tuple(index)]) + ", " + return result.strip() + else: + result = "" + if depth > 0: + result += "\n" + " " * ((depth - 1) * 4) + for i in range(tensor.shape[depth]): + index[depth] = i + result += "[" + result += print_recursively(tensor, depth + 1, index) + "]," + if i < tensor.shape[depth] - 1: + result += "\n" + " " * (depth * 4) + return result.strip(",") + + index = [0] * self.ndim + result = "tensor([" + result += print_recursively(self, 0, index) + result += f"""], device="{self.device}", requires_grad={self.requires_grad})""" + return result + + def __repr__(self): + return self.__str__() + + def __add__(self, other): + if isinstance(other, (int, float)): + other = other * self.ones_like() + + broadcasted_shape_add = [] + + # 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 + + + 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_add.append(max(dim1, dim2)) + return broadcasted_shape_add, True + + broadcasted_shape_add, needs_broadcasting = broadcast_shape(self.shape, other.shape) + + if needs_broadcasting: + # Call add_broadcasted_tensor if broadcasting is needed + if other.ndim == self.ndim - 1: + other = other.reshape([1] + other.shape) + + elif self.ndim == other.ndim - 1: + self = self.reshape([1] + self.shape) + + + + 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_add.copy() + result_data.ndim = len(broadcasted_shape_add) + + 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 or other.requires_grad + if result_data.requires_grad: + result_data.grad_fn = AddBroadcastedBackward(self, other) + + 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_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 # 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)): + other = other * self.ones_like() + + 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) + + result_tensor_ptr = Tensor._C.add_tensor(other.tensor, self.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.requires_grad = self.requires_grad or other.requires_grad + if result_data.requires_grad: + result_data.grad_fn = AddBackward(other, self) + + return result_data + + def __sub__(self, other): + if isinstance(other, (int, float)): + other = other * self.ones_like() + + broadcasted_shape_sub = [] + + # 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 + + + 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_sub.append(max(dim1, dim2)) + return broadcasted_shape_sub, True + + broadcasted_shape_sub, needs_broadcasting = broadcast_shape(self.shape, other.shape) + + if needs_broadcasting: + if other.ndim == self.ndim - 1: + other = other.reshape([1] + other.shape) + + elif self.ndim == other.ndim - 1: + self = self.reshape([1] + self.shape) + + # 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_sub.copy() + result_data.ndim = len(broadcasted_shape_sub) + + 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) + + 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_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 # 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 + + def __rsub__(self, other): + 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") + + 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(other.tensor, self.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.requires_grad = self.requires_grad or other.requires_grad + if result_data.requires_grad: + result_data.grad_fn = SubBackward(other, self) + + return result_data + + def __mul__(self, other): + if isinstance(other, (int, float)): + result_data = Tensor() + result_data.shape = self.shape.copy() + result_data.ndim = self.ndim + result_data.device = self.device + result_data.numel = self.numel + + Tensor._C.scalar_mul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_float] + Tensor._C.scalar_mul_tensor.restype = ctypes.POINTER(CTensor) + + result_data.tensor = Tensor._C.scalar_mul_tensor(self.tensor, ctypes.c_float(other)) + + result_data.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = ScalarMulBackward(self, other) + + return result_data + elif isinstance(other, Tensor): + if self.shape != other.shape: + raise ValueError("Tensors must have the same shape for element-wise multiplication") + + Tensor._C.elementwise_mul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.elementwise_mul_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.elementwise_mul_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.requires_grad = self.requires_grad or other.requires_grad + if result_data.requires_grad: + result_data.grad_fn = ElementwiseMulBackward(self, other) + + return result_data + else: + raise TypeError("Unsupported operand type(s) for *: '{}' and '{}'".format(type(self), type(other))) + + def __rmul__(self, other): + return self.__mul__(other) + + def __neg__(self): + return self.__mul__(-1) + + def __pos__(self): + return self + + def __matmul__(self, other): + if self.ndim < 3 and other.ndim == 3: + #broadcasted 2D x 3D matmul + + Tensor._C.broadcasted_batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.broadcasted_batched_matmul_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.broadcasted_batched_matmul_tensor(self.tensor, other.tensor) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = [other.shape[0], self.shape[0], other.shape[2]] + result_data.ndim = 3 + result_data.device = self.device + result_data.numel = 1 + for s in result_data.shape: + result_data.numel *= s + + + elif self.ndim == 3 and other.ndim == 3: + #broadcasted 3D x 3D matmul + + Tensor._C.batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.batched_matmul_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.batched_matmul_tensor(self.tensor, other.tensor) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = [other.shape[0], self.shape[1], other.shape[2]] + result_data.ndim = 3 + result_data.device = self.device + result_data.numel = 1 + for s in result_data.shape: + result_data.numel *= s + else: + #2D matmul + if self.ndim != 2 or other.ndim != 2: + raise ValueError("Matrix multiplication requires 2D tensors") + + if self.shape[1] != other.shape[0]: + raise ValueError("Incompatible shapes for matrix multiplication") + + Tensor._C.matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.matmul_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.matmul_tensor(self.tensor, other.tensor) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = [self.shape[0], other.shape[1]] + result_data.ndim = 2 + 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 or other.requires_grad + if result_data.requires_grad: + result_data.grad_fn = MatmulBackward(self, other) + + return result_data + + def __pow__(self, other): + other = float(other) + Tensor._C.tensor_pow_scalar.argtypes = [ctypes.POINTER(CTensor), ctypes.c_float] + Tensor._C.tensor_pow_scalar.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.tensor_pow_scalar(self.tensor, ctypes.c_float(other)) + + 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.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = PowBackward(self, other) + + return result_data + + def __rpow__(self, other): + other = float(other) + Tensor._C.scalar_pow_tensor.argtypes = [ctypes.c_float, ctypes.POINTER(CTensor)] + Tensor._C.scalar_pow_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.scalar_pow_tensor(ctypes.c_float(other), self.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.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = PowBackward(other, self) + + return result_data + + def __truediv__(self, other): + if isinstance(other, (int, float)): + other = float(other) + Tensor._C.tensor_div_scalar.argtypes = [ctypes.POINTER(CTensor), ctypes.c_float] + Tensor._C.tensor_div_scalar.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.tensor_div_scalar(self.tensor, ctypes.c_float(other)) + + 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.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = DivisionBackward(self, other) + + elif isinstance(self, Tensor) and isinstance(other, Tensor): + if other.numel == 1: + return self.__truediv__(other.tensor.contents.data[0]) + + Tensor._C.tensor_div_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.tensor_div_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.tensor_div_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.requires_grad = self.requires_grad or other.requires_grad + if result_data.requires_grad: + result_data.grad_fn = DivisionBackward(self, other) + + return result_data + + + + def __rtruediv__(self, other): + other = float(other) + + Tensor._C.scalar_div_tensor.argtypes = [ctypes.c_float, ctypes.POINTER(CTensor)] + Tensor._C.scalar_div_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.scalar_div_tensor(ctypes.c_float(other), self.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.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = DivisionBackward(other, self) + + return result_data + + + def equal(self, other): + + if isinstance(other, Tensor) and other.numel == 1: + # other is a single value tensor + other = self.zeros_like() + other + return self.equal(other) + + if not isinstance(other, Tensor): + # other is a single value + if isinstance(other, (int, float)): + other = self.zeros_like() + other + + return self.equal(other) + else: + return False + + broadcasted_shape_add = [] + + # 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 + + + 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_add.append(max(dim1, dim2)) + return broadcasted_shape_add, True + + broadcasted_shape_add, needs_broadcasting = broadcast_shape(self.shape, other.shape) + + if needs_broadcasting: + if other.ndim == self.ndim - 1: + other = other.reshape([1] + other.shape) + + elif self.ndim == other.ndim - 1: + self = self.reshape([1] + self.shape) + + # Call equal_broadcasted_tensor if broadcasting is needed + Tensor._C.equal_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.equal_broadcasted_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.equal_broadcasted_tensor(self.tensor, other.tensor) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = broadcasted_shape_add.copy() + result_data.ndim = len(broadcasted_shape_add) + + result_data.device = self.device + result_data.numel = 1 + for s in result_data.shape: + result_data.numel *= s + + else: + Tensor._C.equal_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)] + Tensor._C.equal_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.equal_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 + + return result_data + + def log(self): + Tensor._C.log_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.log_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.log_tensor(self.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.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = LogBackward(self) + + return result_data + + def sum(self, axis=None, keepdim=False): + if axis is not None and axis < 0: + axis = self.ndim + axis + + if axis == None: + axis = -1 + + if axis > self.ndim - 1: + raise ValueError(f"Error: axis argument {axis} cannot be higher than tensor dimension {self.ndim}") + + Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool] + Tensor._C.sum_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.sum_tensor(self.tensor, axis, keepdim) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + + if axis == -1: + if keepdim: + result_data.ndim = self.ndim + result_data.shape = [1] * self.ndim + + else: + result_data.shape = [1] + result_data.ndim = 1 + else: + if keepdim: + result_data.shape = self.shape[:axis] + [1] + self.shape[axis+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, axis, keepdim=keepdim) + + return result_data + + def max(self, axis=None, keepdim=False): + if axis is not None and axis < 0: + axis = self.ndim + axis + + if axis == None: + axis = -1 + + if axis > self.ndim - 1: + raise ValueError(f"Error: axis argument {axis} cannot be higher than tensor dimension {self.ndim}") + + Tensor._C.max_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool] + Tensor._C.max_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.max_tensor(self.tensor, axis, keepdim) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + + if axis == -1: + if keepdim: + result_data.ndim = self.ndim + result_data.shape = [1] * self.ndim + + else: + result_data.shape = [1] + result_data.ndim = 1 + else: + if keepdim: + result_data.shape = self.shape[:axis] + [1] + self.shape[axis+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 = MaxBackward(self, axis, keepdim=keepdim) + + return result_data + + def min(self, axis=None, keepdim=False): + if axis is not None and axis < 0: + axis = self.ndim + axis + + if axis == None: + axis = -1 + + if axis > self.ndim - 1: + raise ValueError(f"Error: axis argument {axis} must be smaller than tensor dimension {self.ndim}") + + Tensor._C.min_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_bool] + Tensor._C.min_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.min_tensor(self.tensor, axis, keepdim) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + + if axis == -1: + if keepdim: + result_data.ndim = self.ndim + result_data.shape = [1] * self.ndim + + else: + result_data.shape = [1] + result_data.ndim = 1 + else: + if keepdim: + result_data.shape = self.shape[:axis] + [1] + self.shape[axis+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 = MinBackward(self, axis, keepdim=keepdim) + + return result_data + + + def sin(self): + Tensor._C.sin_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.sin_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.sin_tensor(self.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.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = SinBackward(self) + + return result_data + + def cos(self): + Tensor._C.cos_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.cos_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.cos_tensor(self.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.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = CosBackward(self) + + return result_data + + def sigmoid(self): + Tensor._C.sigmoid_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.sigmoid_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.sigmoid_tensor(self.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.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = SigmoidBackward(self) + + return result_data + + + + def transpose(self, axis1, axis2): + if axis1 < 0: + axis1 = self.ndim + axis1 + if axis2 < 0: + axis2 = self.ndim + axis2 + + Tensor._C.transpose_axes_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_int] + Tensor._C.transpose_axes_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.transpose_axes_tensor(self.tensor, axis1, axis2) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = self.shape.copy() + result_data.shape[axis1] = self.shape[axis2] + result_data.shape[axis2] = self.shape[axis1] + result_data.ndim = self.ndim + result_data.device = self.device + result_data.numel = self.numel + + result_data.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = TransposeBackward(self, axis1, axis2) + + return result_data + + @property + def T(self): + Tensor._C.transpose_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.transpose_tensor.restype = ctypes.POINTER(CTensor) + + result_tensor_ptr = Tensor._C.transpose_tensor(self.tensor) + + result_data = Tensor() + result_data.tensor = result_tensor_ptr + result_data.shape = self.shape.copy()[::-1] + result_data.ndim = self.ndim + result_data.device = self.device + result_data.numel = self.numel + + result_data.requires_grad = self.requires_grad + if result_data.requires_grad: + result_data.grad_fn = TBackward(self) + + return result_data + + def detach(self): + self.grad = None + self.grad_fn = None + + return self \ No newline at end of file diff --git a/build/lib/norch/utils/__init__.py b/build/lib/norch/utils/__init__.py new file mode 100644 index 0000000..c6adfd1 --- /dev/null +++ b/build/lib/norch/utils/__init__.py @@ -0,0 +1,2 @@ +from .utils import * +from .data import * \ No newline at end of file diff --git a/build/lib/norch/utils/data/__init__.py b/build/lib/norch/utils/data/__init__.py new file mode 100644 index 0000000..2397554 --- /dev/null +++ b/build/lib/norch/utils/data/__init__.py @@ -0,0 +1,4 @@ +from .dataset import * +from .example import * +from .dataloader import * +from .batch import * \ No newline at end of file diff --git a/build/lib/norch/utils/data/batch.py b/build/lib/norch/utils/data/batch.py new file mode 100644 index 0000000..3c9b92c --- /dev/null +++ b/build/lib/norch/utils/data/batch.py @@ -0,0 +1,13 @@ +class Batch(object): + """ + A ``Batch``depends on a ``Dataset`` and is made of ``batch_size`` ``Example``. + """ + def __init__(self, example, batch_size): + self.example = example + self.batch_size = batch_size + + def __getitem__(self, item): + return self.example[item] + + def __len__(self): + return self.batch_size \ No newline at end of file diff --git a/build/lib/norch/utils/data/dataloader.py b/build/lib/norch/utils/data/dataloader.py new file mode 100644 index 0000000..01bd41e --- /dev/null +++ b/build/lib/norch/utils/data/dataloader.py @@ -0,0 +1,20 @@ +import numpy as np +from .batch import Batch + + +class Dataloader(object): + + def __init__(self, dataset, batch_size=32): + self.dataset = dataset + self.batch_size = batch_size + + def __iter__(self): + starts = np.arange(0, len(self.dataset), self.batch_size) + + for start in starts: + end = start + self.batch_size + batch_size = min(end, len(self.dataset)) - start + yield Batch(self.dataset[start:end], batch_size) + + def __len__(self): + return len(self.dataset) // self.batch_size diff --git a/build/lib/norch/utils/data/dataset.py b/build/lib/norch/utils/data/dataset.py new file mode 100644 index 0000000..d38c1ea --- /dev/null +++ b/build/lib/norch/utils/data/dataset.py @@ -0,0 +1,74 @@ +from abc import ABC, abstractmethod +import os +from norch.utils import extract_to_dir, download_from_url +from .example import Example +import norch + + +class Dataset(ABC): + r""" + Abstract Dataset class. All dataset for machine learning purposes can inherits from this architecture, + for convenience. + """ + urls = [] + name = '' + dirname = '' + + def __init__(self, examples, fields): + self.examples = examples + self.fields = fields + + @classmethod + def splits(cls, train=None, test=None, valid=None, root='.'): + raise NotImplementedError + + @classmethod + def download(cls, root): + r"""Download and unzip a web archive (.zip, .gz, or .tgz). + + Args: + root (str): Folder to download data to. + + Returns: + string: Path to extracted dataset. + """ + path_dirname = os.path.join(root, cls.dirname) + path_name = os.path.join(path_dirname, cls.name) + if not os.path.isdir(path_dirname): + for url in cls.urls: + filename = os.path.basename(url) + zpath = os.path.join(path_dirname, filename) + if not os.path.isfile(zpath): + if not os.path.exists(os.path.dirname(zpath)): + os.makedirs(os.path.dirname(zpath)) + print(f'Download {filename} from {url} to {zpath}') + download_from_url(url, zpath) + extract_to_dir(zpath, path_name) + + return path_name + + def __repr__(self): + name = self.__class__.__name__ + string = f"Dataset {name}(" + tab = " " + for (key, value) in self.__dict__.items(): + if key[0] != "_": + if isinstance(value, Example): + fields = self.fields + for (name, field) in fields: + string += f"\n{tab}({name}): {field.__class__.__name__}" \ + f"(transform={True if field.transform is not None else None}, dtype={field.dtype})" + elif isinstance(value, norch.Tensor): + string += f"\n{tab}({key}): {value.__class__.__name__}(shape={value.shape}, dtype={value.dtype})" + else: + string += f"\n{tab}({key}): {value.__class__.__name__}" + return f'{string}\n)' + + def __getitem__(self, item): + return self.examples[item] + + def __setitem__(self, key, value): + self.examples[key] = value + + def __len__(self): + return len(self.examples) diff --git a/build/lib/norch/utils/data/example.py b/build/lib/norch/utils/data/example.py new file mode 100644 index 0000000..a4f4ebf --- /dev/null +++ b/build/lib/norch/utils/data/example.py @@ -0,0 +1,65 @@ +# File: example.py +# Creation: Wednesday August 19th 2020 +# Author: Arthur Dujardin +# Contact: arthur.dujardin@ensg.eu +# arthurd@ifi.uio.no +# -------- +# Copyright (c) 2020 Arthur Dujardin + + +class Field(object): + r""" + A ``Field`` defines the data to process from a raw dataset. It will convert the data into a tensor. The data can + be a string, integers, float etc. The data is meant to be preprocessed and the attribute ``transform`` handles + the way the user want to process the raw data. + """ + + def __init__(self, transform=None, dtype=None): + self.transform = transform + self.dtype = dtype + + def process(self, value): + """Applies the transformation and changes the data's type if necessary. + + Args: + value (Tensor): + + Returns: + + """ + if self.transform is not None: + value = self.transform(value) + if self.dtype is not None: + value = value.astype(self.dtype) + return value + + +class Example(object): + r""" + Store a single training / testing example, and store it as an attribute. + Highly inspired from PyTorch `Example `__. + + """ + @classmethod + def fromlist(cls, values, fields): + """ + Add an example from a list of data with respect to the fields. + + Args: + values: raw data + fields (tuple(string, Field)): fields to preprocess the data on + + Returns: + None + """ + example = cls() + for (value, field) in zip(values, fields): + assert len(field) == 2, f"expected a field template similar to \ + ('name_field', Field()) but got {format(field)}" + assert isinstance(field[1], Field), f"expected a field template similar to \ + ('name_field', Field()) but got {format(field)}" + name = field[0] + value = field[1].process(value) + setattr(example, name, value) + + return example diff --git a/build/lib/norch/utils/utils.py b/build/lib/norch/utils/utils.py new file mode 100644 index 0000000..f8b9b30 --- /dev/null +++ b/build/lib/norch/utils/utils.py @@ -0,0 +1,120 @@ +import random +import requests +import tarfile +import zipfile +import shutil +import os + +def generate_random_list(shape): + """ + Generate a list with random numbers and shape 'shape' + """ + if len(shape) == 0: + return [] + else: + inner_shape = shape[1:] + if len(inner_shape) == 0: + return [random.uniform(-1, 1) for _ in range(shape[0])] + else: + return [generate_random_list(inner_shape) for _ in range(shape[0])] + + +def download_from_url(url, save_path, chunk_size=128): + """Download a file from an URL. + + Original answer from https://stackoverflow.com/questions/9419162/download-returned-zip-file-from-url + + Args: + url (str): path to the URL. + save_path (str): path to the saving directory. + chunk_size (int): download chunk. + + Returns: + None + """ + response = requests.get(url, stream=True) + total = response.headers.get('content-length') + with open(save_path, 'wb') as f: + if total is None: + f.write(response.content) + else: + downloaded = 0 + total = int(total) + for data in response.iter_content(chunk_size=max(int(total / 1000), 1024 * 1024)): + downloaded += len(data) + f.write(data) + progress_bar(downloaded, total, "Downloading...") + + +def extract_to_dir(filename, dirpath='.'): + # Does not create folder twice with the same name + name, ext = os.path.splitext(filename) + # if os.path.basename(name) == os.path.basename(dirpath): + # dirpath = '.' + # Extract + print(dirpath) + print("Extracting...", end="") + if tarfile.is_tarfile(filename): + tarfile.open(filename, 'r').extractall(dirpath) + elif zipfile.is_zipfile(filename): + zipfile.ZipFile(filename, 'r').extractall(dirpath) + elif ext == '.gz': + if not os.path.exists(dirpath): + os.mkdir(dirpath) + shutil.move(filename, os.path.join(dirpath, os.path.basename(filename))) + print(f" | NOTE: gzip files are not extracted, and moved to {dirpath}", end="") + # Return the path where the file was extracted + print(" | Done !") + return os.path.abspath(dirpath) + +def progress_bar(current_index, max_index, prefix=None, suffix=None, start_time=None): + """Display a progress bar and duration. + + Args: + current_index (int): current state index (or epoch number). + max_index (int): maximal numbers of state. + prefix (str, optional): prefix of the progress bar. The default is None. + suffix (str, optional): suffix of the progress bar. The default is None. + start_time (float, optional): starting time of the progress bar. If not None, it will display the time + spent from the beginning to the current state. The default is None. + + Returns: + None. Display the progress bar in the console. + """ + # Add a prefix to the progress bar + prefix = "" if prefix is None else str(prefix) + " " + + # Get the percentage + percentage = current_index * 100 // max_index + loading = "[" + "=" * (percentage // 2) + " " * (50 - percentage // 2) + "]" + progress_display = "\r{0}{1:3d}% | {2}".format(prefix, percentage, loading) + + # Add a suffix to the progress bar + progress_display += "" if suffix is None else " | " + str(suffix) + + # Add a timer + if start_time is not None: + time_min, time_sec = get_time(start_time, time.time()) + time_display = " | Time: {0}m {1}s".format(time_min, time_sec) + progress_display += time_display + + # Print the progress bar + # TODO: return a string instead + print(progress_display, end="{}".format("" if current_index < max_index else " | Done !\n")) + +def get_time(start_time, end_time): + """Get ellapsed time in minutes and seconds. + + Args: + start_time (float): strarting time + end_time (float): ending time + + Returns: + elapsed_mins (float): elapsed time in minutes + elapsed_secs (float): elapsed time in seconds. + """ + elapsed_time = end_time - start_time + elapsed_mins = int(elapsed_time / 60) + elapsed_secs = int(elapsed_time - (elapsed_mins * 60)) + + return elapsed_mins, elapsed_secs diff --git a/build/lib/norch/utils/utils_unittests.py b/build/lib/norch/utils/utils_unittests.py new file mode 100644 index 0000000..4e0a12e --- /dev/null +++ b/build/lib/norch/utils/utils_unittests.py @@ -0,0 +1,24 @@ +import torch + +def to_torch(custom_tensor): + shape = custom_tensor.shape + pytorch_tensor = torch.zeros(shape) + + def _iterate_indices(shape): + if len(shape) == 0: + yield () + else: + for index in range(shape[0]): + for sub_indices in _iterate_indices(shape[1:]): + yield (index,) + sub_indices + + # Iterate over all elements using the custom tensor's __getitem__ method + for indices in _iterate_indices(shape): + value = custom_tensor[indices] + pytorch_tensor[tuple(indices)] = value + + return pytorch_tensor + +def compare_torch(tensor1, tensor2, epsilon=1e-5): + diff = torch.abs(tensor1 - tensor2) + return torch.all(diff < epsilon) \ No newline at end of file diff --git a/build/lib/tests/__init__.py b/build/lib/tests/__init__.py new file mode 100644 index 0000000..4750076 --- /dev/null +++ b/build/lib/tests/__init__.py @@ -0,0 +1,3 @@ +from .test_operations import * +from .test_autograd import * +from .test_nn import * \ No newline at end of file diff --git a/build/lib/tests/test_autograd.py b/build/lib/tests/test_autograd.py new file mode 100644 index 0000000..3be772c --- /dev/null +++ b/build/lib/tests/test_autograd.py @@ -0,0 +1,993 @@ +import unittest +import norch +from norch.utils import utils_unittests as utils +import torch +import os + +class TestTensorAutograd(unittest.TestCase): + def setUp(self): + self.device = os.environ.get('device') + if self.device is None or self.device != 'cuda': + self.device = 'cpu' + + print(f"Running tests on: {self.device}") + + + def test_addition(self): + """ + Test autograd from addition two tensors: tensor1 + tensor2 + """ + 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() + 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, device=self.device) + torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True, device=self.device) + 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_sum_axis(self): + """ + Test autograd from sum specifying axis + """ + 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, device=self.device) + torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True, device=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)) + + 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=1).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, device=self.device) + torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True, device=self.device) + + torch_result = (torch_tensor1 + torch_tensor2).sum(axis=1).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_max(self): + """ + Test autograd from max + """ + norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + norch_result = norch_tensor.max() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True, device=self.device) + + torch_result = torch_tensor.max() + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + def test_max_axis(self): + """ + Test autograd from max specifying axis + """ + norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + norch_max_axis = norch_tensor.max(axis=1) + norch_result = norch_max_axis.sum() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True, device=self.device) + + torch_max_axis, _ = torch_tensor.max(axis=1) + torch_result = torch_max_axis.sum() + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + norch_tensor = norch.Tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + norch_max_axis = norch_tensor.max(axis=2) + norch_result = norch_max_axis.sum() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True, device=self.device) + + torch_max_axis, _ = torch_tensor.max(axis=2) + torch_result = torch_max_axis.sum() + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + + ## evaluate case with some repeated values and axis 2 + + #def test_max_axis(self): + # """ + # Test autograd from max specifying axis + # """ + # + # norch_tensor = norch.Tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + # norch_max_axis = norch_tensor.max(axis=2) + # norch_result = norch_max_axis.sum() + # + # norch_result.backward() + # norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + # + # torch_tensor = torch.tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + # + # torch_max_axis, _ = torch_tensor.max(axis=2) + # torch_result = torch_max_axis.sum() + # torch_result.backward() + # torch_tensor_grad = torch_tensor.grad + # self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + + #def test_min_axis(self): + # """ + # Test autograd from max specifying axis + # """ + # + # norch_tensor = norch.Tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + # norch_min_axis = norch_tensor.min(axis=2) + # norch_result = norch_min_axis.sum() + # + # norch_result.backward() + # norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + # + # torch_tensor = torch.tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + # + # torch_min_axis, _ = torch_tensor.min(axis=2) + # torch_result = torch_min_axis.sum() + # torch_result.backward() + # torch_tensor_grad = torch_tensor.grad + # self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + + + def test_min(self): + """ + Test autograd from min + """ + norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + norch_result = norch_tensor.min() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True, device=self.device) + + torch_result = torch_tensor.min() + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + def test_min_axis(self): + """ + Test autograd from min specifying axis + """ + norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + norch_min = norch_tensor.min(axis=1) + norch_result = norch_min.sum() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True, device=self.device) + + torch_min, _ = torch_tensor.min(axis=1) + torch_result = torch_min.sum() + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + norch_tensor = norch.Tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + norch_min = norch_tensor.min(axis=2) + norch_result = norch_min.sum() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True, device=self.device) + + torch_min, _ = torch_tensor.min(axis=2) + torch_result = torch_min.sum() + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + + def test_broadcasted_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).to(self.device) + norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device) + + torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True, device=self.device) # Shape (1, 2, 3) + torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True, device=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)) + + ## reversed order broadcasting + 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_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_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True, device=self.device) # Shape (1, 2, 3) + torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True, device=self.device) # Shape (3) + + 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): + """ + Test autograd from subtraction two tensors: tensor1 - tensor2 + """ + norch_tensor1_sub = norch.Tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device) + norch_tensor2_sub = norch.Tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device) + norch_result_sub = (norch_tensor1_sub - norch_tensor2_sub).sum() + norch_result_sub.backward() + norch_tensor1_grad_sub = utils.to_torch(norch_tensor1_sub.grad).to(self.device) + norch_tensor2_grad_sub = utils.to_torch(norch_tensor2_sub.grad).to(self.device) + + torch_tensor1_sub = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_tensor2_sub = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True, device=self.device) + torch_result_sub = (torch_tensor1_sub - torch_tensor2_sub).sum() + torch_result_sub.backward() + torch_tensor1_grad_sub = torch_tensor1_sub.grad + torch_tensor2_grad_sub = torch_tensor2_sub.grad + + 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_broadcasted_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).to(self.device) + norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device) + + torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True, device=self.device) # Shape (1, 2, 3) + torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True, device=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)) + + # reversed order broadcasting + 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_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_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True, device=self.device) # Shape (1, 2, 3) + torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True, device=self.device) # Shape (3) + + 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_division(self): + """ + Test autograd from dividing two tensors: tensor1 / tensor2 + """ + norch_tensor1_div = norch.Tensor([[[2, 5.1], [6, -8]], [[10, 12], [14, 16]]], requires_grad=True).to(self.device) + norch_tensor2_div = norch.Tensor([[[1, 1], [2, 2.2]], [[3, 3], [4, 4]]], requires_grad=True).to(self.device) + norch_result_div = (norch_tensor1_div / norch_tensor2_div).sum() + norch_result_div.backward() + norch_tensor1_grad_div = utils.to_torch(norch_tensor1_div.grad).to(self.device) + norch_tensor2_grad_div = utils.to_torch(norch_tensor2_div.grad).to(self.device) + + torch_tensor1_div = torch.tensor([[[2, 5.1], [6, -8]], [[10, 12], [14, 16]]], requires_grad=True, device=self.device) + torch_tensor2_div = torch.tensor([[[1, 1], [2, 2.2]], [[3, 3], [4, 4]]], requires_grad=True, device=self.device) + torch_result_div = (torch_tensor1_div / torch_tensor2_div).sum() + torch_result_div.backward() + torch_tensor1_grad_div = torch_tensor1_div.grad + torch_tensor2_grad_div = torch_tensor2_div.grad + + self.assertTrue(utils.compare_torch(norch_tensor1_grad_div, torch_tensor1_grad_div)) + self.assertTrue(utils.compare_torch(norch_tensor2_grad_div, torch_tensor2_grad_div)) + + + def test_tensor_division_scalar(self): + """ + Test autograd from dividing tensor by scalar: tensor / scalar + """ + norch_tensor_div_scalar = norch.Tensor([[[2, 4.7], [6, 8]], [[10, 12], [14, 16]]], requires_grad=True).to(self.device) + scalar = 2 + norch_result_div_scalar = (norch_tensor_div_scalar / scalar).sum() + norch_result_div_scalar.backward() + norch_tensor_grad_div_scalar = utils.to_torch(norch_tensor_div_scalar.grad).to(self.device) + + torch_tensor_div_scalar = torch.tensor([[[2, 4.7], [6, 8]], [[10, 12], [14, 16]]], requires_grad=True, device=self.device) + torch_result_div_scalar = (torch_tensor_div_scalar / scalar).sum() + torch_result_div_scalar.backward() + torch_tensor_grad_div_scalar = torch_tensor_div_scalar.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_div_scalar, torch_tensor_grad_div_scalar)) + + + def test_scalar_division_tensor(self): + """ + Test autograd from dividing scalar by tensor: scalar / tensor + """ + scalar = 2 + norch_tensor_scalar_div = norch.Tensor([[[1, 2.23], [3, 4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device) + norch_result_scalar_div = (scalar / norch_tensor_scalar_div).sum() + norch_result_scalar_div.backward() + norch_tensor_grad_scalar_div = utils.to_torch(norch_tensor_scalar_div.grad).to(self.device) + + torch_tensor_scalar_div = torch.tensor([[[1, 2.23], [3, 4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_result_scalar_div = (scalar / torch_tensor_scalar_div).sum() + torch_result_scalar_div.backward() + torch_tensor_grad_scalar_div = torch_tensor_scalar_div.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_scalar_div, torch_tensor_grad_scalar_div)) + + + def test_power_scalar_tensor(self): + """ + Test autograd from scalar raised to tensor: scalar ** tensor + """ + scalar = 2 + norch_tensor_power_st = norch.Tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device) + norch_result_power_st = (scalar ** norch_tensor_power_st).sum() + norch_result_power_st.backward() + norch_tensor_grad_power_st = utils.to_torch(norch_tensor_power_st.grad).to(self.device) + + torch_tensor_power_st = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True, device=self.device) + torch_result_power_st = (scalar ** torch_tensor_power_st).sum() + torch_result_power_st.backward() + torch_tensor_grad_power_st = torch_tensor_power_st.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_power_st, torch_tensor_grad_power_st)) + + def test_power_tensor_scalar(self): + """ + Test autograd from tensor raised to scalar: tensor ** scalar + """ + scalar = 2 + norch_tensor_power_ts = norch.Tensor([[[2, 3], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device) + norch_result_power_ts = (norch_tensor_power_ts ** scalar).sum() + norch_result_power_ts.backward() + norch_tensor_grad_power_ts = utils.to_torch(norch_tensor_power_ts.grad).to(self.device) + + torch_tensor_power_ts = torch.tensor([[[2, 3], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True, device=self.device) + torch_result_power_ts = (torch_tensor_power_ts ** scalar).sum() + torch_result_power_ts.backward() + torch_tensor_grad_power_ts = torch_tensor_power_ts.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_power_ts, torch_tensor_grad_power_ts)) + + def test_matmul(self): + """ + Test autograd from matrix multiplication: matmul(tensor1, tensor2) + """ + norch_tensor1_matmul = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device) + norch_tensor2_matmul = norch.Tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device) + norch_result_matmul = (norch_tensor1_matmul @ norch_tensor2_matmul).sum() + norch_result_matmul.backward() + norch_tensor1_grad_matmul = utils.to_torch(norch_tensor1_matmul.grad).to(self.device) + norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad).to(self.device) + + torch_tensor1_matmul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_tensor2_matmul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True, device=self.device) + torch_result_matmul = (torch_tensor1_matmul @ torch_tensor2_matmul).sum() + torch_result_matmul.backward() + torch_tensor1_grad_matmul = torch_tensor1_matmul.grad + torch_tensor2_grad_matmul = torch_tensor2_matmul.grad + + self.assertTrue(utils.compare_torch(norch_tensor1_grad_matmul, torch_tensor1_grad_matmul)) + self.assertTrue(utils.compare_torch(norch_tensor2_grad_matmul, torch_tensor2_grad_matmul)) + + def test_batched_matmul(self): + """ + Test autograd from batched matrix multiplication: BxMxP = BxNxM @ BxMxP + """ + B = 3 # Batch size + + norch_tensor1_matmul = norch.Tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)], requires_grad=True).to(self.device) + norch_tensor2_matmul = norch.Tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(self.device) + + norch_result_matmul = norch_tensor1_matmul @ norch_tensor2_matmul + norch_result_matmul_sum = norch_result_matmul.sum() # Sum over all elements + norch_result_matmul_sum.backward() + + # Convert gradients to torch tensors + norch_tensor1_grad_matmul = utils.to_torch(norch_tensor1_matmul.grad).to(self.device) + norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad).to(self.device) + + # Repeat the same process with torch tensors + torch_tensor1_matmul = torch.tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)], requires_grad=True, device=self.device) + torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True, device=self.device) + torch_result_matmul = torch.matmul(torch_tensor1_matmul, torch_tensor2_matmul) + torch_result_matmul_sum = torch_result_matmul.sum() + torch_result_matmul_sum.backward() + + # Extract gradients from torch tensors + torch_tensor1_grad_matmul = torch_tensor1_matmul.grad + torch_tensor2_grad_matmul = torch_tensor2_matmul.grad + + # Assertions to compare the gradients + self.assertTrue(utils.compare_torch(norch_tensor1_grad_matmul, torch_tensor1_grad_matmul)) + self.assertTrue(utils.compare_torch(norch_tensor2_grad_matmul, torch_tensor2_grad_matmul)) + + def test_broadcasted_batched_matmul(self): + """ + Test autograd from broadcasted batched matrix multiplication: BxMxP = NxM @ BxMxP + """ + B = 3 # Batch size + + norch_tensor1_matmul = norch.Tensor([[1., 2], [3, -4], [5, 6], [7, 8]], requires_grad=True).to(self.device) + norch_tensor2_matmul = norch.Tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(self.device) + + norch_result_matmul = norch_tensor1_matmul @ norch_tensor2_matmul + norch_result_matmul_sum = norch_result_matmul.sum() # Sum over all elements + norch_result_matmul_sum.backward() + + # Convert gradients to torch tensors + norch_tensor1_grad_matmul = utils.to_torch(norch_tensor1_matmul.grad).to(self.device) + norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad).to(self.device) + + # Repeat the same process with torch tensors + torch_tensor1_matmul = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]], requires_grad=True, device=self.device) + torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True, device=self.device) + torch_result_matmul = torch.matmul(torch_tensor1_matmul, torch_tensor2_matmul) + torch_result_matmul_sum = torch_result_matmul.sum() + torch_result_matmul_sum.backward() + + # Extract gradients from torch tensors + torch_tensor1_grad_matmul = torch_tensor1_matmul.grad + torch_tensor2_grad_matmul = torch_tensor2_matmul.grad + + # Assertions to compare the gradients + self.assertTrue(utils.compare_torch(norch_tensor1_grad_matmul, torch_tensor1_grad_matmul)) + self.assertTrue(utils.compare_torch(norch_tensor2_grad_matmul, torch_tensor2_grad_matmul)) + + + def test_elementwise_mul_scalar(self): + """ + Test autograd from elementwise multiplication with scalar: scalar * tensor + """ + scalar = 2 + norch_tensor_elemwise_mul_scalar = norch.Tensor([[[1.1, 2], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device) + norch_result_elemwise_mul_scalar = (scalar * norch_tensor_elemwise_mul_scalar).sum() + norch_result_elemwise_mul_scalar.backward() + norch_tensor_grad_elemwise_mul_scalar = utils.to_torch(norch_tensor_elemwise_mul_scalar.grad).to(self.device) + + torch_tensor_elemwise_mul_scalar = torch.tensor([[[1.1, 2], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_result_elemwise_mul_scalar = (scalar * torch_tensor_elemwise_mul_scalar).sum() + torch_result_elemwise_mul_scalar.backward() + torch_tensor_grad_elemwise_mul_scalar = torch_tensor_elemwise_mul_scalar.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_elemwise_mul_scalar, torch_tensor_grad_elemwise_mul_scalar)) + + + def test_elementwise_mul_tensor(self): + """ + Test autograd from elementwise multiplication between two tensors: tensor1 * tensor2 + """ + norch_tensor1_elemwise_mul = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device) + norch_tensor2_elemwise_mul = norch.Tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device) + norch_result_elemwise_mul = (norch_tensor1_elemwise_mul * norch_tensor2_elemwise_mul).sum() + norch_result_elemwise_mul.backward() + norch_tensor1_grad_elemwise_mul = utils.to_torch(norch_tensor1_elemwise_mul.grad).to(self.device) + norch_tensor2_grad_elemwise_mul = utils.to_torch(norch_tensor2_elemwise_mul.grad).to(self.device) + + torch_tensor1_elemwise_mul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_tensor2_elemwise_mul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True, device=self.device) + torch_result_elemwise_mul = (torch_tensor1_elemwise_mul * torch_tensor2_elemwise_mul).sum() + torch_result_elemwise_mul.backward() + torch_tensor1_grad_elemwise_mul = torch_tensor1_elemwise_mul.grad + torch_tensor2_grad_elemwise_mul = torch_tensor2_elemwise_mul.grad + + self.assertTrue(utils.compare_torch(norch_tensor1_grad_elemwise_mul, torch_tensor1_grad_elemwise_mul)) + self.assertTrue(utils.compare_torch(norch_tensor2_grad_elemwise_mul, torch_tensor2_grad_elemwise_mul)) + + def test_sin_tensor(self): + """ + Test autograd from sin operation: sin(tensor) + """ + norch_sin_tensor = norch.Tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device) + norch_result_sin_tensor = (norch_sin_tensor.sin()).sum() + norch_result_sin_tensor.backward() + torch_result_sin_tensor_grad = utils.to_torch(norch_sin_tensor.grad).to(self.device) + + torch_sin_tensor = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True, device=self.device) + torch_expected_sin_tensor = (torch.sin(torch_sin_tensor)).sum() + torch_expected_sin_tensor.backward() + torch_expected_sin_tensor_grad = torch_sin_tensor.grad + + self.assertTrue(utils.compare_torch(torch_result_sin_tensor_grad, torch_expected_sin_tensor_grad)) + + def test_cos_tensor(self): + """ + Test autograd from cosine operation: cos(tensor) + """ + norch_cos_tensor = norch.Tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device) + norch_result_cos_tensor = (norch_cos_tensor.sin()).sum() + norch_result_cos_tensor.backward() + torch_result_cos_tensor_grad = utils.to_torch(norch_cos_tensor.grad).to(self.device) + + torch_cos_tensor = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True, device=self.device) + torch_expected_cos_tensor = (torch.sin(torch_cos_tensor)).sum() + torch_expected_cos_tensor.backward() + torch_expected_cos_tensor_grad = torch_cos_tensor.grad + + self.assertTrue(utils.compare_torch(torch_result_cos_tensor_grad, torch_expected_cos_tensor_grad)) + + def test_sigmoid(self): + """ + Test autograd from sigmoid + """ + norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + norch_sigmoid = norch.sigmoid(norch_tensor) + norch_result = norch_sigmoid.sum() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True, device=self.device) + torch_sigmoid = torch.sigmoid(torch_tensor) + torch_result = torch_sigmoid.sum() + + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + def test_mse_loss_autograd(self): + """ + Test the MSELoss with autograd functionality + """ + loss_fn_norch = norch.nn.MSELoss() + loss_fn_torch = torch.nn.MSELoss() + + predictions_norch = norch.Tensor([1.1, 2, 3, 4], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([4, 3, 2.1, 1]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([1.1, 2, 3, 4], requires_grad=True, device=self.device) + labels_torch = torch.tensor([4, 3, 2.1, 1], device=self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + def test_cross_entropy_loss_autograd(self): + """ + Test the CrossEntropyLoss with autograd functionality + """ + loss_fn_norch = norch.nn.CrossEntropyLoss() + loss_fn_torch = torch.nn.CrossEntropyLoss() + + # Test case 1: Single class, single sample + predictions_norch = norch.Tensor([2.0, 1.0, 0.1], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([0]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([2.0, 1.0, 0.1], requires_grad=True, device=self.device) + labels_torch = torch.tensor(0, device=self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + # Test case 2: Multiple classes, multiple samples + predictions_norch = norch.Tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([2, 1]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]], requires_grad=True, device=self.device) + labels_torch = torch.tensor([2, 1], device=self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + # Test case 3: Edge case - all predictions are zero + predictions_norch = norch.Tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([1, 2]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], requires_grad=True, device=self.device) + labels_torch = torch.tensor([1, 2], device=self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + # Test case 4: Batched class probabilities instead of class index + predictions_norch = norch.Tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([[1., 0, 0], [0, 1, 0]]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]], requires_grad=True, device=self.device) + labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]], device=self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + + # implement grad pure softmax --> 0 + # def test_softmax(self): + # """ + # Test autograd from softmax + # """ + # norch_tensor = norch.Tensor([[[-5, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + # norch_softmax = norch.softmax(norch_tensor, dim=1) + # norch_result = norch_softmax.sum() + + # norch_result.backward() + # norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + # torch_tensor = torch.tensor([[[-5, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + # torch_softmax = torch.softmax(torch_tensor, dim=1) + # torch_result = torch_softmax.sum() + + # torch_result.backward() + # torch_tensor_grad = torch_tensor.grad + + # self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + # norch_tensor = norch.Tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + # norch_softmax = norch.softmax(norch_tensor, dim=2) + # norch_result = norch_softmax.sum() + + # norch_result.backward() + # norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + # torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + + # torch_softmax = torch.softmax(torch_tensor, dim=2) + # torch_result = torch_softmax.sum() + # torch_result.backward() + # torch_tensor_grad = torch_tensor.grad + + # self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + + def test_reshape(self): + """ + Test autograd from reshaping a tensor: tensor.reshape(shape) + """ + norch_tensor_reshape = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device) + new_shape = [2, 4] + norch_result_reshape = norch_tensor_reshape.reshape(new_shape).sum() + norch_result_reshape.backward() + norch_tensor_grad_reshape = utils.to_torch(norch_tensor_reshape.grad).to(self.device) + + torch_tensor_reshape = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_result_reshape = torch_tensor_reshape.reshape(new_shape).sum() + torch_result_reshape.backward() + torch_tensor_grad_reshape = torch_tensor_reshape.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_reshape, torch_tensor_grad_reshape)) + + + def test_transpose_axes(self): + """ + Test autograd from transposing a tensor with specific axes: tensor.transpose(axis1, axis2) + """ + norch_tensor_transpose = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device) + axis1, axis2 = 0, 2 + norch_result_transpose = norch_tensor_transpose.transpose(axis1, axis2).sum() + norch_result_transpose.backward() + norch_tensor_grad_transpose = utils.to_torch(norch_tensor_transpose.grad).to(self.device) + + torch_tensor_transpose = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_result_transpose = torch_tensor_transpose.transpose(axis1, axis2).sum() + torch_result_transpose.backward() + torch_tensor_grad_transpose = torch_tensor_transpose.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_transpose, torch_tensor_grad_transpose)) + + + def test_T(self): + """ + Test autograd from transposing a tensor using .T attribute + """ + norch_tensor_T = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device) + norch_result_T = norch_tensor_T.T.sum() + norch_result_T.backward() + norch_tensor_grad_T = utils.to_torch(norch_tensor_T.grad).to(self.device) + + torch_tensor_T = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True, device=self.device) + torch_result_T = torch_tensor_T.mT.sum() + torch_result_T.backward() + torch_tensor_grad_T = torch_tensor_T.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_T, torch_tensor_grad_T)) + + def test_reshape_then_matmul(self): + """ + Test autograd from reshaping a tensor then performing matrix multiplication: matmul(tensor1.reshape(shape), tensor2) + """ + norch_tensor1 = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True).to(self.device) + norch_tensor2 = norch.Tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], requires_grad=True).to(self.device) + + new_shape = [2, 4] + + norch_result_reshape_matmul = (norch_tensor1.reshape(new_shape) @ norch_tensor2).sum() + norch_result_reshape_matmul.backward() + norch_tensor_grad_reshape_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device) + norch_tensor_grad_reshape_matmul2 = utils.to_torch(norch_tensor2.grad).to(self.device) + + torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True, device=self.device) + torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True, device=self.device) + + torch_result_reshape_matmul = (torch_tensor1.reshape(new_shape) @ torch_tensor2).sum() + torch_result_reshape_matmul.backward() + torch_tensor_grad_reshape_matmul1 = torch_tensor1.grad + torch_tensor_grad_reshape_matmul2 = torch_tensor2.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_reshape_matmul1, torch_tensor_grad_reshape_matmul1)) + self.assertTrue(utils.compare_torch(norch_tensor_grad_reshape_matmul2, torch_tensor_grad_reshape_matmul2)) + + def test_unsqueeze(self): + """ + Test autograd from unsqueezing a tensor: tensor.unsqueeze(dim) + """ + + # Unsqueeze at dim=0 + norch_tensor_unsqueeze = norch.Tensor([[1., 2], [3, 4]], requires_grad=True).to(self.device) + norch_result_unsqueeze_0 = norch_tensor_unsqueeze.unsqueeze(0).sum() + norch_result_unsqueeze_0.backward() + norch_tensor_grad_unsqueeze_0 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device) + + torch_tensor_unsqueeze = torch.tensor([[1., 2], [3, 4]], requires_grad=True, device=self.device) + torch_result_unsqueeze_0 = torch_tensor_unsqueeze.unsqueeze(0).sum() + torch_result_unsqueeze_0.backward() + torch_tensor_grad_unsqueeze_0 = torch_tensor_unsqueeze.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_0, torch_tensor_grad_unsqueeze_0)) + + # Unsqueeze at dim=1 + norch_tensor_unsqueeze = norch.Tensor([[1., 2.], [3, 4]], requires_grad=True).to(self.device) + norch_result_unsqueeze_1 = norch_tensor_unsqueeze.unsqueeze(1).sum() + norch_result_unsqueeze_1.backward() + norch_tensor_grad_unsqueeze_1 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device) + + torch_tensor_unsqueeze = torch.tensor([[1., 2.], [3, 4]], requires_grad=True, device=self.device) + torch_result_unsqueeze_1 = torch_tensor_unsqueeze.unsqueeze(1).sum() + torch_result_unsqueeze_1.backward() + torch_tensor_grad_unsqueeze_1 = torch_tensor_unsqueeze.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_1, torch_tensor_grad_unsqueeze_1)) + + # Unsqueeze at dim=2 + norch_tensor_unsqueeze = norch.Tensor([[1., 2], [3, 4]], requires_grad=True).to(self.device) + norch_result_unsqueeze_2 = norch_tensor_unsqueeze.unsqueeze(2).sum() + norch_result_unsqueeze_2.backward() + norch_tensor_grad_unsqueeze_2 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device) + + torch_tensor_unsqueeze = torch.tensor([[1., 2], [3, 4]], requires_grad=True, device=self.device) + torch_result_unsqueeze_2 = torch_tensor_unsqueeze.unsqueeze(2).sum() + torch_result_unsqueeze_2.backward() + torch_tensor_grad_unsqueeze_2 = torch_tensor_unsqueeze.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_2, torch_tensor_grad_unsqueeze_2)) + + def test_unsqueeze_then_matmul(self): + """ + Test autograd from unsqueezing a tensor then performing matrix multiplication: matmul(tensor1.unsqueeze(dim), tensor2) + """ + norch_tensor1 = norch.Tensor([[1, 2], [3, 4]], requires_grad=True).to(self.device) + norch_tensor2 = norch.Tensor([[1, 2], [3, 4]], requires_grad=True).to(self.device) + + # Unsqueeze at dim=0 then matmul + norch_result_unsqueeze_matmul = (norch_tensor1 @ norch_tensor2.unsqueeze(0)).sum() + norch_result_unsqueeze_matmul.backward() + norch_tensor_grad_unsqueeze_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device) + norch_tensor_grad_unsqueeze_matmul2 = utils.to_torch(norch_tensor2.grad).to(self.device) + + torch_tensor1 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True, device=self.device) + torch_tensor2 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True, device=self.device) + + torch_result_unsqueeze_matmul = (torch_tensor1 @ torch_tensor2.unsqueeze(0)).sum() + torch_result_unsqueeze_matmul.backward() + torch_tensor_grad_unsqueeze_matmul1 = torch_tensor1.grad + torch_tensor_grad_unsqueeze_matmul2 = torch_tensor2.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_matmul1, torch_tensor_grad_unsqueeze_matmul1)) + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_matmul2, torch_tensor_grad_unsqueeze_matmul2)) + + def test_squeeze(self): + """ + Test autograd from squeezing a tensor: tensor.squeeze(dim) + """ + # Squeeze at dim=0 + norch_tensor_squeeze = norch.Tensor([[[1., 2], [3, 4]]], requires_grad=True).to(self.device) + norch_result_squeeze_0 = norch_tensor_squeeze.squeeze(0).sum() + norch_result_squeeze_0.backward() + norch_tensor_grad_squeeze_0 = utils.to_torch(norch_tensor_squeeze.grad).to(self.device) + + torch_tensor_squeeze = torch.tensor([[[1., 2], [3, 4]]], requires_grad=True, device=self.device) + torch_result_squeeze_0 = torch_tensor_squeeze.squeeze(0).sum() + torch_result_squeeze_0.backward() + torch_tensor_grad_squeeze_0 = torch_tensor_squeeze.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_squeeze_0, torch_tensor_grad_squeeze_0)) + + def test_squeeze_then_matmul(self): + """ + Test autograd from squeezing a tensor then performing matrix multiplication: matmul(tensor1.squeeze(dim), tensor2) + """ + norch_tensor1 = norch.Tensor([[[1., 2], [3, 4]]], requires_grad=True).to(self.device) + norch_tensor2 = norch.Tensor([[[1., 2], [3, 4]]], requires_grad=True).to(self.device) + + # Squeeze at dim=0 then matmul + norch_result_squeeze_matmul = (norch_tensor1.squeeze(0) @ norch_tensor2).sum() + norch_result_squeeze_matmul.backward() + norch_tensor_grad_squeeze_matmul1 = utils.to_torch(norch_tensor1.grad,).to(self.device) + norch_tensor_grad_squeeze_matmul2 = utils.to_torch(norch_tensor2.grad).to(self.device) + + torch_tensor1 = torch.tensor([[[1., 2], [3, 4]]], dtype=torch.float32, requires_grad=True, device=self.device) + torch_tensor2 = torch.tensor([[[1., 2], [3, 4]]], dtype=torch.float32, requires_grad=True, device=self.device) + + torch_result_squeeze_matmul = (torch_tensor1.squeeze(0) @ torch_tensor2).sum() + torch_result_squeeze_matmul.backward() + torch_tensor_grad_squeeze_matmul1 = torch_tensor1.grad + torch_tensor_grad_squeeze_matmul2 = torch_tensor2.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_squeeze_matmul1, torch_tensor_grad_squeeze_matmul1)) + self.assertTrue(utils.compare_torch(norch_tensor_grad_squeeze_matmul2, torch_tensor_grad_squeeze_matmul2)) + + + def test_T_then_matmul(self): + """ + Test autograd from transposing a tensor then performing matrix multiplication: matmul(tensor.T, tensor) + """ + norch_tensor1 = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True) + norch_tensor2 = norch.Tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], requires_grad=True) + + norch_result_T_matmul = (norch_tensor1.T @ norch_tensor2).sum() + norch_result_T_matmul.backward() + norch_tensor_grad_T_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device) + norch_tensor_grad_T_matmult2 = utils.to_torch(norch_tensor2.grad).to(self.device) + + torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True, device=self.device) + torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True, device=self.device) + + torch_result_T_matmul = (torch_tensor1.T @ torch_tensor2).sum() + torch_result_T_matmul.backward() + torch_tensor_grad_T_matmul1 = torch_tensor1.grad + torch_tensor_grad_T_matmul2 = torch_tensor2.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_T_matmul1, torch_tensor_grad_T_matmul1)) + self.assertTrue(utils.compare_torch(norch_tensor_grad_T_matmult2, torch_tensor_grad_T_matmul2)) + + def todo(self): + """ + The code has a problem on the following operation + tensor1.reshape(..) @ tensor1 + print(tensor1.grad) + (also transpsoe and .T) + """ + pass + + def test_transpose_axes_then_matmul(self): + """ + Test autograd from transposing a tensor with specific axes then performing matrix multiplication: matmul(tensor.transpose(axis1, axis2), tensor) + """ + norch_tensor1 = norch.Tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], requires_grad=True).to(self.device) + norch_tensor2 = norch.Tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], requires_grad=True).to(self.device) + + norch_result_transpose_matmul = (norch_tensor1.transpose(0, 1) @ norch_tensor2).sum() + norch_result_transpose_matmul.backward() + norch_tensor_grad_transpose_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device) + norch_tensor_grad_transpose_matmult2 = utils.to_torch(norch_tensor2.grad).to(self.device) + + torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True, device=self.device) + torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True, device=self.device) + + torch_result_transpose_matmul = (torch_tensor1.T @ torch_tensor2).sum() + torch_result_transpose_matmul.backward() + torch_tensor_grad_transpose_matmul1 = torch_tensor1.grad + torch_tensor_grad_transpose_matmul2 = torch_tensor2.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_transpose_matmul1, torch_tensor_grad_transpose_matmul1)) + self.assertTrue(utils.compare_torch(norch_tensor_grad_transpose_matmult2, torch_tensor_grad_transpose_matmul2)) + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/build/lib/tests/test_nn.py b/build/lib/tests/test_nn.py new file mode 100644 index 0000000..532f5b6 --- /dev/null +++ b/build/lib/tests/test_nn.py @@ -0,0 +1,204 @@ +import unittest +import norch +from norch.utils import utils_unittests as utils +import torch +import os + +class TestNNModuleLoss(unittest.TestCase): + + def setUp(self): + self.device = os.environ.get('device') + if self.device is None or self.device != 'cuda': + self.device = 'cpu' + + print(f"Running tests on: {self.device}") + + def test_mse_loss(self): + """ + Test the MSELoss + """ + loss_fn_norch = norch.nn.MSELoss() + loss_fn_torch = torch.nn.MSELoss() + + # Test case 1: Predictions and labels are equal + predictions_norch = norch.Tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 4]]).to(self.device) + labels_norch = norch.Tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 3]]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_torch_result = utils.to_torch(loss_norch).to(self.device) + + predictions_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 4]], device=self.device) + labels_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 3]], device=self.device) + + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + + self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected)) + + # Test case 2: Predictions and labels are different + predictions_norch = norch.Tensor([1.1, 2, 3, 4]).to(self.device) + labels_norch = norch.Tensor([4, 3, 2.1, 1]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_torch_result = utils.to_torch(loss_norch).to(self.device) + + predictions_torch = torch.tensor([1.1, 2, 3, 4], device=self.device) + labels_torch = torch.tensor([4, 3, 2.1, 1], device=self.device) + + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + + self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected)) + + def test_cross_entropy_loss(self): + """ + Test the CrossEntropyLoss + """ + loss_fn_norch = norch.nn.CrossEntropyLoss() + loss_fn_torch = torch.nn.CrossEntropyLoss() + + # Test case 1: Single class, single sample + predictions_norch = norch.Tensor([2.0, 1.0, 0.1]).to(self.device) + labels_norch = norch.Tensor([0]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + + loss_torch_result = utils.to_torch(loss_norch).to(self.device) + + predictions_torch = torch.tensor([2.0, 1.0, 0.1], device=self.device) + labels_torch = torch.tensor(0, device=self.device) + + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + + self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected)) + + # Test case 2: Multiple classes, multiple samples + predictions_norch = norch.Tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]]).to(self.device) + labels_norch = norch.Tensor([2, 1]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_torch_result = utils.to_torch(loss_norch).to(self.device) + + + predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]], device=self.device) + labels_torch = torch.tensor([2, 1], device=self.device) + + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + + self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected)) + + # Test case 3: Edge case - all predictions are zero + predictions_norch = norch.Tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]).to(self.device) + labels_norch = norch.Tensor([1, 2]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_torch_result = utils.to_torch(loss_norch).to(self.device) + + predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], device=self.device) + labels_torch = torch.tensor([1, 2], device=self.device) + + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + + self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected)) + + # Test case 4: Class probabilities instead of class index + predictions_norch = norch.Tensor([0.5, 0.2, 0.1]).to(self.device) + labels_norch = norch.Tensor([1., 0, 0]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_torch_result = utils.to_torch(loss_norch).to(self.device) + + predictions_torch = torch.tensor([0.5, 0.2, 0.1]) + labels_torch = torch.tensor([1., 0, 0]) + + predictions_torch.to(self.device) + labels_torch.to(self.device) + + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + + self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected)) + + # Test case 4: Batched class probabilities instead of class index + predictions_norch = norch.Tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]]).to(self.device) + labels_norch = norch.Tensor([[1., 0, 0], [0, 1, 0]]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_torch_result = utils.to_torch(loss_norch).to(self.device) + + predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]], device=self.device) + labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]], device=self.device) + + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + + self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected)) + + +class TestNNModuleActivationFn(unittest.TestCase): + + def setUp(self): + self.device = os.environ.get('device') + if self.device is None or self.device != 'cuda': + self.device = 'cpu' + + def test_sigmoid_activation(self): + """ + Test Sigmoid activation function + """ + sigmoid_fn_norch = norch.nn.Sigmoid() + sigmoid_fn_torch = torch.nn.Sigmoid() + + # Test case 1: Positive input + x = norch.Tensor([[1, 2, 3]]).to(self.device) + sigmoid_norch = sigmoid_fn_norch.forward(x) + sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device) + + x = torch.tensor([[1, 2, 3]], device=self.device) + + sigmoid_torch_expected = sigmoid_fn_torch.forward(x) + + self.assertTrue(utils.compare_torch(sigmoid_torch_result, sigmoid_torch_expected)) + + # Test case 1: Negative input + x = norch.Tensor([-1, 2, -3]).to(self.device) + sigmoid_norch = sigmoid_fn_norch.forward(x) + sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device) + + x = torch.tensor([-1, 2, -3], device=self.device) + + sigmoid_torch_expected = sigmoid_fn_torch.forward(x) + + self.assertTrue(utils.compare_torch(sigmoid_torch_result, sigmoid_torch_expected)) + + # Test case 1: Zero input + x = norch.Tensor([0, 0, 0]).to(self.device) + sigmoid_norch = sigmoid_fn_norch.forward(x) + sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device) + + x = torch.tensor([0, 0, 0], device=self.device) + + sigmoid_torch_expected = sigmoid_fn_torch.forward(x) + + self.assertTrue(utils.compare_torch(sigmoid_torch_result, sigmoid_torch_expected)) + + def test_softmax_activation(self): + """ + Test Softmax activation function + """ + + # Test different axes + axes = [0, 1, 2, -1] + + # Define the input tensors for different test cases + test_cases = [ + (norch.Tensor([[[1., 2, 3], [4, 5, 6]]]).to(self.device), torch.tensor([[[1., 2, 3], [4, 5, 6]]], device=self.device)), + (norch.Tensor([[[1., -1, 0], [2, -2, 0]]]).to(self.device), torch.tensor([[[1., -1, 0], [2, -2, 0]]], device=self.device)), + (norch.Tensor([[[0., 0, 0], [0, 0, 0]]]).to(self.device), torch.tensor([[[0., 0, 0], [0, 0, 0]]], device=self.device)) + ] + + for dim in axes: + softmax_fn_norch = norch.nn.Softmax(dim=dim) + softmax_fn_torch = torch.nn.Softmax(dim=dim) + + for norch_input, torch_input in test_cases: + + # Forward pass using norch + softmax_norch = softmax_fn_norch.forward(norch_input) + softmax_torch_result = utils.to_torch(softmax_norch).to(self.device) + + # Forward pass using torch + softmax_torch_expected = softmax_fn_torch.forward(torch_input) + + # Compare the results + self.assertTrue(utils.compare_torch(softmax_torch_result, softmax_torch_expected)) + diff --git a/build/lib/tests/test_operations.py b/build/lib/tests/test_operations.py new file mode 100644 index 0000000..262cee4 --- /dev/null +++ b/build/lib/tests/test_operations.py @@ -0,0 +1,753 @@ +import unittest +import norch +from norch.utils import utils_unittests as utils +import torch +import sys +import os + +class TestTensorOperations(unittest.TestCase): + + def setUp(self): + self.device = os.environ.get('device') + if self.device is None or self.device != 'cuda': + self.device = 'cpu' + + print(f"Running tests on: {self.device}") + + def test_creation_and_conversion(self): + """ + Test creation and convertion of norch tensor to pytorch + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_tensor = utils.to_torch(norch_tensor) + self.assertTrue(torch.is_tensor(torch_tensor)) + + def test_addition(self): + """ + Test addition two tensors: tensor1 + tensor2 + """ + norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + norch_result = norch_tensor1 + norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + torch_expected = torch_tensor1 + torch_tensor2 + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_addition_broadcasted(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).to(self.device) + + 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)) + + norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) # Shape (1, 2, 3) + norch_tensor2 = norch.Tensor([[10, 10], [5, 6]]).to(self.device) # Shape (3) + norch_result = norch_tensor1 + norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) # Shape (1, 2, 3) + torch_tensor2 = torch.tensor([[[10, 10], [5, 6]]]).to(self.device) # Shape (3) + torch_expected = torch_tensor1 + torch_tensor2 + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + # reversed order broadcasting + 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 + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + norch_result = norch_tensor2 + norch_tensor1 + torch_expected = torch_tensor2 + torch_tensor1 + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + + + def test_subtraction(self): + """ + Test subtraction of two tensors: tensor1 - tensor2 + """ + norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + norch_result = norch_tensor1 - norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + torch_expected = torch_tensor1 - torch_tensor2 + + 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).to(self.device) + + 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)) + + # reversed order broadcasting + norch_result = norch_tensor2 - norch_tensor1 + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_expected = torch_tensor2 - torch_tensor1 + + 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 + """ + norch_tensor = norch.Tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device) + scalar = 2 + norch_result = norch_tensor / scalar + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device) + torch_expected = torch_tensor / scalar + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_scalar_division_by_tensor(self): + """ + Test scalar division by a tensor: scalar / tensor + """ + scalar = 10 + norch_tensor = norch.Tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device) + norch_result = scalar / norch_tensor + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[2, 4], [6, -8]], [[10, 12], [14, 16]]]).to(self.device) + torch_expected = scalar / torch_tensor + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_matrix_multiplication(self): + """ + Test matrix multiplication: tensor1 @ tensor2 + """ + norch_tensor1 = norch.Tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[1., 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device) + norch_result = norch_tensor1 @ norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[1., 0], [0, 1]], [[-1, 0], [0, -1]]]).to(self.device) + torch_expected = torch_tensor1 @ torch_tensor2 + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_elementwise_multiplication_by_scalar(self): + """ + Test elementwise multiplication of a tensor by a scalar: tensor * scalar + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + scalar = 2 + norch_result = norch_tensor * scalar + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch_tensor * scalar + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_elementwise_multiplication_by_tensor(self): + """ + Test elementwise multiplication of two tensors: tensor1 * tensor2 + """ + norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[2, 2], [2, 2]], [[2, 2], [2, 2]]]).to(self.device) + norch_result = norch_tensor1 * norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[2, 2], [2, 2]], [[2, 2], [2, 2]]]).to(self.device) + torch_expected = torch_tensor1 * torch_tensor2 + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_reshape(self): + """ + Test reshaping of a tensor: tensor.reshape(shape) + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + new_shape = [2, 4] + norch_result = norch_tensor.reshape(new_shape) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch_tensor.reshape(new_shape) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_unsqueeze(self): + """ + Test unsqueeze operation on a tensor + """ + norch_tensor = norch.Tensor([[1, 2], [3, 4]]).to(self.device) + + # Unsqueeze at dim=0 + norch_unsqueeze_0 = norch_tensor.unsqueeze(0) + torch_unsqueeze_0 = utils.to_torch(norch_unsqueeze_0).to(self.device) + torch_tensor = torch.tensor([[1, 2], [3, 4]]).to(self.device) + torch_expected_0 = torch_tensor.unsqueeze(0) + self.assertTrue(utils.compare_torch(torch_unsqueeze_0, torch_expected_0)) + + # Unsqueeze at dim=1 + norch_unsqueeze_1 = norch_tensor.unsqueeze(1) + torch_unsqueeze_1 = utils.to_torch(norch_unsqueeze_1).to(self.device) + torch_expected_1 = torch_tensor.unsqueeze(1) + self.assertTrue(utils.compare_torch(torch_unsqueeze_1, torch_expected_1)) + + # Unsqueeze at dim=2 + norch_unsqueeze_2 = norch_tensor.unsqueeze(2) + torch_unsqueeze_2 = utils.to_torch(norch_unsqueeze_2).to(self.device) + torch_expected_2 = torch_tensor.unsqueeze(2) + self.assertTrue(utils.compare_torch(torch_unsqueeze_2, torch_expected_2)) + + # Unsqueeze at dim=-1 + norch_unsqueeze_neg_1 = norch_tensor.unsqueeze(-1) + torch_unsqueeze_neg_1 = utils.to_torch(norch_unsqueeze_neg_1).to(self.device) + torch_expected_neg_1 = torch_tensor.unsqueeze(-1) + self.assertTrue(utils.compare_torch(torch_unsqueeze_neg_1, torch_expected_neg_1)) + + # Unsqueeze at dim=-2 + norch_unsqueeze_neg_2 = norch_tensor.unsqueeze(-2) + torch_unsqueeze_neg_2 = utils.to_torch(norch_unsqueeze_neg_2).to(self.device) + torch_expected_neg_2 = torch_tensor.unsqueeze(-2) + self.assertTrue(utils.compare_torch(torch_unsqueeze_neg_2, torch_expected_neg_2)) + + def test_squeeze(self): + """ + Test squeeze operation on a tensor + """ + # Create a tensor with some dimensions of size 1 + norch_tensor = norch.Tensor([[[1, 2], [3, 4]]]).to(self.device) # shape [1, 2, 2] + + # Squeeze at dim=0 + norch_squeeze_0 = norch_tensor.squeeze(0) + torch_squeeze_0 = utils.to_torch(norch_squeeze_0).to(self.device) + torch_tensor = torch.tensor([[[1, 2], [3, 4]]]).to(self.device) + torch_expected_0 = torch_tensor.squeeze(0) + self.assertTrue(utils.compare_torch(torch_squeeze_0, torch_expected_0)) + + # Create a tensor with a dimension of size 1 in the middle + norch_tensor_middle_1 = norch.Tensor([[[1, 2]], [[3, 4]]]).to(self.device) # shape [2, 1, 2] + + # Squeeze at dim=1 + norch_squeeze_1 = norch_tensor_middle_1.squeeze(1) + torch_squeeze_1 = utils.to_torch(norch_squeeze_1).to(self.device) + torch_tensor_middle_1 = torch.tensor([[[1, 2]], [[3, 4]]]).to(self.device) + torch_expected_1 = torch_tensor_middle_1.squeeze(1) + self.assertTrue(utils.compare_torch(torch_squeeze_1, torch_expected_1)) + + # Squeeze at dim=-2 (same as dim=1 in this case) + norch_squeeze_neg_2 = norch_tensor_middle_1.squeeze(-2) + torch_squeeze_neg_2 = utils.to_torch(norch_squeeze_neg_2).to(self.device) + torch_expected_neg_2 = torch_tensor_middle_1.squeeze(-2) + self.assertTrue(utils.compare_torch(torch_squeeze_neg_2, torch_expected_neg_2)) + + # Squeeze all dimensions of size 1 (None) + norch_tensor_all_1 = norch.Tensor([[[[1, 2], [3, 4]]]]).to(self.device) # shape [1, 1, 2, 2] + norch_squeeze_all = norch_tensor_all_1.squeeze() + torch_squeeze_all = utils.to_torch(norch_squeeze_all).to(self.device) + torch_tensor_all_1 = torch.tensor([[[[1, 2], [3, 4]]]]).to(self.device) + torch_expected_all = torch_tensor_all_1.squeeze() + self.assertTrue(utils.compare_torch(torch_squeeze_all, torch_expected_all)) + + # Squeeze no dimensions (no dimensions of size 1) + norch_tensor_no_1 = norch.Tensor([[1, 2], [3, 4]]).to(self.device) # shape [2, 2] + norch_squeeze_none = norch_tensor_no_1.squeeze() + torch_squeeze_none = utils.to_torch(norch_squeeze_none).to(self.device) + torch_tensor_no_1 = torch.tensor([[1, 2], [3, 4]]).to(self.device) + torch_expected_none = torch_tensor_no_1.squeeze() + self.assertTrue(utils.compare_torch(torch_squeeze_none, torch_expected_none)) + + + def test_transpose(self): + """ + Test transposition of a tensor: tensor.transpose(dim1, dim2) + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + dim1, dim2 = 0, 2 + norch_result = norch_tensor.transpose(dim1, dim2) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch_tensor.transpose(dim1, dim2) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_logarithm(self): + """ + Test elementwise logarithm of a tensor: tensor.log() + """ + norch_tensor = norch.Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.log() + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch.log(torch_tensor) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_sum(self): + """ + Test summation of a tensor: tensor.sum() + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.sum() + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch.sum(torch_tensor) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_sum_axis(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).to(self.device) + + 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)) + + # negative axis + + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.sum(axis=-2) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch.sum(torch_tensor, dim=-2) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + + def test_sum_axis_keepdim(self): + """ + Test summation of a tensor along a specific axis with keepdim=True + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.sum(axis=1, keepdim=True) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch.sum(torch_tensor, dim=1, keepdim=True) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_max(self): + """ + Test max of a tensor: tensor.max() + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.max() + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch.max(torch_tensor) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_max_axis(self): + """ + Test max 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.max(axis=1) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected, _ = torch.max(torch_tensor, dim=1) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + # negative axis + + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.max(axis=-1) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected, _ = torch.max(torch_tensor, dim=-1) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + + def test_max_axis_keepdim(self): + """ + Test max of a tensor along a specific axis with keepdim=True + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.max(axis=1, keepdim=True) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected, _ = torch.max(torch_tensor, dim=1, keepdim=True) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_min(self): + """ + Test min of a tensor: tensor.min() + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.min() + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch.min(torch_tensor) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_min_axis(self): + """ + Test min 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.min(axis=1) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected, _ = torch.min(torch_tensor, dim=1) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + # negative axis + + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.min(axis=-1) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected, _ = torch.min(torch_tensor, dim=-1) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + + def test_min_axis_keepdim(self): + """ + Test min of a tensor along a specific axis with keepdim=True + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.min(axis=1, keepdim=True) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected, _ = torch.min(torch_tensor, dim=1, keepdim=True) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_1D_T(self): + """ + Test transposition of a 1D tensor. + """ + norch_tensor = norch.Tensor([1, 2, 3, 4]).to(self.device) + norch_result = norch_tensor.T + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([1, 2, 3, 4]).to(self.device) + torch_expected = torch_tensor.T + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_2D_T(self): + """ + Test transposition of a 2D tensor. + """ + norch_tensor = norch.Tensor([[1, 2, 3], [4, 5, 6]]).to(self.device) + norch_result = norch_tensor.T + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]).to(self.device) + torch_expected = torch_tensor.T + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_3D_T(self): + """ + Test transposition of a tensor: tensor.T + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor.T + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch_tensor.T + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_matmul(self): + """ + 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) + norch_tensor2 = norch.Tensor([[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]]).to(self.device) + + norch_result = norch_tensor1 @ norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + # Converting to PyTorch tensors for comparison + torch_tensor1 = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + torch_tensor2 = torch.tensor([[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]]).to(self.device) + + torch_expected = torch.matmul(torch_tensor1, torch_tensor2) + + # Comparing results + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_reshape_then_matmul(self): + """ + Test reshaping a tensor followed by matrix multiplication: (tensor.reshape(shape) @ other_tensor) + """ + norch_tensor = norch.Tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + new_shape = [2, 4] + norch_reshaped = norch_tensor.reshape(new_shape) + + norch_result = norch_reshaped @ norch_tensor + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + torch_expected = torch_tensor.reshape(new_shape) @ torch_tensor + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_batched_matmul(self): + """ + Test batched matrix multiplication: BxMxP = BxNxM @ BxMxP + """ + B = 3 # Batch size + + # Creating batched tensors for Norch + norch_tensor1 = norch.Tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)]).to(self.device) + norch_tensor2 = norch.Tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + + norch_result = norch_tensor1 @ norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + # Converting to PyTorch tensors for comparison + torch_tensor1 = torch.tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)]).to(self.device) + torch_tensor2 = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + + torch_expected = torch.matmul(torch_tensor1, torch_tensor2) + + # Comparing results + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + + def test_broadcasted_batched_matmul(self): + """ + Test broadcasted batched matrix multiplication: BxMxP = NxM @ BxMxP + """ + B = 3 # Batch size + + # Creating batched tensors for Norch + norch_tensor1 = norch.Tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + norch_tensor2 = norch.Tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + + norch_result = norch_tensor1 @ norch_tensor2 + torch_result = utils.to_torch(norch_result).to(self.device) + + # Converting to PyTorch tensors for comparison + torch_tensor1 = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]]).to(self.device) + torch_tensor2 = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)]).to(self.device) + + torch_expected = torch.matmul(torch_tensor1, torch_tensor2) + + # Comparing results + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + + + def test_transpose_then_matmul(self): + """ + Test transposing a tensor followed by matrix multiplication: (tensor.transpose(dim1, dim2) @ other_tensor) + """ + norch_tensor = norch.Tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + dim1, dim2 = 0, 2 + norch_result = norch_tensor.transpose(dim1, dim2) @ norch_tensor + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1., 2], [3, 4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch_tensor.transpose(dim1, dim2) @ torch_tensor + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_add_div_matmul_then_reshape(self): + """ + Test a combination of operations: (tensor.sum() + other_tensor) / scalar @ another_tensor followed by reshape + """ + norch_tensor1 = norch.Tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[1., 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + scalar = 2 + new_shape = [2, 4] + norch_result = ((norch_tensor1 + norch_tensor2) / scalar) @ norch_tensor1 + norch_result = norch_result.reshape(new_shape) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[1., 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[1., 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + torch_expected = ((torch_tensor1 + torch_tensor2) / scalar) @ torch_tensor1 + torch_expected = torch_expected.reshape(new_shape) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_scalar_power_tensor(self): + """ + Test scalar power of a tensor: scalar ** tensor + """ + scalar = 3 + norch_tensor = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = scalar ** norch_tensor + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = scalar ** torch_tensor + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_tensor_power_scalar(self): + """ + Test tensor power of a scalar: tensor ** scalar + """ + scalar = 3 + norch_tensor = norch.Tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_result = norch_tensor ** scalar + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_expected = torch_tensor ** scalar + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_tensor_sin(self): + """ + Test sine function on tensor + """ + norch_tensor = norch.Tensor([[[0, 30], [45, 60]], [[90, 120], [135, 180]]]).to(self.device) + norch_result = norch_tensor.sin() + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[0, 30], [45, 60]], [[90, 120], [135, 180]]]).to(self.device) + torch_expected = torch.sin(torch_tensor) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_tensor_cos(self): + """ + Test cosine function on tensor + """ + norch_tensor = norch.Tensor([[[0, 30], [45, 60]], [[90, 120], [135, 180]]]).to(self.device) + norch_result = norch_tensor.cos() + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor = torch.tensor([[[0, 30], [45, 60]], [[90, 120], [135, 180]]]).to(self.device) + torch_expected = torch.cos(torch_tensor) + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_equal(self): + """ + Test equal two tensors: tensor1.equal(tensor2) + """ + norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 1], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + norch_result = norch_tensor1.equal(norch_tensor2) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[1, 2], [3, -4]], [[5, 1], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device) + torch_expected = (torch_tensor1 == torch_tensor2).float() + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + def test_broadcasted_equal(self): + """ + Test broadcasted equal two tensors: tensor1.equal(tensor2) + """ + norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[10, 10]], [[5, 6]]]).to(self.device) + norch_result = norch_tensor1.equal(norch_tensor2) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[10, 10]], [[5, 6]]]).to(self.device) + torch_expected = (torch_tensor1 == torch_tensor2).float() + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) + norch_tensor2 = norch.Tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]]).to(self.device) + norch_result = norch_tensor1.equal(norch_tensor2) + torch_result = utils.to_torch(norch_result).to(self.device) + + torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) + torch_tensor2 = torch.tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]]).to(self.device) + torch_expected = (torch_tensor1 == torch_tensor2).float() + + self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + + + def test_zeros_like(self): + """ + Test creating a tensor of zeros with the same shape as another tensor. + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_zeros = norch_tensor.zeros_like() + torch_zeros_result = utils.to_torch(norch_zeros).to(self.device) + + torch_tensor_expected = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_zeros_expected = torch.zeros_like(torch_tensor_expected) + + self.assertTrue(utils.compare_torch(torch_zeros_result, torch_zeros_expected)) + + def test_ones_like(self): + """ + Test creating a tensor of ones with the same shape as another tensor. + """ + norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + norch_ones = norch_tensor.ones_like() + torch_ones_result = utils.to_torch(norch_ones).to(self.device) + + torch_tensor_expected = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device) + torch_ones_expected = torch.ones_like(torch_tensor_expected) + + self.assertTrue(utils.compare_torch(torch_ones_result, torch_ones_expected)) + + +if __name__ == '__main__': + unittest.main() diff --git a/setup.py b/setup.py index 8774eb8..d94cfc4 100644 --- a/setup.py +++ b/setup.py @@ -18,7 +18,7 @@ apt_get_dependencies = [ apt_nccl = [ 'libnccl2=2.21.5-1+cuda12.2', - 'libnccl-dev=2.21.5-1+cuda12.2 ' + 'libnccl-dev=2.21.5-1+cuda12.2' ] subprocess.check_call(['sudo', 'apt-get', 'update']) @@ -27,14 +27,16 @@ for package in apt_dependencies: subprocess.check_call(['sudo', 'apt', 'install', package]) for package in apt_get_dependencies: - subprocess.check_call(['sudo', 'apt-get', 'install', 'y', package]) + subprocess.check_call(['sudo', 'apt-get', 'install', '-y', package]) -subprocess.check_call(['wget', 'https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb ']) +subprocess.check_call(['wget', 'https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb']) subprocess.check_call(['sudo', 'dpkg', '-i', 'cuda-keyring_1.0-1_all.deb']) for package in apt_nccl: subprocess.check_call(['sudo', 'apt', 'install', package]) +subprocess.check_call(['rm', 'cuda-keyring_1.0-1_all.deb']) + class CustomInstall(install): def run(self): subprocess.call(['make', '-C', 'build'])