autograd add function

norch/autograd/functions.py

@@ -0,0 +1,6 @@
+class AddBackward:
+    def __init__(self, x, y):
+        self.tensors = [x, y]
+
+    def backward(self, gradient):
+        return [gradient, gradient]

norch/csrc/cpu.cpp

@@ -61,3 +61,17 @@ void sum_tensor_cpu(Tensor* tensor, float* result_data) {
 
     *result_data = sum;
 }
+
+void ones_like_tensor_cpu(Tensor* tensor, float* result_data) {
+    
+    for (int i = 0; i < tensor->size; i++) {
+        result_data[i] = 1.0;
+    }
+}
+
+void zeros_like_tensor_cpu(Tensor* tensor, float* result_data) {
+    
+    for (int i = 0; i < tensor->size; i++) {
+        result_data[i] = 0.0;
+    }
+}

norch/csrc/cpu.h

@@ -10,5 +10,7 @@ void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_
 void 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);
 
 #endif /* CPU_H */

norch/csrc/cuda.cu

@@ -223,4 +223,48 @@ __host__ void pow_tensor_cuda(Tensor* tensor, float power, float* result_data) {
     cudaDeviceSynchronize();
 }
 
+__global__ void ones_like_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = 1.0;
+    }
+}
+
+__host__ void ones_like_tensor_cuda(Tensor* tensor, float* result_data) {
+    
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    ones_like_tensor_cuda_kernel<<<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();
+}
+
+__global__ void zeros_like_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = 0.0;
+    }
+}
+
+__host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data) {
+    
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    zeros_like_tensor_cuda_kernel<<<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

@@ -25,4 +25,11 @@
     __global__ void pow_tensor_cuda_kernel(float* data, float power, float* result_data, int size);
     __host__ void pow_tensor_cuda(Tensor* tensor, float power, float* result_data);
 
+    __global__ void ones_like_tensor_cuda_kernel(float* data, float* result_data, int size);
+    __host__ void ones_like_tensor_cuda(Tensor* tensor, float* result_data);
+
+    __global__ void zeros_like_tensor_cuda_kernel(float* data, float* result_data, int size);
+    __host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data);
+
+
 #endif /* CUDA_KERNEL_H_ */

norch/csrc/tensor.cpp

@@ -461,4 +461,78 @@ extern "C" {
             stride *= new_shape[i];
         }
     }
