From ae112c9dbe315f2f4ea2526f206608b5f2631179 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 3 Jan 2024 23:53:33 -0500 Subject: [PATCH] fix some long lines in tests (#3006) * fix some long lines in tests * better --- sz.py | 6 ++- test/models/test_onnx.py | 12 ++++-- test/test_linearizer.py | 35 ++++++++++------ test/test_ops.py | 89 +++++++++++++++++++++++++--------------- test/test_randomness.py | 16 +++++--- test/test_uops.py | 1 - test/test_zero_copy.py | 1 - tinygrad/device.py | 2 +- 8 files changed, 104 insertions(+), 58 deletions(-) diff --git a/sz.py b/sz.py index bc521587f5..8d465b7fa4 100755 --- a/sz.py +++ b/sz.py @@ -38,7 +38,8 @@ def gen_diff(table_old, table_new): file_stat_old = [stats for stats in table_old if file in stats] file_stat_new = [stats for stats in table_new if file in stats] if file_stat_new[0][1]-file_stat_old[0][1] != 0 or file_stat_new[0][2]-file_stat_old[0][2] != 0: - table.append([file_stat_new[0][0], file_stat_new[0][1], file_stat_new[0][1]-file_stat_old[0][1], file_stat_new[0][2], file_stat_new[0][2]-file_stat_old[0][2]]) # noqa: E501 + table.append([file_stat_new[0][0], file_stat_new[0][1], file_stat_new[0][1]-file_stat_old[0][1], file_stat_new[0][2], + file_stat_new[0][2]-file_stat_old[0][2]]) return table def display_diff(diff): return "+"+str(diff) if diff > 0 else str(diff) @@ -58,7 +59,8 @@ if __name__ == "__main__": if len(sys.argv) == 3: print("### Changes") print("```") - print(tabulate([headers] + sorted(table, key=lambda x: -x[1]), headers="firstrow", intfmt=(..., "d", "+d"), floatfmt=(..., ..., ..., ".1f", "+.1f"))+"\n") # noqa: E501 + print(tabulate([headers] + sorted(table, key=lambda x: -x[1]), headers="firstrow", intfmt=(..., "d", "+d"), + floatfmt=(..., ..., ..., ".1f", "+.1f"))+"\n") print(f"\ntotal lines changes: {display_diff(sum([x[2] for x in table]))}") print("```") else: diff --git a/test/models/test_onnx.py b/test/models/test_onnx.py index db7b708823..b4d691f6af 100644 --- a/test/models/test_onnx.py +++ b/test/models/test_onnx.py @@ -99,17 +99,23 @@ class TestOnnxModel(unittest.TestCase): @unittest.skip("slow") def test_efficientnet(self): input_name, input_new = "images:0", True - self._test_model(fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx"), input_name, input_new) # noqa: E501 + self._test_model( + fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx"), + input_name, input_new) def test_shufflenet(self): input_name, input_new = "gpu_0/data_0", False - self._test_model(fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/shufflenet/model/shufflenet-9.onnx"), input_name, input_new) # noqa: E501 + self._test_model( + fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/shufflenet/model/shufflenet-9.onnx"), + input_name, input_new) @unittest.skip("test is very slow") def test_resnet(self): # NOTE: many onnx models can't be run right now due to max pool with strides != kernel_size input_name, input_new = "data", False - self._test_model(fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet18-v2-7.onnx"), input_name, input_new) # noqa: E501 + self._test_model( + fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet18-v2-7.onnx"), + input_name, input_new) def _test_model(self, fn, input_name, input_new, debug=False): onnx_model = onnx.load(fn) diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 3e3b5142c8..6f1031a99d 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -1,4 +1,3 @@ -# ruff: noqa: E501 import numpy as np import unittest, os @@ -110,8 +109,10 @@ class TestLinearizer(unittest.TestCase): def test_simplify_uop(self): def helper_test_simplify(uop, dtype, vin, arg=None): - ast = LazyOp(op=BufferOps.CONST, src=(), arg=ConstBuffer(val=42, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)))) - ast = LazyOp(BufferOps.STORE, (ast,), MemBuffer(0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)))) + ast = LazyOp(BufferOps.CONST, (), + ConstBuffer(42, dtypes.float, ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)))) + ast = LazyOp(BufferOps.STORE, (ast,), + MemBuffer(0, dtypes.float, ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)))) lin = Linearizer(ast=ast) # this is a dummy ast lin.uops = [] @@ -122,13 +123,15 @@ class TestLinearizer(unittest.TestCase): c0 = UOp(UOps.CONST, dtypes.float, vin=(), arg=0.0) c1 = UOp(UOps.CONST, dtypes.float, vin=(), arg=1.0) - assert helper_test_simplify(UOps.ALU, dtypes.float, vin=(UOp(UOps.CONST, dtypes.bool, vin=(), arg=True), c0, c1), arg=TernaryOps.WHERE).uop == UOps.ALU + assert helper_test_simplify(UOps.ALU, dtypes.float, vin=(UOp(UOps.CONST, dtypes.bool, vin=(), arg=True), c0, c1), + arg=TernaryOps.WHERE).uop == UOps.ALU def helper_realized_ast(r:Tensor): s = r.lazydata.schedule() run_schedule(s[:-1]) # run all kernels except the last one # now all input LazyBuffers buffers in s[-1] should be realized - output_buffer = Buffer(s[-1].out.device, prod((s if isinstance(s, int) else s.max for s in s[-1].out.shape)), s[-1].out.dtype) # allocate an output buffer + # allocate an output buffer + output_buffer = Buffer(s[-1].out.device, prod((s if isinstance(s, int) else s.max for s in s[-1].out.shape)), s[-1].out.dtype) return s[-1].ast, [output_buffer] + [l.realized for l in s[-1].inputs] @unittest.skipUnless(isinstance(Device[Device.DEFAULT], Compiled) and Device[Device.DEFAULT].linearizer_opts.supports_float4, @@ -409,8 +412,10 @@ class TestLinearizerOpts(unittest.TestCase): [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 2)], [Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.GROUPTOP, 0, 16)], [Opt(OptOps.LOCAL, 0, 32), Opt(OptOps.GROUPTOP, 0, 2)], - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 64)], # Checking how it works with locals + grouped reduce - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UNROLL, 1, 4)], # Checking how it works with locals + grouped reduce + upcasts + # Checking how it works with locals + grouped reduce + [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 64)], + # Checking how it works with locals + grouped reduce + upcasts + [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UNROLL, 1, 4)], ]) def test_upcasts(self): @@ -457,8 +462,11 @@ class TestLinearizerOpts(unittest.TestCase): [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 0, 32)], [Opt(OptOps.LOCAL, 0, 8), Opt(OptOps.GROUPTOP, 0, 32)], [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 8), Opt(OptOps.GROUPTOP, 0, 4)], # Checking how it works with local+grouped_reduce - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)], # Checking all together - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 8)], # Full global upcast + local + # Checking all together + [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), + Opt(OptOps.UPCAST, 1, 2)], + # Full global upcast + local + [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 8)], ]) def test_double_reduce(self): @@ -479,10 +487,13 @@ class TestLinearizerOpts(unittest.TestCase): [Opt(OptOps.GROUPTOP, 0, 16), Opt(OptOps.GROUPTOP, 1, 2), Opt(OptOps.UNROLL, 0, 4)], [Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.GROUPTOP, 1, 32), Opt(OptOps.UNROLL, 2, 4)], # Checking how it works with 2 grouped_reduces + upcasts. [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 0, 4), Opt(OptOps.GROUPTOP, 1, 4)], - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.GROUPTOP, 1, 32), Opt(OptOps.UNROLL, 1, 4)], # Checking how it works with 2 grouped_reduces + upcasts + locals. + # Checking how it works with 2 grouped_reduces + upcasts + locals. + [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 0, 2), Opt(OptOps.GROUPTOP, 1, 32), Opt(OptOps.UNROLL, 1, 4)], [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.UPCAST, 0, 2)], - [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)], # Checking how it works with 2 grouped_reduces + upcasts + locals. - [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 0, 4), Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 0, 2)], # No globals + [Opt(OptOps.LOCAL, 0, 2), Opt(OptOps.LOCAL, 1, 2), Opt(OptOps.GROUPTOP, 0, 8), Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.UPCAST, 0, 2), + Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)], # Checking how it works with 2 grouped_reduces + upcasts + locals. + [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4), Opt(OptOps.GROUPTOP, 0, 4), Opt(OptOps.GROUPTOP, 1, 4), Opt(OptOps.UPCAST, 0, 2), + Opt(OptOps.UPCAST, 0, 2)], # No globals ]) def test_tensor_core_opts(self): diff --git a/test/test_ops.py b/test/test_ops.py index c2bd63d8ef..30e0935bed 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1,9 +1,6 @@ -# ruff: noqa: E501 import torch -import time -import math +import time, math, unittest import numpy as np -import unittest from tinygrad.tensor import Tensor from tinygrad.helpers import getenv, IMAGE, DEBUG, CI from tinygrad import Device, dtypes @@ -14,7 +11,8 @@ if CI: FORWARD_ONLY = getenv("FORWARD_ONLY", 0) PRINT_TENSORS = getenv("PRINT_TENSORS", 0) -def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, grad_atol=1e-4, grad_rtol=1e-3, forward_only=False, vals=None, a=-0.5, b=3): +def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, grad_atol=1e-4, grad_rtol=1e-3, + forward_only=False, vals=None, a=-0.5, b=3): if tinygrad_fxn is None: tinygrad_fxn = torch_fxn ts, tst = prepare_test_op(a, b, shps, vals, forward_only) @@ -54,7 +52,9 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, gra for i, (t, tt) in enumerate(zip(ts, tst)): compare(f"backward pass tensor {i}", tt.grad.numpy(), t.grad.detach().numpy(), atol=grad_atol, rtol=grad_rtol) - if not CI: print("\ntesting %40r torch/tinygrad fp: %.2f / %.2f ms bp: %.2f / %.2f ms " % (shps, torch_fp*1000, tinygrad_fp*1000, torch_fbp*1000, tinygrad_fbp*1000), end="") + if not CI: + print("\ntesting %40r torch/tinygrad fp: %.2f / %.2f ms bp: %.2f / %.2f ms " % \ + (shps, torch_fp*1000, tinygrad_fp*1000, torch_fbp*1000, tinygrad_fbp*1000), end="") def prepare_test_op(a, b, shps, vals, forward_only=False): torch.manual_seed(0) @@ -165,11 +165,13 @@ class TestOps(unittest.TestCase): helper_test_op([], lambda: (-torch.ones(3,3)).sum(axis=1), lambda: (-Tensor.ones(3,3)).sum(axis=1), forward_only=True) def test_sum_pad_collapse(self): - helper_test_op([], lambda: torch.nn.functional.pad(torch.ones(256,256), pad=(0,64,0,0)).sum(axis=1), lambda: Tensor.ones(256,256).pad(((0,0), (0,64))).sum(axis=1), forward_only=True) + helper_test_op([], lambda: torch.nn.functional.pad(torch.ones(256,256), pad=(0,64,0,0)).sum(axis=1), + lambda: Tensor.ones(256,256).pad(((0,0), (0,64))).sum(axis=1), forward_only=True) # this is more complex and won't fold for a while def test_sum_cat_collapse(self): - helper_test_op([], lambda: torch.cat([torch.ones(256,256), torch.zeros(256,64)], dim=1).sum(axis=1), lambda: Tensor.cat(Tensor.ones(256,256), Tensor.zeros(256,64), dim=1).sum(axis=1), forward_only=True) + helper_test_op([], lambda: torch.cat([torch.ones(256,256), torch.zeros(256,64)], dim=1).sum(axis=1), + lambda: Tensor.cat(Tensor.ones(256,256), Tensor.zeros(256,64), dim=1).sum(axis=1), forward_only=True) def test_max_dont_collapse(self): helper_test_op([], lambda: torch.ones(256,256).max(1)[0], lambda: Tensor.ones(256,256).max(1), forward_only=True) @@ -449,7 +451,8 @@ class TestOps(unittest.TestCase): def test_multinomial(self): # NOTE: this is random, so it has a very large atol - helper_test_op([(1000,)], lambda x: torch.multinomial(x.clip(0,1), num_samples=1), lambda x: Tensor.multinomial(x.clip(0,1)), forward_only=True, atol=1000.) + helper_test_op([(1000,)], lambda x: torch.multinomial(x.clip(0,1), num_samples=1), + lambda x: Tensor.multinomial(x.clip(0,1)), forward_only=True, atol=1000.) def test_small_cumsum(self): helper_test_op([(10)], lambda x: torch.cumsum(x, dim=0), lambda x: Tensor.cumsum(x, axis=0), atol=1e-6) @@ -508,9 +511,11 @@ class TestOps(unittest.TestCase): # batch matrix multiplication, result & input permuted helper_test_op([(20,10,25),(10,25,32)], lambda a,b: torch.einsum('jik,ikl->jil', [a, b]), lambda a,b: Tensor.einsum('jik,ikl->jil', [a, b])) # tensor contraction - helper_test_op([(3,5,8,10),(11,13,5,16,8)], lambda a,b: torch.einsum('pqrs,tuqvr->pstuv', a,b), lambda a,b: Tensor.einsum('pqrs,tuqvr->pstuv', a,b), atol=1e-5) + helper_test_op([(3,5,8,10),(11,13,5,16,8)], lambda a,b: torch.einsum('pqrs,tuqvr->pstuv', a,b), + lambda a,b: Tensor.einsum('pqrs,tuqvr->pstuv', a,b), atol=1e-5) # tensor contraction, input permuted - helper_test_op([(3,8,10,5),(11,5,13,16,8)], lambda a,b: torch.einsum('prsq,tquvr->pstuv', a,b), lambda a,b: Tensor.einsum('prsq,tquvr->pstuv', a,b), atol=1e-5) + helper_test_op([(3,8,10,5),(11,5,13,16,8)], lambda a,b: torch.einsum('prsq,tquvr->pstuv', a,b), + lambda a,b: Tensor.einsum('prsq,tquvr->pstuv', a,b), atol=1e-5) # bilinear transformation helper_test_op([(2,3),(5,3,7),(2,7)], lambda a,b,c: torch.einsum('ik,jkl,il->ij', [a,b,c]), lambda a,b,c: Tensor.einsum('ik,jkl,il->ij', [a,b,c])) @@ -611,7 +616,8 @@ class TestOps(unittest.TestCase): helper_test_op([(45, 65, 85)], lambda x: torch.std(x, correction=0, dim=None), lambda x: Tensor.std(x, axis=None, correction=0)) def test_std_keepdim(self): helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=None, keepdim=True), lambda x: Tensor.std(x, keepdim=True)) - helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=0, keepdim=True, correction=0), lambda x: Tensor.std(x, keepdim=True, correction=0, axis=0)) + helper_test_op([(45, 65, 85)], lambda x: torch.std(x, dim=0, keepdim=True, correction=0), + lambda x: Tensor.std(x, keepdim=True, correction=0, axis=0)) def test_log_softmax(self): helper_test_op([(45,65)], lambda x: torch.nn.LogSoftmax(dim=1)(x), Tensor.log_softmax, atol=1e-7, grad_atol=1e-7) helper_test_op([()], lambda x: torch.nn.LogSoftmax(dim=0)(x), Tensor.log_softmax, atol=1e-7, grad_atol=1e-7) @@ -766,8 +772,8 @@ class TestOps(unittest.TestCase): def test_pad(self): helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4)),lambda x: x.pad(((3,4),(1,2)))) helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4), value=5), lambda x: x.pad(((3,4), (1,2)), value=5)) - helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4), value=float("inf")), lambda x: x.pad(((3,4), (1,2)), value=float("inf"))) - helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4), value=float("-inf")), lambda x: x.pad(((3,4), (1,2)), value=float("-inf"))) + helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4), value=math.inf), lambda x: x.pad(((3,4), (1,2)), value=math.inf)) + helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4), value=-math.inf), lambda x: x.pad(((3,4), (1,2)), value=-math.inf)) helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (0,0,3,4), value=1), lambda x: x.pad(((3,4), None), value=1)) helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x, (0,0,0,0), value=1), lambda x: x.pad((None, None), value=1)) @@ -790,8 +796,10 @@ class TestOps(unittest.TestCase): helper_test_op([(4,4)], lambda x: torch.nn.functional.pad(x,(3,4,1,2), value=value)[:,4], lambda x: x.pad(((1,2),(3,4)), value=value)[:,4]) helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x,(0,3,0,0), value=value)[:,4:6], lambda x: x.pad(((0,0),(0,3)), value=value)[:,4:6]) helper_test_op([(3,3)], lambda x: torch.nn.functional.pad(x,(0,1,3,2), value=value)[0:2,:], lambda x: x.pad(((3,2),(0,1)), value=value)[0:2,:]) - helper_test_op([(3,3,3)], lambda x: torch.nn.functional.pad(x,(1,1,0,1,3,2), value=value)[0:2,:,:], lambda x: x.pad(((3,2),(0,1),(1,1)), value=value)[0:2,:,:]) - helper_test_op([(3,3,3)], lambda x: torch.nn.functional.pad(x,(1,1,0,1,3,2), value=value)[2:4,:,:], lambda x: x.pad(((3,2),(0,1),(1,1)), value=value)[2:4,:,:]) + helper_test_op([(3,3,3)], lambda x: torch.nn.functional.pad(x,(1,1,0,1,3,2), value=value)[0:2,:,:], + lambda x: x.pad(((3,2),(0,1),(1,1)), value=value)[0:2,:,:]) + helper_test_op([(3,3,3)], lambda x: torch.nn.functional.pad(x,(1,1,0,1,3,2), value=value)[2:4,:,:], + lambda x: x.pad(((3,2),(0,1),(1,1)), value=value)[2:4,:,:]) def test_stack_slice(self): helper_test_op([(4)], lambda x: torch.stack([x for i in range(3)])[0,:], lambda x: Tensor.stack([x for i in range(3)])[0,:]) @@ -1255,7 +1263,7 @@ class TestOps(unittest.TestCase): x = Tensor.randn(45, 65, 3) for dim in range(-1, 3): - helper_test_op([(45, 65, 3), (45, 65, 3), (45, 65, 3)], lambda x, y, z: torch.stack((x, y, z), dim=dim), lambda x, y, z: Tensor.stack([x, y, z], dim=dim)) + helper_test_op([(45,65,3), (45,65,3), (45,65,3)], lambda x, y, z: torch.stack((x, y, z), dim), lambda x, y, z: Tensor.stack([x, y, z], dim)) with self.assertRaises(IndexError): Tensor.stack([x], dim=77) @@ -1353,10 +1361,14 @@ class TestOps(unittest.TestCase): def test_slice_fancy_indexing_with_idx(self): # indexing using idx with different dim - helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,0,0],[0,0,0]]), torch.tensor(1)], lambda x: x[Tensor([[0,0,0],[0,0,0]]), Tensor(1)]) - helper_test_op([(2,3)], lambda x: x[torch.tensor([1]), torch.tensor([[0,0,0],[0,0,0]])], lambda x: x[Tensor([1]), Tensor([[0,0,0],[0,0,0]])]) - helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,0,0],[0,0,0]]), torch.tensor([2,1,1])], lambda x: x[Tensor([[0,0,0],[0,0,0]]), Tensor([2,1,1])]) - helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,1,-1],[-1,-2,0]]), torch.tensor([2,1,-1])], lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])]) + helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,0,0],[0,0,0]]), torch.tensor(1)], + lambda x: x[Tensor([[0,0,0],[0,0,0]]), Tensor(1)]) + helper_test_op([(2,3)], lambda x: x[torch.tensor([1]), torch.tensor([[0,0,0],[0,0,0]])], + lambda x: x[Tensor([1]), Tensor([[0,0,0],[0,0,0]])]) + helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,0,0],[0,0,0]]), torch.tensor([2,1,1])], + lambda x: x[Tensor([[0,0,0],[0,0,0]]), Tensor([2,1,1])]) + helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,1,-1],[-1,-2,0]]), torch.tensor([2,1,-1])], + lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])]) def test_slice_fancy_indexing_list_indices(self): a,b,c,d,e,i,j,k,o,p = self._get_index_randoms() @@ -1389,8 +1401,10 @@ class TestOps(unittest.TestCase): def test_slice_fancy_indexing_tuple_with_tensors(self): a,b,c,d,e,i,j,k,o,p = self._get_index_randoms() - # helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,),], lambda x: x[(i,),]) TypeError: only integer tensors of a single element can be converted to an index - # helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,1),], lambda x: x[(i,1),]) TypeError: only integer tensors of a single element can be converted to an index + # # TypeError: only integer tensors of a single element can be converted to an index + # helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,),], lambda x: x[(i,),]) + # # TypeError: only integer tensors of a single element can be converted to an index + # helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,1),], lambda x: x[(i,1),]) helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,[1,1])], lambda x: x[(i,[1,1])]) helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,(1,1))], lambda x: x[(i,(1,1))]) helper_test_op([(2,5,6,5,3,4)], lambda x: x[(a,b,c,d,e)], lambda x: x[(i,j,k,o,p)]) @@ -1419,21 +1433,32 @@ class TestOps(unittest.TestCase): helper_test_op([(4,5,6)], lambda x: x.gather(index=b, dim=1), lambda x: x.gather(idx=a, dim=1)) helper_test_op([(4,5,6)], lambda x: x.gather(index=b, dim=2), lambda x: x.gather(idx=a, dim=2)) helper_test_op([(3,4,5)], lambda x: x.gather(index=b, dim=0), lambda x: x.gather(idx=a, dim=0)) - self.helper_test_exception([(4,5,6)], lambda x: x.gather(index=torch.tensor([1], dtype=torch.int64), dim=0), lambda x: x.gather(idx=Tensor([1], dtype=dtypes.int32), dim=0), expected=(RuntimeError, AssertionError)) - self.helper_test_exception([(2,1,1)], lambda x: x.gather(index=b, dim=0), lambda x: x.gather(idx=a, dim=0), expected=(RuntimeError, AssertionError)) + self.helper_test_exception([(4,5,6)], lambda x: x.gather(index=torch.tensor([1], dtype=torch.int64), dim=0), + lambda x: x.gather(idx=Tensor([1], dtype=dtypes.int32), dim=0), expected=(RuntimeError, AssertionError)) + self.helper_test_exception([(2,1,1)], lambda x: x.gather(index=b, dim=0), + lambda x: x.gather(idx=a, dim=0), expected=(RuntimeError, AssertionError)) def test_scaled_product_attention(self): - helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64)], lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z), lambda x,y,z: Tensor.scaled_dot_product_attention(x,y,z)) - helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64), (32,8,16,16)], lambda x,y,z,m: torch.nn.functional.scaled_dot_product_attention(x,y,z,attn_mask=m), lambda x,y,z,m: Tensor.scaled_dot_product_attention(x,y,z,attn_mask=m)) + helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64)], lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z), + lambda x,y,z: Tensor.scaled_dot_product_attention(x,y,z)) + helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64), (32,8,16,16)], + lambda x,y,z,m: torch.nn.functional.scaled_dot_product_attention(x,y,z,attn_mask=m), + lambda x,y,z,m: Tensor.scaled_dot_product_attention(x,y,z,attn_mask=m)) def test_scaled_product_attention_causal(self): - helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64)], lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z,is_causal=True), lambda x,y,z: Tensor.scaled_dot_product_attention(x,y,z,is_causal=True)) + helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64)], + lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z,is_causal=True), + lambda x,y,z: Tensor.scaled_dot_product_attention(x,y,z,is_causal=True)) def test_binary_crossentropy(self): - helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),torch.clip(y,0,1)), lambda x,y: x.sigmoid().binary_crossentropy(y.clip(0,1))) - helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,torch.clip(y,0,1)), lambda x,y: x.binary_crossentropy_logits(y.clip(0,1))) - helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,torch.clip(y,0,1)), lambda x,y: x.sigmoid().binary_crossentropy(y.clip(0,1))) - helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),torch.clip(y,0,1)), lambda x,y: x.binary_crossentropy_logits(y.clip(0,1))) + helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),torch.clip(y,0,1)), + lambda x,y: x.sigmoid().binary_crossentropy(y.clip(0,1))) + helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,torch.clip(y,0,1)), + lambda x,y: x.binary_crossentropy_logits(y.clip(0,1))) + helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy_with_logits(x,torch.clip(y,0,1)), + lambda x,y: x.sigmoid().binary_crossentropy(y.clip(0,1))) + helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),torch.clip(y,0,1)), + lambda x,y: x.binary_crossentropy_logits(y.clip(0,1))) if __name__ == '__main__': np.random.seed(1337) diff --git a/test/test_randomness.py b/test/test_randomness.py index c73ee4d960..1646eeb04c 100644 --- a/test/test_randomness.py +++ b/test/test_randomness.py @@ -1,4 +1,3 @@ -# ruff: noqa: E501 import math import unittest import numpy as np @@ -49,7 +48,8 @@ def equal_distribution(tiny_func, torch_func=None, numpy_func=None, shape=(20, 2 x2 = tiny_func(shape).numpy().flatten() if numpy_func is not None: y = numpy_func(shape).flatten() if torch_func is not None: z = torch_func(shape).numpy().flatten() - return (numpy_func is None or (kstest(x1, y) >= alpha and kstest(x2, y) >= alpha)) and (torch_func is None or (kstest(x1, z) >= alpha and kstest(x2, z) >= alpha)) + return (numpy_func is None or (kstest(x1, y) >= alpha and kstest(x2, y) >= alpha)) and \ + (torch_func is None or (kstest(x1, z) >= alpha and kstest(x2, z) >= alpha)) def normal_test(func, shape=(20, 23), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha) @@ -68,20 +68,24 @@ class TestRandomness(unittest.TestCase): def test_normal(self): self.assertTrue(normal_test(Tensor.normal)) - self.assertTrue(equal_distribution(Tensor.normal, lambda x: torch.nn.init.normal_(torch.empty(x), mean=0, std=1), lambda x: np.random.normal(loc=0, scale=1, size=x))) + self.assertTrue(equal_distribution(Tensor.normal, lambda x: torch.nn.init.normal_(torch.empty(x), mean=0, std=1), + lambda x: np.random.normal(loc=0, scale=1, size=x))) def test_uniform(self): self.assertFalse(normal_test(Tensor.uniform)) self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x)), lambda x: np.random.uniform(size=x))) - self.assertTrue(equal_distribution(partial(Tensor.uniform, low=-100, high=100, dtype=dtypes.int32), numpy_func=lambda x: np.random.randint(low=-100, high=100, size=x))) + self.assertTrue(equal_distribution(partial(Tensor.uniform, low=-100, high=100, dtype=dtypes.int32), + numpy_func=lambda x: np.random.randint(low=-100, high=100, size=x))) def test_scaled_uniform(self): self.assertFalse(normal_test(Tensor.scaled_uniform)) - self.assertTrue(equal_distribution(Tensor.scaled_uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1) / math.sqrt(math.prod(x)), lambda x: np.random.uniform(-1, 1, size=x) / math.sqrt(math.prod(x)))) + self.assertTrue(equal_distribution(Tensor.scaled_uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1) / math.sqrt(math.prod(x)), + lambda x: np.random.uniform(-1, 1, size=x) / math.sqrt(math.prod(x)))) def test_glorot_uniform(self): self.assertFalse(normal_test(Tensor.glorot_uniform)) - self.assertTrue(equal_distribution(Tensor.glorot_uniform, lambda x: torch.nn.init.xavier_uniform_(torch.empty(x)), lambda x: np.random.uniform(-1, 1, size=x) * math.sqrt(6 / (x[0] + math.prod(x[1:]))))) + self.assertTrue(equal_distribution(Tensor.glorot_uniform, lambda x: torch.nn.init.xavier_uniform_(torch.empty(x)), + lambda x: np.random.uniform(-1, 1, size=x) * math.sqrt(6 / (x[0] + math.prod(x[1:]))))) def test_kaiming_uniform(self): Tensor.manual_seed(1337) diff --git a/test/test_uops.py b/test/test_uops.py index 33e111509a..69a24d486a 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -1,4 +1,3 @@ -# ruff: noqa: E501 from typing import Optional, Tuple, Any, List import unittest, math import numpy as np diff --git a/test/test_zero_copy.py b/test/test_zero_copy.py index a462073a72..da8ca955aa 100644 --- a/test/test_zero_copy.py +++ b/test/test_zero_copy.py @@ -1,6 +1,5 @@ import unittest from tinygrad import Tensor, Device -#from tinygrad.helpers import CI import time def time_tensor_numpy(out:Tensor): diff --git a/tinygrad/device.py b/tinygrad/device.py index 74d15b185b..10957d6873 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -218,7 +218,7 @@ def _get_interpreted_fxn(fxn_for_op:Dict[Op, Callable], ast:LazyOp) -> Interpret if ast.src[0].op == UnaryOps.CAST and ast.src[0].src[0].op == BinaryOps.MUL: ast = LazyOp(TernaryOps.MULACC, ast.src[0].src[0].src, ast.arg) if ast.op in BufferOps: - if ast.op == ast.op == BufferOps.CONST: tmp = f"{gstr(fxn_for_op[ast.op], ast.op)}({gstr(ast.arg.val)}, {gstr(ast.arg.dtype)})" + if ast.op == BufferOps.CONST: tmp = f"{gstr(fxn_for_op[ast.op], ast.op)}({gstr(ast.arg.val)}, {gstr(ast.arg.dtype)})" else: tmp = f"{gstr(fxn_for_op[UnaryOps.CAST], UnaryOps.CAST)}(inputs[{ast.arg.idx-1}], ({gstr(ast.arg.dtype)}, True))" for mop,arg in ast.arg.st.to_movement_ops(): tmp = f"{gstr(fxn_for_op[mop], mop)}({tmp}, {gstr(arg)})" else: