Merge pull request #59 from lucasdelimanogueira/tmp

Tmp

norch/__init__.py

@@ -2,6 +2,7 @@ from norch.tensor import Tensor
 from .nn import *
 from .optim import *
 from .utils import *
+from .datasets import *
 
 __version__ = "0.0.1"
 __author__ = 'Lucas de Lima Nogueira'

norch/autograd/functions.py

@@ -1,4 +1,5 @@
 import math
+import norch
 
 class AddBackward:
     def __init__(self, x, y):
@@ -15,6 +16,7 @@ class AddBroadcastedBackward:
         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):
@@ -25,11 +27,17 @@ class AddBroadcastedBackward:
         # 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)
+                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]
@@ -49,15 +57,7 @@ class SubBroadcastedBackward:
         for i in range(len(shape)):
             if shape[i] == 1:
                 gradient = gradient.sum(axis=i)
-
         return gradient
-
-class SubBackward:
-    def __init__(self, x, y):
-        self.input = [x, y]
-
-    def backward(self, gradient):
-        return [gradient, -gradient]
     
 class ScalarMulBackward:
     def __init__(self, x, scalar):
@@ -82,7 +82,13 @@ class MatmulBackward:
 
     def backward(self, gradient):
         x, y = self.input
-        return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
+        
+        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):
@@ -121,12 +127,29 @@ class LogBackward:
         return [grad_input]
        
 class SumBackward:
-    def __init__(self, x):
+    def __init__(self, x, axis=None, keepdim=False):
         self.input = [x]
+        self.axis = axis
+        self.keepdim = keepdim
 
     def backward(self, gradient):
-        # Since sum reduces a tensor to a scalar, gradient is broadcasted to match the original shape.
-        return [float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()]
+        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.tensor.contents.data[0]) * 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):
@@ -177,5 +200,95 @@ class CosBackward:
     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.tensor.contents.data[0]) * self.input[0].ones_like()
+
+            grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[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.tensor.contents.data[0]) * self.input[0].ones_like()
+
+            grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[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
 
 

norch/csrc/cpu.cpp

@@ -11,7 +11,7 @@ void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
     }
 }
 
-void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
+void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size) {
     int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
 
     // Calculate strides for broadcasting
@@ -28,12 +28,12 @@ void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_
         int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
         strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
         strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
-        stride1 *= broadcasted_shape[i];
-        stride2 *= broadcasted_shape[i];
+        stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
+        stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
     }
 
     // Perform element-wise addition with broadcasting
-    for (int i = 0; i < tensor1->size; i++) {
+    for (int i = 0; i < broadcasted_size; i++) {
         int index1 = 0, index2 = 0;
         int linear_index = i;
         for (int j = max_ndim - 1; j >= 0; j--) {
@@ -58,7 +58,7 @@ void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
     }
 }
 
-void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
+void 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
@@ -75,12 +75,12 @@ void sub_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_
         int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
         strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
         strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
-        stride1 *= broadcasted_shape[i];
-        stride2 *= broadcasted_shape[i];
+        stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
+        stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
     }
 
     // Perform element-wise addition with broadcasting
-    for (int i = 0; i < tensor1->size; i++) {
+    for (int i = 0; i < broadcasted_size; i++) {
         int index1 = 0, index2 = 0;
         int linear_index = i;
         for (int j = max_ndim - 1; j >= 0; j--) {
@@ -205,7 +205,7 @@ void log_tensor_cpu(Tensor* tensor, float* result_data) {
     }
 }
 
-void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis) {
+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;
@@ -219,26 +219,14 @@ void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis) {
             return;
         }
         
-        int result_shape[tensor->ndim - 1];
-        int result_size = 1;
         int axis_stride = tensor->strides[axis];
 
-        int idx = 0;
-        for (int i = 0; i < tensor->ndim; i++) {
-            if (i != axis) {
-                result_shape[idx++] = tensor->shape[i];
-                result_size *= tensor->shape[i];
-            }
-        }
-
-        memset(result_data, 0, result_size * sizeof(float));
-
         for (int i = 0; i < tensor->shape[axis]; i++) {
-            for (int j = 0; j < result_size; j++) {
+            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];
+                    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];
@@ -247,8 +235,115 @@ void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis) {
     }
 }
 
