refactor tinygrad to be more tiny
MODIFY
README.md
README.md
@@ -38,6 +38,7 @@ print(y.grad) # dz/dy
### TODO (to make real neural network library)
+* Implement gradcheck (numeric)
* Implement convolutions
* Implement Adam optimizer

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
from functools import partialmethod
import numpy as np
-# **** start with three base classes ****
-
-class Context:
- def __init__(self, arg, *tensors):
- self.arg = arg
- self.parents = tensors
- self.saved_tensors = []
-
- def save_for_backward(self, *x):
- self.saved_tensors.extend(x)
+# **** start with two base classes ****
class Tensor:
def __init__(self, data):
@@ -35,17 +26,18 @@ class Tensor:
if self.grad is None and allow_fill:
# fill in the first grad with one
+ # this is "implicit gradient creation"
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)
+ grads = self._ctx.backward(self._ctx, self.grad)
if len(self._ctx.parents) == 1:
grads = [grads]
for t,g in zip(self._ctx.parents, grads):
if g.shape != t.data.shape:
- print("grad shape must match tensor shape in %r, %r != %r" % (self._ctx.arg, g.shape, t.data.shape))
+ print("grad shape must match tensor shape in %r, %r != %r" % (self._ctx, g.shape, t.data.shape))
assert(False)
t.grad = g
t.backward(False)
@@ -54,9 +46,18 @@ class Tensor:
div = Tensor(np.array([1/self.data.size]))
return self.sum().mul(div)
+# The Function is the Context
class Function:
+ def __init__(self, *tensors):
+ self.parents = tensors
+ self.saved_tensors = []
+
+ def save_for_backward(self, *x):
+ self.saved_tensors.extend(x)
+
+ # note that due to how partialmethod works, self and arg are switched
def apply(self, arg, *x):
- ctx = Context(arg, self, *x)
+ ctx = arg(self, *x)
ret = Tensor(arg.forward(ctx, self.data, *[t.data for t in x]))
ret._ctx = ctx
return ret

refactor tinygrad to be more tiny