Reshape autograd

norch/autograd/functions.py

@@ -51,4 +51,11 @@ class SumBackward:
 
     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()]
+        return [float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()]
+    
+class ReshapeBackward:
+    def __init__(self, x):
+        self.input = [x]
+
+    def backward(self, gradient):
+        return [gradient.reshape(self.input[0].shape)]

norch/csrc/cpu.cpp

@@ -86,3 +86,10 @@ void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
         }
     }
 }
+
+void assign_tensor_cpu(Tensor* tensor, float* result_data) {
+
+    for (int i = 0; i < tensor->size; i++) {
+        result_data[i] = tensor->data[i];
+    }
+}

norch/csrc/cpu.h

@@ -13,5 +13,6 @@ 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 assign_tensor_cpu(Tensor* tensor, float* result_data);
 
 #endif /* CPU_H */

norch/csrc/cuda.cu

@@ -314,4 +314,26 @@ __host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data) {
     cudaDeviceSynchronize();
 }
 
+__global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = data[i];
+    }
+}
+
+__host__ void assign_tensor_cuda(Tensor* tensor, float* result_data) {
+    
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    assign_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
 

norch/csrc/cuda.h

@@ -34,6 +34,10 @@
     __global__ void transpose_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols);
     __host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data);
 
+    __global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size);
+    __host__ void assign_tensor_cuda(Tensor* tensor, float* result_data);
+
+
 
 
 

norch/csrc/tensor.cpp

@@ -425,7 +425,15 @@ extern "C" {
         }
     }
 
-    void reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim) {
+    Tensor* reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim) {
+        char* device = (char*)malloc(strlen(tensor->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+
         // Calculate the total number of elements in the new shape
         int new_size = 1;
         for (int i = 0; i < new_ndim; i++) {
@@ -438,28 +446,21 @@ extern "C" {
             exit(1);
         }
 
-        // Update the shape
-        tensor->shape = (int*)malloc(new_ndim * sizeof(int));
-        if (tensor->shape == NULL) {
-            fprintf(stderr, "Memory allocation failed\n");
-            exit(1);
-        }
-        for (int i = 0; i < new_ndim; i++) {
-            tensor->shape[i] = new_shape[i];
-        }
-        tensor->ndim = new_ndim;
+        if (strcmp(tensor->device, "cuda") == 0) {
 
-        // Update the strides
-        tensor->strides = (int*)malloc(new_ndim * sizeof(int));
-        if (tensor->strides == NULL) {
-            fprintf(stderr, "Memory allocation failed\n");
-            exit(1);
-        }
-
-        int stride = 1;
-        for (int i = new_ndim - 1; i >= 0; i--) {
-            tensor->strides[i] = stride;
-            stride *= new_shape[i];
+            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);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            assign_tensor_cpu(tensor, result_data);
+            return create_tensor(result_data, new_shape, new_ndim, device);
         }
     }
 

norch/csrc/tensor.h

@@ -22,7 +22,7 @@ extern "C" {
     Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2);
     Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2);
     Tensor* scalar_mul_tensor(Tensor* tensor, float scalar);
-    void reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim);
+    Tensor* reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim);
     Tensor* matmul_tensor(Tensor* tensor1, Tensor* tensor2);
     Tensor* pow_tensor(Tensor* tensor, float power);
     void to_device(Tensor* tensor, char* device);

norch/tensor.py

@@ -93,7 +93,7 @@ class Tensor:
         Tensor._C.zeros_like_tensor.restype = ctypes.POINTER(CTensor)
         Tensor._C.zeros_like_tensor(self.tensor)   
 
-        result_tensor_ptr = Tensor._C.ones_like_tensor(self.tensor)
+        result_tensor_ptr = Tensor._C.zeros_like_tensor(self.tensor)
 
         result_data = Tensor()
         result_data.tensor = result_tensor_ptr
@@ -103,17 +103,24 @@ class Tensor:
         
         return result_data
     
-    def reshape(self, new_shape):
+    def reshape(self, new_shape, requires_grad=None):
 
         new_shape_ctype = (ctypes.c_int * len(new_shape))(*new_shape)
         new_ndim_ctype = ctypes.c_int(len(new_shape))
         
         Tensor._C.reshape_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(ctypes.c_int), ctypes.c_int]
-        Tensor._C.reshape_tensor.restype = None
-        Tensor._C.reshape_tensor(self.tensor, new_shape_ctype, new_ndim_ctype)   
+        Tensor._C.reshape_tensor.restype = ctypes.POINTER(CTensor)
+        result_tensor_ptr = Tensor._C.reshape_tensor(self.tensor, new_shape_ctype, new_ndim_ctype)   
 
-        self.shape = new_shape
-        self.ndim = len(new_shape)
+        result_data = Tensor()
+        result_data.tensor = result_tensor_ptr
+        result_data.shape = new_shape.copy()
+        result_data.ndim = len(new_shape)
+        result_data.device = self.device
+
+        result_data.requires_grad = self.requires_grad
+        if requires_grad:
+            self.grad_fn = ReshapeBackward(self)
 
         return self
     

test.py

@@ -30,12 +30,11 @@ if __name__ == "__main__":
     a = norch.Tensor([[1, 2], [1, 2], [1, 2]], requires_grad=True)#.to("cuda")
     b = norch.Tensor([[1, 400, 3], [1, 2, 3]], requires_grad=True)
 
-    print(a.sum())
-    #c = (b @ a) * 5
-    #d = c.sum()
-    #d.backward()
+    c = (a@b).reshape([9])
+    d = c.sum()
+    d.backward()
 
-    #print(a.grad)
+    print(a.grad)
     
     """#print(a)
     N = 1000