have you considered using tinygrad
@@ -5,7 +5,9 @@ import numpy as np
# **** start with three base classes ****
class Context:
- def __init__(self):
+ def __init__(self, arg, *tensors):
+ self.arg = arg
+ self.parents = tensors
self.saved_tensors = []
def save_for_backward(self, *x):
@@ -14,16 +16,36 @@ class Context:
class Tensor:
def __init__(self, data, _children=()):
self.data = data
- self.grad = np.zeros_like(data)
+ self.grad = None
# internal variables used for autograd graph construction
- self._prev = set(_children)
+ self._ctx = None
+
+ def __str__(self):
+ return "Tensor of shape %r with grad %r" % (self.data.shape, self.grad)
+
+ def backward(self, allow_fill=True):
+ #print("running backward on", self)
+ if self._ctx is None:
+ return
+
+ if self.grad is None and allow_fill:
+ # fill in the first grad with one
+ assert self.data.size == 1
+ self.grad = np.ones_like(self.data)
+
+ assert(self.grad is not None)
+
+ grads = self._ctx.arg.backward(self._ctx, self.grad)
+ for t,g in zip(self._ctx.parents, grads):
+ t.grad = g
+ t.backward(False)
class Function:
def apply(self, arg, *x):
- ctx = Context()
- x = [self]+list(x)
- ret = Tensor(arg.forward(ctx, *[t.data for t in x]))
+ ctx = Context(arg, self, *x)
+ ret = Tensor(arg.forward(ctx, self.data, *[t.data for t in x]))
+ ret._ctx = ctx
return ret
def register(name, fxn):
@@ -31,7 +53,6 @@ def register(name, fxn):
# **** implement a few functions ****
-"""
class ReLU(Function):
@staticmethod
def forward(ctx, input):
@@ -41,11 +62,10 @@ class ReLU(Function):
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
- grad_input = grad_output.clone()
+ grad_input = grad_output.copy()
grad_input[input < 0] = 0
- return grad_input
-setattr(Tensor, 'relu', partialmethod(run, ReLU))
-"""
+ return grad_input,
+register('relu', ReLU)
class Dot(Function):
@staticmethod
@@ -57,11 +77,10 @@ class Dot(Function):
def backward(ctx, grad_output):
input, weight = ctx.saved_tensors
grad_input = grad_output.dot(weight.T)
- grad_weight = grad_output.dot(input)
+ grad_weight = grad_output.T.dot(input).T
return grad_input, grad_weight
register('dot', Dot)
-# may be wrong
class Sum(Function):
@staticmethod
def forward(ctx, input):
@@ -71,6 +90,6 @@ class Sum(Function):
@staticmethod
def backward(ctx, grad_output):
input = ctx.saved_tensors
- return grad_output * input
+ return grad_output * np.ones_like(input)
register('sum', Sum)