+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) {

norch/csrc/cpu.h

@@ -4,10 +4,12 @@
 #include "tensor.h"
 
 void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
-void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape);
-void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis);
+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);
+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);
@@ -19,6 +21,8 @@ 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);

norch/csrc/cuda.cu

@@ -71,7 +71,7 @@ __global__ void add_broadcasted_tensor_cuda_kernel(float* data1, float* data2, f
     result_data[i] = data1[index1] + data2[index2];
 }
 
-__host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
+__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));
@@ -87,8 +87,8 @@ __host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, floa
         int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
         strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
         strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
-        stride1 *= broadcasted_shape[i];
-        stride2 *= broadcasted_shape[i];
+        stride1 *= (dim1 == broadcasted_shape[i]) ? dim1 : 1;
+        stride2 *= (dim2 == broadcasted_shape[i]) ? dim2 : 1;
     }
     
     int* d_broadcasted_shape;
@@ -104,8 +104,8 @@ __host__ void add_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, floa
     cudaMalloc((void**)&d_strides2, max_ndim * sizeof(int));
     cudaMemcpy(d_strides2, strides2, max_ndim * sizeof(int), cudaMemcpyHostToDevice);
 
-    int number_of_blocks = (tensor1->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
-    add_broadcasted_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, d_broadcasted_shape, d_strides1, d_strides2, max_ndim, tensor1->size);
+    int number_of_blocks = (broadcasted_size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    add_broadcasted_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor1->data, tensor2->data, result_data, d_broadcasted_shape, d_strides1, d_strides2, max_ndim, broadcasted_size);
 
     cudaError_t error = cudaGetLastError();
     if (error != cudaSuccess) {
@@ -187,6 +187,67 @@ __host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
     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<<<number_of_blocks, THREADS_PER_BLOCK>>>(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;
@@ -481,6 +542,89 @@ __host__ void log_tensor_cuda(Tensor* tensor, float* result_data) {
     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<<<number_of_blocks, THREADS_PER_BLOCK>>>(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<<<number_of_blocks, THREADS_PER_BLOCK>>>(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) {
     

norch/csrc/cuda.h

@@ -8,7 +8,10 @@
     __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);
+    __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);
     __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data);
@@ -46,6 +49,12 @@
     __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);
 

norch/csrc/tensor.cpp

@@ -149,26 +149,30 @@ extern "C" {
             broadcasted_shape[max_ndim - 1 - i] = dim1 > dim2 ? dim1 : dim2;
         }
 
+        int broadcasted_size = 1;
+        for (int i = 0; i < max_ndim; i++) {
+            broadcasted_size *= broadcasted_shape[i];
+        }
+
         if (strcmp(tensor1->device, "cuda") == 0) {
             float* result_data;
-            cudaMalloc((void **)&result_data, tensor1->size * sizeof(float));
-            add_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape);
+            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(tensor1->size * sizeof(float));
+            float* result_data = (float*)malloc(broadcasted_size * sizeof(float));
             if (result_data == NULL) {
                 fprintf(stderr, "Memory allocation failed\n");
                 exit(1);
             }
 
-            add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
+            add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size);
             return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
         }
     }
 
