autograd other functions

norch/autograd/functions.py

@@ -4,3 +4,36 @@ class AddBackward:
 
     def backward(self, gradient):
         return [gradient, gradient]
+    
+class SubBackward:
+    def __init__(self, x, y):
+        self.tensors = [x, y]
+
+    def backward(self, gradient):
+        return [gradient, -gradient]
+    
+class ScalarMulBackward:
+    def __init__(self, x, scalar):
+        self.tensors = [x]
+        self.scalar = scalar
+
+    def backward(self, gradient):
+        return [gradient * self.scalar]
+
+
+class ElementwiseMulBackward:
+    def __init__(self, x, y):
+        self.tensors = [x, y]
+
+    def backward(self, gradient):
+        return [gradient * self.tensors[1], gradient * self.tensors[0]]
+    
+class SumBackward:
+    def __init__(self, x):
+        self.tensor = x
+
+    def backward(self, gradient):
+        # Since sum reduces a tensor to a scalar, gradient is broadcasted to match the original shape.
+        return self.tensor.ones_like() * float(gradient.tensor.contents.data.value)
+
+

norch/tensor.py

@@ -133,8 +133,13 @@ class Tensor:
             return
         
         if gradient is None:
+            if self.shape != [1]:
+                raise RuntimeError("Gradient argument must be specified for non-scalar tensors.")
             gradient = self.ones_like()
         
+        if self.shape != [1]:
+            raise RuntimeError("Only scalar tensors can be used to trigger backward propagation.")
+        
         if self.grad is None:
             self.grad = gradient
         else:
@@ -229,6 +234,10 @@ class Tensor:
         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 = SubBackward(self, other)
+
         return result_data
     
     def __mul__(self, other):
@@ -243,6 +252,10 @@ class Tensor:
 
             result_data.tensor = Tensor._C.scalar_mul_tensor(self.tensor, ctypes.c_float(other))
 
+            result_data.requires_grad = self.requires_grad
+            if result_data.requires_grad:
+                result_data.grad_fn = ScalarMulBackward(self, other)
+
             return result_data
         elif isinstance(other, Tensor):
             if self.shape != other.shape:
@@ -259,6 +272,10 @@ class Tensor:
             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 = ElementwiseMulBackward(self, other)
+
             return result_data
         else:
             raise TypeError("Unsupported operand type(s) for *: '{}' and '{}'".format(type(self), type(other)))
@@ -321,6 +338,8 @@ class Tensor:
         result_data.ndim = 1
         result_data.device = self.device
 
-        
+        result_data.requires_grad = self.requires_grad
+        if result_data.requires_grad:
+            result_data.grad_fn = SumBackward(self)
 
         return result_data

test.py

@@ -17,7 +17,7 @@ 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, 2, 3], [1, 2, 3], [1, 2, 3]], requires_grad=True)
+    b = norch.Tensor([[1, 400, 3], [1, 2, 3], [1, 2, 3]], requires_grad=True)
     import time
     import random
     import numpy as np
@@ -29,7 +29,7 @@ if __name__ == "__main__":
 
     #d = b-c
 
-    c = (a + b)
+    c = a.sum()
     c.backward()
     print(a.grad)
     print(b.grad)