diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index 335bdd6..ab543df 100644 Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc index 35e8a00..4169576 100644 Binary files a/norch/autograd/__pycache__/functions.cpython-38.pyc and b/norch/autograd/__pycache__/functions.cpython-38.pyc differ diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py index a2a3344..67bc301 100644 --- a/norch/autograd/functions.py +++ b/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) + + diff --git a/norch/tensor.py b/norch/tensor.py index 265cb6b..5103ec1 100644 --- a/norch/tensor.py +++ b/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 \ No newline at end of file diff --git a/test.py b/test.py index 0dc6a80..56e4173 100644 --- a/test.py +++ b/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)