make default parameters work for functions
MODIFY
test/test_mnist.py
test/test_mnist.py
@@ -100,4 +100,3 @@ class TestMNIST(unittest.TestCase):
if __name__ == '__main__':
unittest.main()
-

MODIFY
test/test_tensor.py
test/test_tensor.py
@@ -66,4 +66,3 @@ class TestTinygrad(unittest.TestCase):
if __name__ == '__main__':
unittest.main()
-

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -135,8 +135,8 @@ def stack_for_pool(x, py, px):
class MaxPool2D(Function):
@staticmethod
- def forward(ctx, x):
- stack = stack_for_pool(x, 2, 2)
+ def forward(ctx, x, pool_size=(2, 2)):
+ stack = stack_for_pool(x, *pool_size)
idxs = np.argmax(stack, axis=0)
ctx.save_for_backward(idxs, x.shape)
return np.max(stack, axis=0)
@@ -144,29 +144,30 @@ class MaxPool2D(Function):
@staticmethod
def backward(ctx, grad_output):
idxs,s = ctx.saved_tensors
- my, mx = (s[2]//2)*2, (s[3]//2)*2
+ py, px = ctx.pool_size
+ my, mx = (s[2]//py)*py, (s[3]//px)*px
ret = np.zeros(s, dtype=grad_output.dtype)
- for Y in range(2):
- for X in range(2):
- ret[:, :, Y:my:2, X:mx:2] = grad_output * (idxs == (Y*2+X))
+ for Y in range(py):
+ for X in range(px):
+ ret[:, :, Y:my:py, X:mx:px] = grad_output * (idxs == (Y*px+X))
return ret
register('max_pool2d', MaxPool2D)
class AvgPool2D(Function):
@staticmethod
- def forward(ctx, x):
- stack = stack_for_pool(x, 2, 2)
+ def forward(ctx, x, pool_size=(2, 2)):
+ stack = stack_for_pool(x, *pool_size)
ctx.save_for_backward(x.shape)
return np.mean(stack, axis=0)
@staticmethod
def backward(ctx, grad_output):
s, = ctx.saved_tensors
- my, mx = (s[2]//2)*2, (s[3]//2)*2
+ py, px = ctx.pool_size
+ my, mx = (s[2]//py)*py, (s[3]//px)*px
ret = np.zeros(s, dtype=grad_output.dtype)
- for Y in range(2):
- for X in range(2):
- ret[:, :, Y:my:2, X:mx:2] = grad_output/4
+ for Y in range(py):
+ for X in range(px):
+ ret[:, :, Y:my:py, X:mx:px] = grad_output/py/px
return ret
register('avg_pool2d', AvgPool2D)
-

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
# inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
from functools import partialmethod
+from inspect import signature
import numpy as np
# **** start with two base classes ****
@@ -74,7 +75,7 @@ class Function:
self.saved_tensors.extend(x)
# note that due to how partialmethod works, self and arg are switched
- def apply(self, arg, *x):
+ def apply(self, arg, *x, **kwargs):
# support the args in both orders
if type(arg) == Tensor:
op = self
@@ -83,7 +84,15 @@ class Function:
op = arg
x = [self]+list(x)
ctx = op(*x)
- ret = Tensor(op.forward(ctx, *[t.data for t in x]))
+ # use default params
+ params = signature(op.forward).parameters
+ for p in params.values():
+ if p.default is not p.empty:
+ setattr(ctx, p.name, p.default)
+ # overwrite with passed params
+ for k, v in kwargs.items():
+ setattr(ctx, k, v)
+ ret = Tensor(op.forward(ctx, *[t.data for t in x], **kwargs))
ret._ctx = ctx
return ret

make default parameters work for functions