-    Tensor* sum_tensor(Tensor* tensor, int axis) {
-
+    Tensor* sum_tensor(Tensor* tensor, int axis, bool keepdim) {
         char* device = (char*)malloc(strlen(tensor->device) + 1);
         if (device != NULL) {
             strcpy(device, tensor->device);
@@ -179,11 +183,8 @@ extern "C" {
         int ndim;
         int* shape;
         if (axis == -1) {
+            
             shape = (int*) malloc(sizeof(int));
-            if (shape == NULL) {
-                fprintf(stderr, "Memory allocation failed\n");
-                exit(1);
-            }
             shape[0] = 1;
             ndim = 1;
         } else {
@@ -195,21 +196,179 @@ extern "C" {
             }
             ndim = tensor->ndim - 1;
         }
+
+        int size = 1;
+        for (int i = 0; i < ndim; i++) {
+            size *= shape[i];
+        }
   
         if (strcmp(tensor->device, "cuda") == 0) {
 
             float* result_data;
-            cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
+            cudaMalloc((void**)&result_data, size * sizeof(float));
+            cudaMemset(result_data, 0, size * sizeof(float));
             sum_tensor_cuda(tensor, result_data);
             return create_tensor(result_data, shape, ndim, device);
         } 
         else {
-            float* result_data = (float*)malloc(1 * sizeof(float));
+            float* result_data = (float*)calloc(size, sizeof(float));
             if (result_data == NULL) {
                 fprintf(stderr, "Memory allocation failed\n");
                 exit(1);
             }
-            sum_tensor_cpu(tensor, result_data, axis);
+
+            sum_tensor_cpu(tensor, result_data, 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 size = 1;
+        for (int i = 0; i < ndim; i++) {
+            size *= shape[i];
+        }
+  
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void**)&result_data, size * sizeof(float));
+            //cudaMemset(result_data, -INFINITY, size * sizeof(float));
+            //max_tensor_cuda(tensor, 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);
+            }
+            max_tensor_cpu(tensor, result_data, 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 size = 1;
+        for (int i = 0; i < ndim; i++) {
+            size *= shape[i];
+        }
+  
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void**)&result_data, size * sizeof(float));
+            //cudaMemset(result_data, INFINITY, size * sizeof(float));
+            //min_tensor_cuda(tensor, 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);
+            }
+            min_tensor_cpu(tensor, result_data, 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);
         }     
     }
@@ -290,16 +449,27 @@ extern "C" {
             broadcasted_shape[max_ndim - 1 - i] = dim1 > dim2 ? dim1 : dim2;
         }
 
-        // Allocate memory for result tensor
-        float* result_data = (float*)malloc(tensor1->size * sizeof(float));
-        if (result_data == NULL) {
-            fprintf(stderr, "Memory allocation failed\n");
-            exit(1);
+        int broadcasted_size = 1;
+        for (int i = 0; i < max_ndim; i++) {
+            broadcasted_size *= broadcasted_shape[i];
         }
 
-        sub_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
+        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);
+            }
 
-        return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
+            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) {
@@ -469,7 +639,7 @@ extern "C" {
 
     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 multiplication\n", 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);
         }
 
@@ -494,7 +664,7 @@ extern "C" {
 
         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);
+                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];
@@ -872,6 +1042,106 @@ extern "C" {
         }
     }
 
+    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) {

norch/csrc/tensor.h

@@ -16,7 +16,9 @@ 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);
+    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);
@@ -28,6 +30,8 @@ extern "C" {
     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);

norch/datasets/__init__.py

@@ -0,0 +1 @@
+from .mnist import *

norch/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)

norch/nn/__init__.py

@@ -1,3 +1,4 @@
 from .modules import *
 from .activation import *
-from .loss import *
+from .loss import *
+from .functional import *

norch/nn/activation.py

@@ -1,4 +1,5 @@
 from .module import Module
+from . import functional as F
 import math
 
 class Activation(Module):
@@ -17,4 +18,13 @@ class Sigmoid(Activation):
         super().__init__()
 
     def forward(self, x):
-        return 1.0 / (1.0 + (math.e) ** (-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)

norch/nn/functional.py

@@ -0,0 +1,33 @@
+import math
+import norch
+import numpy as np
+
+def sigmoid(x):
+    return 1.0 / (1.0 + (math.e) ** (-x)) 
+
+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.tensor.contents.data[i])
+        one_hot[i][target_idx] = 1
+
+    if x.numel < 2:
+        one_hot = one_hot[0]
+
+    return norch.Tensor(one_hot)

