Merge pull request #32 from lucasdelimanogueira/tmp

Tmp

norch/autograd/functions.py

@@ -33,10 +33,8 @@ class MatmulBackward:
 
     def backward(self, gradient):
         x, y = self.input
-        return [gradient @ y.T, x.T @ gradient]
+        return [gradient @ y.transpose(-1,-2), x.transpose(-1,-2) @ gradient]
         
-
-    
 class PowBackward:
     def __init__(self, x, power):
         self.input = [x]
@@ -58,4 +56,13 @@ class ReshapeBackward:
         self.input = [x]
 
     def backward(self, gradient):
-        return [gradient.reshape(self.input[0].shape)]
+        return [gradient.reshape(self.input[0].shape)]
+    
+class TransposeBackward:
+    def __init__(self, x, axis1, axis2):
+        self.input = [x]
+        self.axis1 = axis1
+        self.axis2 = axis2
+
+    def backward(self, gradient):
+        return [gradient.transpose(self.axis2, self.axis1)]

norch/csrc/cpu.cpp

@@ -45,6 +45,46 @@ void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
     }
 }
 
+
+void broadcasted_batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+
+    int tensor2_offset = tensor2->shape[1] * tensor2->shape[2];
+    int result_data_offset = tensor1->shape[0] * tensor2->shape[2];
+
+    for (int batch = 0; batch < tensor2->shape[0]; batch++) {
+    
+        for (int i = 0; i < tensor1->shape[0]; i++) {
+            for (int j = 0; j < tensor2->shape[2]; j++) {
+                float sum = 0.0;
+                for (int k = 0; k < tensor1->shape[1]; k++) {
+                    sum += tensor1->data[i * tensor1->shape[1] + k] * tensor2->data[batch*tensor2_offset + (k * tensor2->shape[2] + j)];
+                }
+                result_data[(batch * result_data_offset) + (i * tensor2->shape[2] + j)] = sum;
+            }
+        }
+    }
+} 
+
+void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
+
+    int tensor1_offset = tensor1->shape[1] * tensor1->shape[2];
+    int tensor2_offset = tensor2->shape[1] * tensor2->shape[2];
+    int result_data_offset = tensor1->shape[1] * tensor2->shape[2];
+
+    for (int batch = 0; batch < tensor2->shape[0]; batch++) {
+    
+        for (int i = 0; i < tensor1->shape[1]; i++) {
+            for (int j = 0; j < tensor2->shape[2]; j++) {
+                float sum = 0.0;
+                for (int k = 0; k < tensor1->shape[2]; k++) {
+                    sum += tensor1->data[(batch * tensor1_offset) + i * tensor1->shape[2] + k] * tensor2->data[batch*tensor2_offset + (k * tensor2->shape[2] + j)];
+                }
+                result_data[(batch * result_data_offset) + (i * tensor2->shape[2] + j)] = sum;
+            }
+        }
+    }
+} 
+
 void pow_tensor_cpu(Tensor* tensor, float power, float* result_data) {
     
     for (int i = 0; i < tensor->size; i++) {
@@ -76,7 +116,14 @@ void zeros_like_tensor_cpu(Tensor* tensor, float* result_data) {
     }
 }
 
-void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
+void transpose_1D_tensor_cpu(Tensor* tensor, float* result_data) {
+
+    for (int i = 0; i < tensor->shape[0]; i++) {
+        result_data[i] = tensor->data[i];
+    }
+}
+
+void transpose_2D_tensor_cpu(Tensor* tensor, float* result_data) {
     int rows = tensor->shape[0];
     int cols = tensor->shape[1];
 
@@ -87,6 +134,20 @@ void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
     }
 }
 
