transpose n dimensions cpu

norch/csrc/cpu.cpp

@@ -116,7 +116,7 @@ void zeros_like_tensor_cpu(Tensor* tensor, float* result_data) {
     }
 }
 
-void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
+/*void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
     int rows = tensor->shape[0];
     int cols = tensor->shape[1];
 
@@ -125,6 +125,36 @@ void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
             result_data[j * rows + i] = tensor->data[i * cols + j];
         }
     }
+}*/
+
+void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
+    int* shape = tensor->shape;
+    int ndim = tensor->ndim;
+    int* strides = (int*)malloc(ndim * sizeof(int));
+    int* indices = (int*)calloc(ndim, sizeof(int));
+
+    strides[ndim - 1] = 1;
+    for (int i = ndim - 2; i >= 0; i--) {
+        strides[i] = strides[i + 1] * shape[i + 1];
+    }
+
+    int idx_result;
+    for (int idx_source = 0; idx_source < tensor->size; idx_source++) {
+        idx_result = 0;
+        for (int dim = 0; dim < ndim; dim++) {
+            idx_result += indices[dim] * strides[dim];
+        }
+        result_data[idx_result] = tensor->data[idx_source];
+
+        // Update indices
+        indices[ndim - 1]++;
+        for (int dim = ndim - 1; dim > 0; dim--) {
+            if (indices[dim] == shape[dim]) {
+                indices[dim] = 0;
+                indices[dim - 1]++;
+            }
+        }
+    }
 }
 
 void assign_tensor_cpu(Tensor* tensor, float* result_data) {

norch/tensor.py

@@ -385,9 +385,6 @@ class Tensor:
     
     @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)
 
@@ -395,8 +392,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[::-1]
+        result_data.ndim = self.ndim
         result_data.device = self.device
 
         return result_data

test.py

@@ -55,8 +55,12 @@ if __name__ == "__main__":
         [7.890, 8.901, 5.91],
     ]])
 
-    result = b @ a
-    print(result)
+    b = norch.Tensor([
+        [1.234, 2.123, 1.5]])
+
+    print(a.shape)
+    result = a.T
+    print(result.shape)
     #c = result.sum()
     #c.backward()
     #print(a.grad)