-}
+}
+
+    Tensor* ones_like_tensor(Tensor* tensor) {
+        char* device = (char*)malloc(strlen(tensor->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            shape[i] = tensor->shape[i];
+        }
+
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
+            ones_like_tensor_cuda(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            ones_like_tensor_cpu(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }
+
+    Tensor* zeros_like_tensor(Tensor* tensor) {
+        char* device = (char*)malloc(strlen(tensor->device) + 1);
+        if (device != NULL) {
+            strcpy(device, tensor->device);
+        } else {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(-1);
+        }
+        int ndim = tensor->ndim;
+        int* shape = (int*)malloc(ndim * sizeof(int));
+        if (shape == NULL) {
+            fprintf(stderr, "Memory allocation failed\n");
+            exit(1);
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            shape[i] = tensor->shape[i];
+        }
+
+        if (strcmp(tensor->device, "cuda") == 0) {
+
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
+            zeros_like_tensor_cuda(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        } 
+        else {
+            float* result_data = (float*)malloc(tensor->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            zeros_like_tensor_cpu(tensor, result_data);
+            return create_tensor(result_data, shape, ndim, device);
+        }
+    }

norch/csrc/tensor.h

@@ -24,6 +24,8 @@ extern "C" {
     Tensor* matmul_tensor(Tensor* tensor1, Tensor* tensor2);
     Tensor* pow_tensor(Tensor* tensor, float power);
     void to_device(Tensor* tensor, char* device);
+    Tensor* ones_like_tensor(Tensor* tensor);
+    Tensor* zeros_like_tensor(Tensor* tensor);
 }
 
 #endif /* TENSOR_H */

norch/tensor.py

@@ -1,5 +1,6 @@
 import ctypes
 import os
+from .autograd.functions import *
 
 class CTensor(ctypes.Structure):
     _fields_ = [
@@ -15,7 +16,7 @@ class Tensor:
     os.path.abspath(os.curdir)
     _C = ctypes.CDLL(os.path.join(os.path.abspath(os.curdir), "build/libtensor.so"))
 
-    def __init__(self, data=None, device="cpu"):
+    def __init__(self, data=None, device="cpu", requires_grad=False):
 
         if data != None:
             data, shape = self.flatten(data)
@@ -29,6 +30,10 @@ class Tensor:
             self.ndim = len(shape)
             self.device = device
 
+            self.requires_grad = requires_grad
+            self.grad = None
+            self.grad_fn = None
+
             Tensor._C.create_tensor.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_char_p]
             Tensor._C.create_tensor.restype = ctypes.POINTER(CTensor)
             
@@ -45,6 +50,10 @@ class Tensor:
             self.shape = None,
             self.ndim = None,
             self.device = device
+            self.requires_grad = None
+            self.grad = None
+            self.grad_fn = None
+
 
     def flatten(self, nested_list):
         def flatten_recursively(nested_list):
@@ -63,6 +72,38 @@ class Tensor:
         flat_data, shape = flatten_recursively(nested_list)
         return flat_data, shape
     
+    def ones_like(self):
+        
+        Tensor._C.ones_like_tensor.argtypes = [ctypes.POINTER(CTensor)]
+        Tensor._C.ones_like_tensor.restype = ctypes.POINTER(CTensor)
+        Tensor._C.ones_like_tensor(self.tensor)   
+
+        result_tensor_ptr = Tensor._C.ones_like_tensor(self.tensor)
+
+        result_data = Tensor()
+        result_data.tensor = result_tensor_ptr
+        result_data.shape = self.shape.copy()
+        result_data.ndim = self.ndim
+        result_data.device = self.device
+        
+        return result_data
+    
+    def zeros_like(self):
+        
+        Tensor._C.zeros_like_tensor.argtypes = [ctypes.POINTER(CTensor)]
+        Tensor._C.zeros_like_tensor.restype = ctypes.POINTER(CTensor)
+        Tensor._C.zeros_like_tensor(self.tensor)   
+
+        result_tensor_ptr = Tensor._C.ones_like_tensor(self.tensor)
+
+        result_data = Tensor()
+        result_data.tensor = result_tensor_ptr
+        result_data.shape = self.shape.copy()
+        result_data.ndim = self.ndim
+        result_data.device = self.device
+        
+        return result_data
+    
     def reshape(self, new_shape):
 
         new_shape_ctype = (ctypes.c_int * len(new_shape))(*new_shape)
@@ -86,6 +127,30 @@ class Tensor:
         Tensor._C.to_device(self.tensor, self.device_ctype)
 
         return self
+    
+    def backward(self, gradient=None):
+        if not self.requires_grad:
+            return
+        
+        if gradient is None:
+            gradient = self.ones_like()
+        
+        if self.grad is None:
+            self.grad = gradient
+        else:
+            self.grad += gradient
+
+        if self.grad_fn is not None:
+            grads = self.grad_fn.backward(gradient)
+            if len(grads) == 1:
+                self.grad = grads[0]
+            else:
+                for tensor, grad in zip(self.grad_fn.tensors, grads):
+                    tensor.backward(grad)
+
+    
+    def zero_grad(self):
+        self.grad = None
 
     def __getitem__(self, indices):
         if len(indices) != self.ndim:
@@ -122,7 +187,7 @@ class Tensor:
         index = [0] * self.ndim
         result = "tensor(["
         result += print_recursively(self, 0, index)
-        result += f"""], device="{self.device}")"""
+        result += f"""], device="{self.device}", requires_grad={self.requires_grad})"""
         return result
 
     def __repr__(self):
@@ -142,6 +207,10 @@ class Tensor:
         result_data.shape = self.shape.copy()
         result_data.ndim = self.ndim
         result_data.device = self.device
+        
+        result_data.requires_grad = self.requires_grad or other.requires_grad
+        if result_data.requires_grad:
+            result_data.grad_fn = AddBackward(self, other)
 
         return result_data
     
@@ -252,4 +321,6 @@ class Tensor:
         result_data.ndim = 1
         result_data.device = self.device
 
+        
+
         return result_data

test.py

@@ -16,8 +16,8 @@ def matrix_sum(matrix1, matrix2):
 
 if __name__ == "__main__":
     import norch
-    a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])#.to("cuda")
-    b = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
+    a = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]], requires_grad=True)#.to("cuda")
+    b = norch.Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]], requires_grad=True)
     import time
     import random
     import numpy as np
@@ -29,8 +29,10 @@ if __name__ == "__main__":
 
     #d = b-c
 
-    b =    a.sum()
-    print(b)
+    c = (a + b)
+    c.backward()
+    print(a.grad)
+    print(b.grad)
     #print(a ** 2)
     
     """#print(a)