norch/nn/loss.py

@@ -1,4 +1,6 @@
 from .module import Module
+from norch.autograd.functions import *
+import norch
 from abc import ABC
 
 class Loss(Module, ABC):
@@ -18,8 +20,66 @@ class MSELoss(Loss):
     def __init__(self):
         super().__init__()
 
-    def forward(self, predictions, labels):
-        assert labels.shape == predictions.shape, \
-            "Labels and predictions shape does not match: {} and {}".format(labels.shape, predictions.shape)
+    def forward(self, predictions, target):
+        assert target.shape == predictions.shape, \
+            "Labels and predictions shape does not match: {} and {}".format(target.shape, predictions.shape)
         
-        return ((predictions - labels) ** 2).sum() / predictions.numel
+        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 == 1:
+            if target.numel == 1:
+                num_classes = input.shape[0]
+                target = norch.one_hot_encode(target, num_classes)
+                
+                logits = norch.softmax(input, dim=0)
+                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)
+                cost = -(logits.log() * target).sum()
+
+
+        elif input.ndim == 2:
+            # batched 
+            if target.ndim == 1:
+                # target -> Ground truth class indices:
+                num_classes = input.shape[1]
+
+                target = norch.one_hot_encode(target, num_classes)
+                
+                batch_size = input.shape[0]
+                logits = norch.softmax(input, dim=1)
+                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)
+                cost = -(logits.log() * target).sum() / batch_size
+
+        if input.requires_grad:
+            cost.grad_fn = CrossEntropyLossBackward(input, target)
+
+        return cost
+            
+
+

norch/nn/modules/linear.py

@@ -2,15 +2,23 @@ from ..module import Module
 from ..parameter import Parameter
 
 class Linear(Module):
-    def __init__(self, input_dim, output_dim):
+    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])
-        self.bias = Parameter(shape=[self.output_dim, 1])
+
+        if bias:
+            self.bias = Parameter(shape=[self.output_dim, 1])
+        else:
+            self.bias = None
 
     def forward(self, x):
-        z = self.weight @ x + self.bias
+        if self.bias:
+            z = self.weight @ x + self.bias
+        else:
+            z = self.weight @ x
+
         return z
 
     def inner_repr(self):

norch/tensor.py

@@ -19,13 +19,18 @@ class Tensor:
     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.data_ctype = (ctypes.c_float * len(data))(*data)
-            self.shape_ctype = (ctypes.c_int * len(shape))(*shape)
+            
+            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.shape = shape
             self.ndim = len(shape)
             self.device = device
 
@@ -109,6 +114,25 @@ class Tensor:
         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))
@@ -129,7 +153,41 @@ class Tensor:
             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:
+                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):
         self.device = device
         self.device_ctype = self.device.encode('utf-8')
@@ -224,6 +282,8 @@ class Tensor:
         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:
@@ -233,17 +293,25 @@ class Tensor:
             shape1 = [1] * (max_len - len(shape1)) + shape1
             shape2 = [1] * (max_len - len(shape2)) + shape2
 
-            broadcasted_shape = []
+            
             for dim1, dim2 in zip(shape1, shape2):
                 if dim1 != dim2 and dim1 != 1 and dim2 != 1:
                     raise ValueError("Shapes are not compatible for broadcasting")
-                broadcasted_shape.append(max(dim1, dim2))
-            return broadcasted_shape, True
+                broadcasted_shape_add.append(max(dim1, dim2))
+            return broadcasted_shape_add, True
 
-        broadcasted_shape, needs_broadcasting = broadcast_shape(self.shape, other.shape)
+        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)
 
@@ -251,11 +319,13 @@ class Tensor:
 
             result_data = Tensor()
             result_data.tensor = result_tensor_ptr
-            result_data.shape = broadcasted_shape
-            result_data.ndim = len(broadcasted_shape)
+            result_data.shape = broadcasted_shape_add.copy()
+            result_data.ndim = len(broadcasted_shape_add)
 
             result_data.device = self.device
