some broadcasting, pool test is fail
MODIFY
test/test_ops.py
test/test_ops.py
@@ -39,6 +39,8 @@ class TestOps(unittest.TestCase):
gpu = False
def test_add(self):
helper_test_op([(45,65), (45,65)], lambda x,y: x+y, Tensor.add, gpu=self.gpu)
+ def test_broadcast_add(self):
+ helper_test_op([(1,32,32,32), (1,32,1,1)], lambda x,y: x+y, Tensor.add, gpu=self.gpu, forward_only=True)
def test_sub(self):
helper_test_op([(45,65), (45,65)], lambda x,y: x-y, Tensor.sub, gpu=self.gpu)
def test_mul(self):
@@ -86,16 +88,16 @@ class TestOps(unittest.TestCase):
lambda x,w: Tensor.conv2d(x,w,stride=(2,1)).relu(), atol=2e-5, grad_atol=2e-6, gpu=self.gpu, forward_only=self.gpu)
def test_maxpool2x2(self):
- helper_test_op([(32,2,110,28)], lambda x: torch.nn.functional.max_pool2d(x, (2,2)), Tensor.max_pool2d, gpu=self.gpu)
+ helper_test_op([(32,2,110,28)], lambda x: torch.nn.functional.max_pool2d(x, (2,2)), Tensor.max_pool2d, gpu=self.gpu, forward_only=self.gpu)
def test_maxpool_sizes(self):
for sz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
helper_test_op([(32,2,110,28)],
lambda x: torch.nn.functional.max_pool2d(x, kernel_size=sz),
- lambda x: Tensor.max_pool2d(x, kernel_size=sz), gpu=self.gpu)
+ lambda x: Tensor.max_pool2d(x, kernel_size=sz), gpu=self.gpu, forward_only=self.gpu)
def test_avgpool2x2(self):
- helper_test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, (2,2)), Tensor.avg_pool2d, gpu=self.gpu)
+ helper_test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, (2,2)), Tensor.avg_pool2d, gpu=self.gpu, forward_only=self.gpu)
if GPU:
class TestOpsGPU(TestOps):

MODIFY
tinygrad/nn.py
tinygrad/nn.py
@@ -18,7 +18,7 @@ class BatchNorm2D:
# this work at inference?
x = x.sub(self.running_mean.reshape(shape=[1, -1, 1, 1]))
x = x.mul(self.weight.reshape(shape=[1, -1, 1, 1]))
- x = x.div(self.running_var.add(Tensor([self.eps])).reshape(shape=[1, -1, 1, 1]).sqrt())
+ x = x.div(self.running_var.add(Tensor([self.eps], gpu=x.gpu)).reshape(shape=[1, -1, 1, 1]).sqrt())
x = x.add(self.bias.reshape(shape=[1, -1, 1, 1]))
return x

MODIFY
tinygrad/opsgpu.py
tinygrad/opsgpu.py
@@ -58,20 +58,32 @@ def subsample_op(ctx, input, kernel_size, iter_op, result_op, init_val=0):
return ret
def binary_op(ctx, code, x, y):
+ assert len(x.shape) == len(y.shape)
+ xdiv = 1
+ ydiv = 1
if x.shape != y.shape:
+ # special case broadcasting
+ # TODO: make general
+ if len(y.shape) == 4 and x.shape[0:2] == y.shape[0:2] and y.shape[2] == 1 and y.shape[3] == 1:
+ ydiv = x.shape[2] * x.shape[3]
+ elif len(y.shape) == 4 and x.shape[0:2] == y.shape[0:2] and x.shape[2] == 1 and x.shape[3] == 1:
+ xdiv = y.shape[2] * y.shape[3]
+ elif np.prod(y.shape) == 1:
+ ydiv = np.prod(x.shape)
+ else:
raise Exception("binary op shape mismatch: %r != %r" % (x.shape, y.shape))
ret = buffer_like(ctx, x)
prg = clbuild(ctx.cl_ctx, """
__kernel void binop(
- __global const float *a_g, __global const float *b_g, __global float *res_g)
+ __global const float *a_g, __global const float *b_g, __global float *res_g, int xdiv, int ydiv)
{
int gid = get_global_id(0);
- float a = a_g[gid];
- float b = b_g[gid];
+ float a = a_g[gid/xdiv];
+ float b = b_g[gid/ydiv];
res_g[gid] = """+code+""";
}
""")
- prg.binop(ctx.cl_queue, [np.prod(ret.shape)], None, x, y, ret)
+ prg.binop(ctx.cl_queue, [np.prod(ret.shape)], None, x, y, ret, np.int32(xdiv), np.int32(ydiv))
return ret
def unary_op(ctx, code, x):

some broadcasting, pool test is fail