pow backward autograd

norch/autograd/functions.py

@@ -28,6 +28,15 @@ class ElementwiseMulBackward:
     def backward(self, gradient):
         return [gradient * self.input[1], gradient * self.input[0]]
     
+class PowBackward:
+    def __init__(self, x, power):
+        self.input = [x]
+        self.power = power
+
+    def backward(self, gradient):
+        print(self.input[0], "@@@")
+        return [(gradient * self.power) * (self.input[0]) ** (self.power - 1)]
+       
 class SumBackward:
     def __init__(self, x):
         self.input = [x]

norch/tensor.py

@@ -306,12 +306,12 @@ class Tensor:
 
     def __pow__(self, power):
         
-        power = ctypes.c_float(power)
+        power_ctypes = ctypes.c_float(power)
 
         Tensor._C.pow_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_float]
         Tensor._C.pow_tensor.restype = ctypes.POINTER(CTensor)
 
-        result_tensor_ptr = Tensor._C.pow_tensor(self.tensor, power)
+        result_tensor_ptr = Tensor._C.pow_tensor(self.tensor, power_ctypes)
 
         result_data = Tensor()
         result_data.tensor = result_tensor_ptr
@@ -319,6 +319,10 @@ class Tensor:
         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 = PowBackward(self, power)
+
         return result_data
     
     def sum(self):

test.py

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