+void transpose_3D_tensor_cpu(Tensor* tensor, float* result_data) {
+    int depth = tensor->shape[0];
+    int rows = tensor->shape[1];
+    int cols = tensor->shape[2];
+
+    for (int i = 0; i < depth; i++) {
+        for (int j = 0; j < rows; j++) {
+            for (int k = 0; k < cols; k++) {
+                result_data[k * rows * depth + j * depth + i] = tensor->data[i * rows * cols + j * cols + k];
+            }
+        }
+    }
+}
+
 void assign_tensor_cpu(Tensor* tensor, float* result_data) {
 
     for (int i = 0; i < tensor->size; i++) {

norch/csrc/cpu.h

@@ -8,11 +8,16 @@ void sum_tensor_cpu(Tensor* tensor1, float* result_data);
 void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
+void broadcasted_batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
+void batched_matmul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
 void pow_tensor_cpu(Tensor* tensor, float power, float* result_data);
 void scalar_mul_tensor_cpu(Tensor* tensor, float scalar, float* result_data);
 void ones_like_tensor_cpu(Tensor* tensor, float* result_data);
 void zeros_like_tensor_cpu(Tensor* tensor, float* result_data);
-void transpose_tensor_cpu(Tensor* tensor, float* result_data);
+void transpose_1D_tensor_cpu(Tensor* tensor, float* result_data);
+void transpose_2D_tensor_cpu(Tensor* tensor, float* result_data);
+void transpose_3D_tensor_cpu(Tensor* tensor, float* result_data);
+void transpose_axes_cpu(Tensor* tensor, float* result_data, int axis1, int axis2, int* new_shape);
 void assign_tensor_cpu(Tensor* tensor, float* result_data);
 
 #endif /* CPU_H */

norch/csrc/cuda.cu

@@ -5,7 +5,6 @@
 
 #define THREADS_PER_BLOCK 128
 #define TILE_SIZE 32
-#define SHMEM_SIZE THREADS_PER_BLOCK * sizeof(float)
 
 __host__ void cpu_to_cuda(Tensor* tensor) {
     
@@ -61,7 +60,7 @@ __host__ void add_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_da
 
 
 __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size) {
-    __shared__ float partial_sum[SHMEM_SIZE];
+    __shared__ float partial_sum[THREADS_PER_BLOCK * sizeof(float)];
 
     int tid = threadIdx.x;
     int i = blockIdx.x * blockDim.x + threadIdx.x;

norch/csrc/tensor.cpp

@@ -60,6 +60,15 @@ extern "C" {
         }
         printf("]\n");
 
+        printf("Strides: [");
+        for (int i = 0; i < ndim; i++) {
+            printf("%d", tensor->strides[i]);
+            if (i < ndim - 1) {
+                printf(", ");
+            }
+        }
+        printf("]\n");
+
         /*printf("Data:\n[");
         for (int i = 0; i < stride; i++) {
             printf("%.2f", tensor->data[i]);
@@ -327,9 +336,10 @@ extern "C" {
     }
 
     Tensor* matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
+        //MxN @ NxP = MxP
         // Check if tensors have compatible shapes for matrix multiplication
         if (tensor1->shape[1] != tensor2->shape[0]) {
-            fprintf(stderr, "Incompatible shapes for matrix multiplication\n");
+            fprintf(stderr, "Incompatible shapes for matrix multiplication %dx%d and %dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1]);
             exit(1);
         }
 
@@ -387,6 +397,137 @@ extern "C" {
         }
     }
 
+    Tensor* broadcasted_batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
+        //MxN @ BATCHxNxP = BATCHxMxP
+        // Check if tensors have compatible shapes for matrix multiplication
+        if (tensor1->shape[1] != tensor2->shape[1]) {
+            fprintf(stderr, "Incompatible shapes for broadcasted batched matrix multiplication %dx%d and %dx%dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1], tensor2->shape[2]);
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        char* device = (char*)malloc(strlen(tensor1->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor1->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+
+        int ndim = 3;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        
+        shape[0] = tensor2->shape[0];;
+        shape[1] = tensor1->shape[0];
+        shape[2] = tensor2->shape[2];  
+
+        int size = 1;
+        for (int i = 0; i < ndim; i++) {
+            size *= shape[i];
+        }
+
+        float* result_data = (float*)malloc(size * sizeof(float));
+        if (result_data == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, size * sizeof(float));
+            ////broadcasted_batched_matmul_tensor_cuda(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            broadcasted_batched_matmul_tensor_cpu(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+
+    }
+
+    Tensor* batched_matmul_tensor(Tensor* tensor1, Tensor* tensor2) {
+        //BATCHxMxN @ BATCHxNxP = BATCHxMxP
+        // Check if tensors have compatible shapes for matrix multiplication
+
+        if (tensor1->shape[0] != tensor2->shape[0]) {
+            fprintf(stderr, "Tensors must have same batch dimension for batch matmul %d and %d\n", tensor1->shape[0], tensor2->shape[0]);
+            exit(1);
+        }
+
+        if (tensor1->shape[2] != tensor2->shape[1]) {
+            fprintf(stderr, "Incompatible shapes for matrix multiplication %dx%d and %dx%d\n", tensor1->shape[0], tensor1->shape[1], tensor2->shape[0], tensor2->shape[1]);
+            exit(1);
+        }
+        
+        if (strcmp(tensor1->device, tensor2->device) != 0) {
+            fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
+            exit(1);
+        }
+
+        char* device = (char*)malloc(strlen(tensor1->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor1->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+
+        int ndim = 3;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+        
+        shape[0] = tensor2->shape[0];;
+        shape[1] = tensor1->shape[1];
+        shape[2] = tensor2->shape[2];  
+
+        int size = 1;
+        for (int i = 0; i < ndim; i++) {
+            size *= shape[i];
+        }
+
+        float* result_data = (float*)malloc(size * sizeof(float));
+        if (result_data == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        if (strcmp(tensor1->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, size * sizeof(float));
+            //batched_matmul_tensor_cuda(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            batched_matmul_tensor_cpu(tensor1, tensor2, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+
+    }
+
+
     Tensor* pow_tensor(Tensor* tensor, float power) {
         char* device = (char*)malloc(strlen(tensor->device) + 1);
         if (device != NULL) {
@@ -433,14 +574,25 @@ extern "C" {
             exit(-1);
         }
 
+        int ndim = new_ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            shape[i] = new_shape[i];
+        }
+
         // Calculate the total number of elements in the new shape
-        int new_size = 1;
+        int size = 1;
         for (int i = 0; i < new_ndim; i++) {
-            new_size *= new_shape[i];
+            size *= shape[i];
         }
 
         // Check if the total number of elements matches the current tensor's size
-        if (new_size != tensor->size) {
+        if (size != tensor->size) {
             fprintf(stderr, "Cannot reshape tensor. Total number of elements in new shape does not match the current size of the tensor.\n");
             exit(1);
         }
@@ -450,7 +602,7 @@ extern "C" {
             float* result_data;
             cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
             assign_tensor_cuda(tensor, result_data);
-            return create_tensor(result_data, new_shape, new_ndim, device);
+            return create_tensor(result_data, shape, ndim, device);
         } 
         else {
             float* result_data = (float*)malloc(tensor->size * sizeof(float));
@@ -459,7 +611,7 @@ extern "C" {
                 exit(1);
             }
             assign_tensor_cpu(tensor, result_data);
-            return create_tensor(result_data, new_shape, new_ndim, device);
+            return create_tensor(result_data, shape, ndim, device);
         }
     }
 
@@ -572,8 +724,112 @@ extern "C" {
                 fprintf(stderr, "Memory allocation failed\n");
                 exit(1);
             }
-            transpose_tensor_cpu(tensor, result_data);
+            switch (ndim) {
+                case 1:
+                    transpose_1D_tensor_cpu(tensor, result_data);
+                    break;
+                case 2:
+                    transpose_2D_tensor_cpu(tensor, result_data);
+                    break;
+                case 3:
+                    transpose_3D_tensor_cpu(tensor, result_data);
+                    break;
+                default:
+                    fprintf(stderr, "Transpose only supports tensors up to 3 dimensions.\n");
+                    exit(-1);
+            }
             return create_tensor(result_data, shape, ndim, device);
         }
     }
-}
+
+    Tensor* transpose_axes_tensor(Tensor* tensor, int axis1, int axis2) {
+        char* device = (char*)malloc(strlen(tensor->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+
+        int ndim = tensor->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            shape[i] = tensor->shape[i];
+        }
+
+        shape[axis1] = tensor->shape[axis2];
+        shape[axis2] = tensor->shape[axis1];
+
+        int size = tensor->size;
+
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, size * sizeof(float));
+            //transpose_axes_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);
+            }
+            //transpose_axes_cpu(tensor, result_data, axis1, axis2, shape);
+            assign_tensor_cpu(tensor, result_data);
+            
+            Tensor* new_tensor = create_tensor(result_data, shape, ndim, device);
+            for (int i = 0; i < ndim; i++) {
+                new_tensor->strides[i] = tensor->strides[i];
+            }
+            new_tensor->strides[axis1] = tensor->strides[axis2];
+            new_tensor->strides[axis2] = tensor->strides[axis1];
+            make_contiguous(new_tensor);
+            return new_tensor;
+        }
+    }
+
+    void make_contiguous(Tensor* tensor) {
+        float* new_data = (float*)malloc(tensor->size * sizeof(float));
+        if (new_data == NULL) {
+            // Handle memory allocation failure
+            return;
+        }
+
+        int* new_strides = (int*)malloc(tensor->ndim * sizeof(int));
+        if (new_strides == NULL) {
+            free(new_data);
+            // Handle memory allocation failure
+            return;
+        }
+
+        // Calculate new strides assuming C-contiguous order
+        int stride = 1;
+        for (int i = tensor->ndim - 1; i >= 0; i--) {
+            new_strides[i] = stride;
+            stride *= tensor->shape[i];
+        }
+
+        // Rearrange data
+        for (int i = 0; i < tensor->size; i++) {
+            int index = 0;
+            int offset = i;
+            for (int j = 0; j < tensor->ndim; j++) {
+                index += (offset / new_strides[j]) * tensor->strides[j];
+                offset %= new_strides[j];
+            }
+            new_data[i] = tensor->data[index];
+        }
+
+        // Free old data and update tensor properties
+        free(tensor->data);
+        free(tensor->strides);
+        tensor->data = new_data;
+        tensor->strides = new_strides;
+        }
+    }

norch/csrc/tensor.h

@@ -26,6 +26,9 @@ extern "C" {
     void to_device(Tensor* tensor, char* device);
     Tensor* ones_like_tensor(Tensor* tensor);
     Tensor* zeros_like_tensor(Tensor* tensor);
+    Tensor* transpose_tensor(Tensor* tensor);
+    Tensor* transpose_axes_tensor(Tensor* tensor, int axis1, int axis2);
+    void make_contiguous(Tensor* tensor);
 }
 
 #endif /* TENSOR_H */

norch/tensor.py

@@ -122,7 +122,7 @@ class Tensor:
         if requires_grad:
             self.grad_fn = ReshapeBackward(self)
 
-        return self
+        return result_data
     
     def to(self, device):
         self.device = device
@@ -291,22 +291,52 @@ class Tensor:
         return self  
     
     def __matmul__(self, other):
-        if self.ndim != 2 or other.ndim != 2:
-            raise ValueError("Matrix multiplication requires 2D tensors")
+        if self.ndim < 3 and other.ndim == 3:
+            #broadcasted 2D x 3D matmul
 
-        if self.shape[1] != other.shape[0]:
-            raise ValueError("Incompatible shapes for matrix multiplication")
+            Tensor._C.broadcasted_batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.broadcasted_batched_matmul_tensor.restype = ctypes.POINTER(CTensor)
+            
+            result_tensor_ptr = Tensor._C.broadcasted_batched_matmul_tensor(self.tensor, other.tensor)
 
-        Tensor._C.matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
-        Tensor._C.matmul_tensor.restype = ctypes.POINTER(CTensor)
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = [other.shape[0], self.shape[0], other.shape[2]]
+            result_data.ndim = 3
+            result_data.device = self.device
 
-        result_tensor_ptr = Tensor._C.matmul_tensor(self.tensor, other.tensor)
+        elif self.ndim == 3 and other.ndim == 3:
+            #broadcasted 3D x 3D matmul
 
-        result_data = Tensor()
-        result_data.tensor = result_tensor_ptr
-        result_data.shape = [self.shape[0], other.shape[1]]
-        result_data.ndim = 2
-        result_data.device = self.device
+            Tensor._C.batched_matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.batched_matmul_tensor.restype = ctypes.POINTER(CTensor)
+            
+            result_tensor_ptr = Tensor._C.batched_matmul_tensor(self.tensor, other.tensor)
+
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = [other.shape[0], self.shape[1], other.shape[2]]
+            result_data.ndim = 3
+            result_data.device = self.device
+            
+        else:
+            #2D matmul
+            if self.ndim != 2 or other.ndim != 2:
+                raise ValueError("Matrix multiplication requires 2D tensors")
+
+            if self.shape[1] != other.shape[0]:
+                raise ValueError("Incompatible shapes for matrix multiplication")
+
+            Tensor._C.matmul_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
+            Tensor._C.matmul_tensor.restype = ctypes.POINTER(CTensor)
+
+            result_tensor_ptr = Tensor._C.matmul_tensor(self.tensor, other.tensor)
+
+            result_data = Tensor()
+            result_data.tensor = result_tensor_ptr
+            result_data.shape = [self.shape[0], other.shape[1]]
+            result_data.ndim = 2
+            result_data.device = self.device
 
         result_data.requires_grad = self.requires_grad or other.requires_grad
         if result_data.requires_grad:
@@ -353,11 +383,34 @@ class Tensor:
 
         return result_data
     
+    def transpose(self, axis1, axis2):
+        if axis1 < 0:
+            axis1 = self.ndim + axis1
+        if axis2 < 0:
+            axis2 = self.ndim + axis2
+
+        Tensor._C.transpose_axes_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int, ctypes.c_int]
+        Tensor._C.transpose_axes_tensor.restype = ctypes.POINTER(CTensor)
+
+        result_tensor_ptr = Tensor._C.transpose_axes_tensor(self.tensor, axis1, axis2)
+
+        result_data = Tensor()
+        result_data.tensor = result_tensor_ptr
+        result_data.shape = self.shape.copy()
+        result_data.shape[axis1] = self.shape[axis2]
+        result_data.shape[axis2] = self.shape[axis1]
+        result_data.ndim = self.ndim
+        result_data.device = self.device
+
+        result_data.requires_grad = self.requires_grad
+        if result_data.requires_grad:
+            result_data.grad_fn = TransposeBackward(self, axis1, axis2)
+
+        return result_data
+
+    
     @property
     def T(self):
-        if self.ndim != 2:
-            raise ValueError("Transpose requires 2D tensors")
-
         Tensor._C.transpose_tensor.argtypes = [ctypes.POINTER(CTensor)]
         Tensor._C.transpose_tensor.restype = ctypes.POINTER(CTensor)
 
@@ -365,8 +418,8 @@ class Tensor:
 
         result_data = Tensor()
         result_data.tensor = result_tensor_ptr
-        result_data.shape = [self.shape[1], self.shape[0]]
-        result_data.ndim = 2
+        result_data.shape = self.shape.copy()[::-1]
+        result_data.ndim = self.ndim
         result_data.device = self.device
 
         return result_data

test.py

@@ -20,8 +20,88 @@ if __name__ == "__main__":
     import random
     import numpy as np
 
+    """a = norch.Tensor([
+        [[1.234, 2.123], [3.635, 4.456], [5.678, 6.789]],
+        [[7.890, 8.901], [9.012, 1.234], [2.345, 3.456]],
+        [[4.567, 5.678], [6.789, 7.890], [8.901, 9.012]],
+        [[1.234, 2.345], [3.456, 4.567], [5.678, 6.789]],
+        [[7.890, 8.901], [9.012, 1.234], [2.345, 3.456]]
+    ], requires_grad=True)
+
+    b = norch.Tensor([[
+        [1.234, 2.123, 1.5],
+        [5.678, 6.789, 1.293],
+        [3.635, 4.456, 1.0202],
+        [7.890, 8.901, 1.91],
+    ],[
+        [1.234, 2.123, 1.5],
+        [5.678, 6.789, 1.293],
+        [3.635, 4.456, 1.0202],
+        [7.890, 8.901, 1.91],
+    ],[
+        [1.234, 2.123, 1.5],
+        [5.678, 6.789, 1.293],
+        [3.635, 4.456, 1.0202],
+        [7.890, 8.901, 1.91],
+    ],[
+        [1.234, 2.123, 1.5],
+        [5.678, 6.789, 1.293],
+        [3.635, 4.456, 1.0202],
+        [7.890, 8.901, 1.91],
+    ],[
+        [1.234, 2.123, 1.5],
+        [5.678, 6.789, 1.293],
+        [3.635, 4.456, 1.0202],
+        [7.890, 8.901, 1.91],
+    ]])
+
+    #print(a.T)
+    #print(a.T.shape)
+    #[5, 3, 2] [4, 3] [5, 4, 2]
+    #print(a.shape, b.shape)
+    #b = norch.Tensor([
+    #    [1.234, 2.123, 1.5]])
+    result = b @ a
+    result = result.sum()
+    result.backward()
+
+    print(a.grad)"""
+
+    #### testar transpose axes!!!! make it contiguous
+
+    tensor1 = norch.Tensor([[[1, 2], [3, 4], [5, 6]],
+                        [[7, 8], [9, 10], [11, 12]],
+                        [[13, 14], [15, 16], [17, 18]],
+                        [[19, 20], [21, 22], [23, 24]],
+                        [[25, 26], [27, 28], [29, 30]]], requires_grad=True)
+
+    # Reshape tensor1 to 2x3x5
+    reshaped_tensor = tensor1.transpose(1, 0)
+
+    # Create a 5x4 tensor
+    tensor2 = norch.Tensor([[1, 2, 3],
+                            [5, 6, 7],
+                            [9, 10, 11],
+                            [13, 14, 15],
+                            [17, 18, 19]])
+
+    tensor2 = tensor2.transpose(1,0)
     
-    #a = Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda")
+
+    # Multiply reshaped_tensor by tensor2
+    result = tensor2 @ reshaped_tensor
+
+    result = result.sum()
+    result.backward()
+    print(tensor1.grad)
+
+    #print(a.shape, b.shape, result.shape)
+    #c = result.sum()
+    #c.backward()
+    #print(a.grad)
+    #print(a.transpose(2,1))
+    
+    #a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda")
     #b = Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda")
     #c = Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda")
 
@@ -37,7 +117,6 @@ if __name__ == "__main__":
     print(a.grad)"""
     
     """#print(a)
-    """
     N = 10
     a = norch.Tensor([[1 for _ in range(N)] for _ in range(N)])
     #b = norch.Tensor([[random.uniform(0, 1) for _ in range(N)] for _ in range(N)])
@@ -56,7 +135,7 @@ if __name__ == "__main__":
     print("\n\n")
 
     """
-    
+    """
     a = [[random.uniform(0, 1) for _ in range(N)] for _ in range(N)]
     b = [[random.uniform(0, 1) for _ in range(N)] for _ in range(N)]
     ini = time.time()