not supporting stride
MODIFY
test/test_ops.py
test/test_ops.py
@@ -63,7 +63,7 @@ class TestOps(unittest.TestCase):
helper_test_op([(45,65)], lambda x: torch.nn.LogSoftmax(dim=1)(x), Tensor.logsoftmax, atol=1e-7, grad_atol=1e-7, gpu=self.gpu)
def test_pad2d(self):
- helper_test_op([(3,3,3,3)], lambda x: torch.nn.functional.pad(x, (1,1,1,1)), lambda x: x.pad2d(padding=(1,1,1,1)), gpu=self.gpu)
+ helper_test_op([(3,3,3,3)], lambda x: torch.nn.functional.pad(x, (1,1,1,1)), lambda x: x.pad2d(padding=(1,1,1,1)), gpu=self.gpu, forward_only=self.gpu)
def test_conv2d(self):
for bs in [1,8]:
@@ -98,11 +98,11 @@ class TestOps(unittest.TestCase):
lambda x: Tensor.max_pool2d(x, kernel_size=ksz), gpu=self.gpu, forward_only=self.gpu)
def test_avgpool2d(self):
- for ksz in [(2,2), ]:#, (3,3), (3,2), (5,5), (5,1)]:
+ for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
with self.subTest(kernel_size=ksz):
helper_test_op([(32,2,111,28)],
lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz),
- lambda x: Tensor.avg_pool2d(x, kernel_size=ksz), gpu=self.gpu)
+ lambda x: Tensor.avg_pool2d(x, kernel_size=ksz), gpu=self.gpu, forward_only=self.gpu)
if GPU:
class TestOpsGPU(TestOps):

MODIFY
tinygrad/opsgpu.py
tinygrad/opsgpu.py
@@ -382,9 +382,8 @@ register('sigmoid', Sigmoid, gpu=True)
class AvgPool2D(Function):
@staticmethod
- def forward(ctx, input, kernel_size=(2, 2), stride=None):
- if not stride:
- ctx.stride = stride = kernel_size
+ def forward(ctx, input, kernel_size=(2, 2)):
+ ctx.stride = stride = ctx.kernel_size
iter_op = "group_res += input[iid]"
result_op = "group_res / (kernel_size.x * kernel_size.y)"
ret = subsample_op(ctx, input, kernel_size, stride, iter_op, result_op)
@@ -393,6 +392,8 @@ class AvgPool2D(Function):
@staticmethod
def backward(ctx, grad_output):
+ raise NotImplementedError("GPU AvgPool2D.backward() is broken")
+
# TODO implement for stride != kernel_size
if ctx.kernel_size != ctx.stride:
raise NotImplementedError("GPU AvgPool2D.backward() with stride != kernel_size not implemented")
@@ -403,9 +404,8 @@ register('avg_pool2d', AvgPool2D, gpu=True)
class MaxPool2D(Function):
@staticmethod
- def forward(ctx, input, kernel_size=(2, 2), stride=None):
- if not stride:
- ctx.stride = stride = kernel_size
+ def forward(ctx, input, kernel_size=(2, 2)):
+ ctx.stride = stride = ctx.kernel_size
init_val = "FLT_MIN"
iter_op = "group_res = max(group_res, input[iid])"
result_op = "group_res"

not supporting stride