add support for add
MODIFY
tensor.py
tensor.py
@@ -78,6 +78,15 @@ class Mul(Function):
return y*grad_output, x*grad_output
register('mul', Mul)
+class Add(Function):
+ @staticmethod
+ def forward(ctx, x, y):
+ return x+y
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ return grad_output, grad_output
+register('add', Add)
class ReLU(Function):
@staticmethod

MODIFY
test.py
test.py
@@ -14,7 +14,8 @@ def test_tinygrad():
outr = out.relu()
outl = outr.logsoftmax()
outm = outl.mul(m)
- outx = outm.sum()
+ outa = outm.add(m)
+ outx = outa.sum()
outx.backward()
return outx.data, x.grad, W.grad
@@ -26,7 +27,8 @@ def test_pytorch():
outr = out.relu()
outl = torch.nn.functional.log_softmax(outr, dim=1)
outm = outl.mul(m)
- outx = outm.sum()
+ outa = outm.add(m)
+ outx = outa.sum()
outx.backward()
return outx.detach().numpy(), x.grad, W.grad

add support for add