-            result_data.numel = self.numel  # Update this to calculate the correct number of elements if broadcasting
+            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:
@@ -312,6 +382,8 @@ class Tensor:
         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:
@@ -321,16 +393,22 @@ class Tensor:
             shape1 = [1] * (max_len - len(shape1)) + shape1
             shape2 = [1] * (max_len - len(shape2)) + shape2
 
-            broadcasted_shape = []
+            
             for dim1, dim2 in zip(shape1, shape2):
                 if dim1 != dim2 and dim1 != 1 and dim2 != 1:
                     raise ValueError("Shapes are not compatible for broadcasting")
-                broadcasted_shape.append(max(dim1, dim2))
-            return broadcasted_shape, True
+                broadcasted_shape_sub.append(max(dim1, dim2))
+            return broadcasted_shape_sub, True
 
-        broadcasted_shape, needs_broadcasting = broadcast_shape(self.shape, other.shape)
+        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)
@@ -339,8 +417,8 @@ class Tensor:
 
             result_data = Tensor()
             result_data.tensor = result_tensor_ptr
-            result_data.shape = broadcasted_shape
-            result_data.ndim = len(broadcasted_shape)
+            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
@@ -569,6 +647,9 @@ class Tensor:
                 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)
 
@@ -611,6 +692,80 @@ class Tensor:
         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)
@@ -630,20 +785,34 @@ class Tensor:
 
         return result_data
     
-    def sum(self, axis=-1):
-        Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int]
+    def sum(self, axis=None, keepdim=False):
+        if axis is not None and axis < 0:
+            axis = self.ndim + axis
+            
+        if axis == None:
+            axis = -1
+
+        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)
+        result_tensor_ptr = Tensor._C.sum_tensor(self.tensor, axis, keepdim)
 
         result_data = Tensor()
         result_data.tensor = result_tensor_ptr
 
         if axis == -1:
-            result_data.shape = [1]
-            result_data.ndim = 1
+            if keepdim:
+                result_data.ndim = self.ndim
+                result_data.shape = [1] * self.ndim
+
+            else:
+                result_data.shape = [1]
+                result_data.ndim = 1
         else:
-            result_data.shape = self.shape[:axis] + self.shape[axis+1:]
+            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
@@ -653,7 +822,89 @@ class Tensor:
 
         result_data.requires_grad = self.requires_grad
         if result_data.requires_grad:
-            result_data.grad_fn = SumBackward(self)
+            result_data.grad_fn = SumBackward(self, axis, 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
+
+        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
+        
+        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
 

norch/utils/__init__.py

@@ -1 +1,2 @@
-from .utils import *
+from .utils import *
+from .data import *

norch/utils/data/__init__.py

@@ -0,0 +1,4 @@
+from .dataset import *
+from .example import *
+from .dataloader import *
+from .batch import *

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

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

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)

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 <https://github.com/pytorch/text/blob/master/torchtext/data/example.py>`__.
+    
+    """
+    @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

norch/utils/utils.py

@@ -1,4 +1,9 @@
 import random
+import requests
+import tarfile
+import zipfile
+import shutil
+import os
 
 def generate_random_list(shape):
     """
@@ -11,4 +16,105 @@ def generate_random_list(shape):
         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])]
+            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

test.py

