diff --git a/test/test_conv_speed.py b/test/test_net_speed.py similarity index 100% rename from test/test_conv_speed.py rename to test/test_net_speed.py diff --git a/test/test_ops.py b/test/test_ops.py index 4a20557277..04b262b469 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -5,7 +5,7 @@ import timeit import functools from tinygrad.tensor import Tensor -def test_op(shps, torch_fxn, tinygrad_fxn, atol=1e-7, grad_atol=1e-7): +def helper_test_op(shps, torch_fxn, tinygrad_fxn, atol=1e-7, grad_atol=1e-7): ts = [torch.rand(x, requires_grad=True) for x in shps] tst = [Tensor(x.detach().numpy()) for x in ts] @@ -36,15 +36,15 @@ class TestOps(unittest.TestCase): for cin in [1,3]: for H in [2,5]: for W in [2,3,5]: - test_op([(bs,cin,11,28), (4,cin,H,W)], + helper_test_op([(bs,cin,11,28), (4,cin,H,W)], lambda x,w: torch.nn.functional.conv2d(x,w).relu(), lambda x,w: Tensor.conv2d(x,w).relu(), atol=2e-5, grad_atol=2e-6) def test_maxpool2x2(self): - test_op([(32,2,110,28)], lambda x: torch.nn.functional.max_pool2d(x, (2,2)), Tensor.max_pool2d) + helper_test_op([(32,2,110,28)], lambda x: torch.nn.functional.max_pool2d(x, (2,2)), Tensor.max_pool2d) def test_avgpool2x2(self): - test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, (2,2)), Tensor.avg_pool2d) + helper_test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, (2,2)), Tensor.avg_pool2d) if __name__ == '__main__': unittest.main(verbosity=2)