diff --git a/test/test_ops.py b/test/test_ops.py index 87ed285dda..d9c6af71ee 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1,3 +1,4 @@ +import os import torch import numpy as np import unittest @@ -41,8 +42,6 @@ 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): @@ -64,6 +63,16 @@ class TestOps(unittest.TestCase): def test_logsoftmax(self): 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_broadcast(self): + if os.getenv('CI') and self.gpu: + raise unittest.SkipTest('GPU broadcasting not fully supported') + for torch_op, tinygrad_op in [(torch.add, Tensor.add), (torch.sub, Tensor.sub), (torch.mul, Tensor.mul), + (torch.div, Tensor.div), (torch.pow, Tensor.pow)]: + for shapes in [((1,32,32,32), (1,32,1,1)), ((5,13,24,16), (5,1,24,1)), + ((4,1), (4,5)), ((1,4), (5,4)), ((1,3,1,7,1), (2,1,5,1,8))]: + with self.subTest(op=torch_op.__name__, shapes=shapes): + helper_test_op(shapes, torch_op, tinygrad_op, gpu=self.gpu, forward_only=True) + def test_pad2d(self): helper_test_op([(3,3,3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4)), lambda x: x.pad2d(padding=(1,2,3,4)), gpu=self.gpu)