@@ -0,0 +1,75 @@
+import norch
+from norch.utils.data.dataloader import Dataloader
+import norch
+import norch.nn as nn
+import norch.optim as optim
+import random
+random.seed(1)
+
+
+to_tensor = lambda x: norch.Tensor(x)
+reshape = lambda x: x.reshape([-1, 784])
+transform = lambda x: reshape(to_tensor(x))
+target_transform = lambda x: to_tensor(x)
+
+train_data, test_data = norch.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
+sample, _ = train_data[0]
+
+BATCH_SIZE = 100
+
+train_loader = Dataloader(train_data, batch_size = BATCH_SIZE)
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+        self.fc1 = nn.Linear(784, 5)
+        self.sigmoid = nn.Sigmoid()
+        self.fc2 = nn.Linear(5, 10)
+
+    def forward(self, x):
+        out = self.fc1(x)
+        out = self.sigmoid(out)
+        out = self.fc2(out)
+        
+        return out
+
+device = "cpu"
+epochs = 10
+
+model = MyModel().to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.SGD(model.parameters(), lr=0.001)
+loss_list = []
+
+for epoch in range(epochs):    
+    for idx, batch in enumerate(train_loader):
+        x, target = batch
+
+        x = x.unsqueeze(-1)
+        target = target
+
+        x = x.to(device)
+        target = target.to(device).unsqueeze(-1)
+
+        outputs = model(x)
+        print(outputs.shape, target.shape, x.shape)
+        loss = criterion(outputs, target)
+        
+        optimizer.zero_grad()
+        loss.backward()
+        print('loss_antes', loss)
+
+        print('f1 antes', model.fc1.bias)
+        print('f1 grad_antes', model.fc1.bias.grad)
+        print('f2 antes', model.fc2.bias)
+        print('f2 grad_antes', model.fc2.bias.grad)
+
+        optimizer.step()
+        print('\n')
+
+        print('f1 depois', model.fc1.bias)
+        print('f2 depois', model.fc2.bias)
+        print('\n\n')
+
+    print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')
+    loss_list.append(loss[0])

tests/test_autograd.py

@@ -34,7 +34,204 @@ class TestTensorAutograd(unittest.TestCase):
         self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
         self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
 
-    def test_broadcasting_addition_autograd(self):
+    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).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        
+        torch_result = (torch_tensor1 + torch_tensor2).sum(axis=0).sum(axis=0).sum()
+        torch_result.backward()
+        torch_tensor1_grad = torch_tensor1.grad
+        torch_tensor2_grad = torch_tensor2.grad
+        
+        self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
+        self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
+
+        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).to(self.device)
+        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        
+        torch_result = (torch_tensor1 + torch_tensor2).sum(axis=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).to(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).to(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).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))
+
+
+    ## 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).to(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).to(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).to(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
         """
@@ -54,6 +251,26 @@ class TestTensorAutograd(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
         self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
+
+        ## 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).to(self.device)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+
+        torch_result = (torch_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):
         """
@@ -76,7 +293,7 @@ class TestTensorAutograd(unittest.TestCase):
         self.assertTrue(utils.compare_torch(norch_tensor1_grad_sub, torch_tensor1_grad_sub))
         self.assertTrue(utils.compare_torch(norch_tensor2_grad_sub, torch_tensor2_grad_sub))
         
-    def test_broadcasting_subtraction_autograd(self):
+    def test_broadcasted_subtraction_autograd(self):
         """
         Test autograd for broadcasting subtraction: tensor1 - tensor2
         """
@@ -96,6 +313,26 @@ class TestTensorAutograd(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
         self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
+
+        # 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).to(self.device)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+        
+        torch_result = (torch_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):
@@ -211,7 +448,71 @@ class TestTensorAutograd(unittest.TestCase):
         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).to(self.device)
+        torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(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).to(self.device)
+        torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(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
@@ -282,7 +583,169 @@ class TestTensorAutograd(unittest.TestCase):
         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).to(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).to(self.device)
+        labels_torch = torch.tensor([4, 3, 2.1, 1]).to(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).to(self.device)
+        labels_torch = torch.tensor(0).to(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).to(self.device)
+        labels_torch = torch.tensor([2, 1]).to(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).to(self.device)
+        labels_torch = torch.tensor([1, 2]).to(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).to(self.device)
+        labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]]).to(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):
         """
@@ -361,6 +824,115 @@ class TestTensorAutograd(unittest.TestCase):
         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).to(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).to(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).to(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).to(self.device)
+        torch_tensor2 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True).to(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).to(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).to(self.device)
+        torch_tensor2 = torch.tensor([[[1., 2], [3, 4]]], dtype=torch.float32, requires_grad=True).to(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):
         """

tests/test_nn.py

@@ -21,13 +21,13 @@ class TestNNModuleLoss(unittest.TestCase):
         loss_fn_torch = torch.nn.MSELoss()
 
         # Test case 1: Predictions and labels are equal
-        predictions_norch = norch.Tensor([1.1, 2, 3, 4]).to(self.device)
-        labels_norch = norch.Tensor([1.1, 2, 3, 4]).to(self.device)
+        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]).to(self.device)
-        labels_torch = torch.tensor([1.1, 2, 3, 4]).to(self.device)
+        predictions_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 4]]).to(self.device)
+        labels_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 3]]).to(self.device)
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)        
         
         self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected))
@@ -44,6 +44,75 @@ class TestNNModuleLoss(unittest.TestCase):
         
         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]).to(self.device)
+        labels_torch = torch.tensor(0).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 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]]).to(self.device)
+        labels_torch = torch.tensor([2, 1]).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 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]]).to(self.device)
+        labels_torch = torch.tensor([1, 2]).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: 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]).to(self.device)
+        labels_torch = torch.tensor([1., 0, 0]).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]]).to(self.device)
+        labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]]).to(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):
 
@@ -60,11 +129,11 @@ class TestNNModuleActivationFn(unittest.TestCase):
         sigmoid_fn_torch = torch.nn.Sigmoid()
 
         # Test case 1: Positive input
-        x = norch.Tensor([1, 2, 3]).to(self.device)
+        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]).to(self.device)
+        x = torch.tensor([[1, 2, 3]]).to(self.device)
         sigmoid_torch_expected = sigmoid_fn_torch.forward(x)
 
         self.assertTrue(utils.compare_torch(sigmoid_torch_result, sigmoid_torch_expected))
@@ -87,4 +156,39 @@ class TestNNModuleActivationFn(unittest.TestCase):
         x = torch.tensor([0, 0, 0]).to(self.device)
         sigmoid_torch_expected = sigmoid_fn_torch.forward(x)
 
-        self.assertTrue(utils.compare_torch(sigmoid_torch_result, sigmoid_torch_expected))
+        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]]]), torch.tensor([[[1., 2, 3], [4, 5, 6]]])),
+            (norch.Tensor([[[1., -1, 0], [2, -2, 0]]]), torch.tensor([[[1., -1, 0], [2, -2, 0]]])),
+            (norch.Tensor([[[0., 0, 0], [0, 0, 0]]]), torch.tensor([[[0., 0, 0], [0, 0, 0]]]))
+        ]
+
+        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:
+                # Move tensors to the correct device
+                norch_input = norch_input.to(self.device)
+                torch_input = torch_input.to(self.device)
+
+                # 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))
+

tests/test_operations.py

@@ -37,7 +37,7 @@ class TestTensorOperations(unittest.TestCase):
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
-    def test_broadcasting_addition(self):
+    def test_addition_broadcasted(self):
         """
         Test addition of two tensors with broadcasting: tensor1 + tensor2
         """
@@ -52,6 +52,29 @@ class TestTensorOperations(unittest.TestCase):
 
         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
 
@@ -89,6 +112,14 @@ class TestTensorOperations(unittest.TestCase):
 
         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):
         """
@@ -176,6 +207,94 @@ class TestTensorOperations(unittest.TestCase):
 
         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))
+
+        # Squeeze at dim=2 (should raise an error because size is not 1)
+        with self.assertRaises(ValueError):
+            norch_tensor.squeeze(2)
+
+        # 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)
@@ -229,6 +348,134 @@ class TestTensorOperations(unittest.TestCase):
 
         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_transpose_T(self):
         """
         Test transposition of a tensor: tensor.T
@@ -242,6 +489,26 @@ class TestTensorOperations(unittest.TestCase):
 
         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)
@@ -258,6 +525,52 @@ class TestTensorOperations(unittest.TestCase):
 
         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):
         """
@@ -346,6 +659,47 @@ class TestTensorOperations(unittest.TestCase):
 
         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):
         """