From 471bd30d16efa78c24fbac33552251e351de6931 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Tue, 14 Oct 2025 17:50:39 +0800 Subject: [PATCH 1/6] cleanup viz/serve.py (#12665) * use load_pickle * update comment --- tinygrad/uop/ops.py | 2 +- tinygrad/viz/serve.py | 13 ++++++------- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index cd18db40bf..079a30b7c7 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -850,7 +850,7 @@ class PatternMatcher: TRACK_MATCH_STATS = ContextVar("TRACK_MATCH_STATS", 2 if VIZ else 0) match_stats:dict[UPat, list[int|float]] = dict() -# TRACK_MATCH_STATS>=3 saves the UOp fields +# TRACK_MATCH_STATS>=2 or VIZ=1 saves all matches ucount = itertools.count() uop_fields:dict[int, tuple] = {} diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 1d88bb4625..0729247f7c 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -287,8 +287,8 @@ def reloader(): os.execv(sys.executable, [sys.executable] + sys.argv) time.sleep(0.1) -def load_pickle(path:pathlib.Path|None) -> list: - if path is None or not path.exists(): return [] +def load_pickle(fp:str) -> list: + if not (path:=pathlib.Path(fp)).exists(): return [] with path.open("rb") as f: return pickle.load(f) # NOTE: using HTTPServer forces a potentially slow socket.getfqdn @@ -296,8 +296,8 @@ class TCPServerWithReuse(socketserver.TCPServer): allow_reuse_address = True if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--kernels', type=pathlib.Path, help='Path to kernels', default=pathlib.Path(temp("rewrites.pkl", append_user=True))) - parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True))) + parser.add_argument('--kernels', type=load_pickle, help='Path to kernels', default=pathlib.Path(temp("rewrites.pkl", append_user=True))) + parser.add_argument('--profile', type=load_pickle, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True))) args = parser.parse_args() with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: @@ -308,9 +308,8 @@ if __name__ == "__main__": st = time.perf_counter() print("*** viz is starting") - ctxs = get_metadata(load_pickle(args.kernels)) - - profile_ret = get_profile(load_pickle(args.profile)) + ctxs = get_metadata(args.kernels) + profile_ret = get_profile(args.profile) server = TCPServerWithReuse(('', PORT), Handler) reloader_thread = threading.Thread(target=reloader) From 1e6e5a0efdd95eab4132e6217975b504765759f8 Mon Sep 17 00:00:00 2001 From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com> Date: Tue, 14 Oct 2025 11:57:38 +0200 Subject: [PATCH 2/6] `parse_valid` returns None instead of raising (#12663) * parse_valid returns None * change there too --- tinygrad/codegen/late/devectorizer.py | 4 ++-- tinygrad/uop/symbolic.py | 11 +++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 6a973c1aed..c0012b73ff 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -20,8 +20,8 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: # can drop valid if idx is out of bound when valid is False drop_stmt = [] for stmt in valid.split_uop(Ops.AND): - try: X, is_upper_bound, c = parse_valid(stmt) - except ValueError: return None + if (res:=parse_valid(stmt)) is None: continue + X, is_upper_bound, c = res # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)): diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 3e43d13161..5111d61ab6 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -386,7 +386,7 @@ symbolic_flat = symbolic+PatternMatcher([ # ******** we take a small aside to "simplify_valid" to rewrite valids ******** -def parse_valid(valid:UOp) -> tuple[UOp, bool, int]: +def parse_valid(valid:UOp) -> tuple[UOp, bool, int]|None: # if it's X <= c, returns X, True, c # if it's X >= c, returns X, False, c @@ -395,7 +395,7 @@ def parse_valid(valid:UOp) -> tuple[UOp, bool, int]: (s0:=valid.src[0]).op is Ops.CMPLT and dtypes.is_int(s0.src[0].dtype): return s0.src[0], False, int(s0.src[1].vmin) # X < c -> X <= c-1 if valid.op is Ops.CMPLT and dtypes.is_int(valid.src[0].dtype): return valid.src[0], True, int((valid.src[1]).vmax)-1 - raise ValueError(f"not able to parse {valid=}") + return None def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp: # return simplified uop (might be the same as input) @@ -403,8 +403,8 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp: # first, parse valid into {expr: (lower_bound, upper_bound)} bounds:defaultdict[UOp, list[ConstType|None]] = defaultdict(lambda: [None, None]) for stmt in valid.split_uop(Ops.AND): - try: expr, is_upper, c = parse_valid(stmt) - except ValueError: continue # give up if we cannot parse the valid + if (res:=parse_valid(stmt)) is None: continue + expr, is_upper, c = res bounds[expr][int(is_upper)] = c # don't simplify any other gates, can lead to OOB, we substitute them back later @@ -444,8 +444,7 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp: def _valid_priority(v: UOp, valids:list[UOp]): # we want valid that's in other valids' parents to be first, so it's more likely the other valids get simplified - try: return sum(-1 if parse_valid(v)[0] in other.toposort() else 0 for other in valids) - except ValueError: return 0 + return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids) def simplify_valid(valid:UOp) -> UOp|None: if valid.op_in_backward_slice_with_self(Ops.LOAD): return None # this should only be for indexing, skip if there's a LOAD From d3bfcd3277be5ec57da88b16259c9420f6f272ed Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Tue, 14 Oct 2025 18:07:46 +0800 Subject: [PATCH 3/6] minor patches for SQTT over usb on gfx12 (#12627) * disable cpu_access in the sqtt buffer allocation not sure if this is required, it results in a very slow call to pcie_mem_write over USB GPU, removing it worked fine. * fix itrace_se_mask on gfx12 on gfx11 it gave 6 se, on gfx11 this value is 2 so no instructions were traced. * Revert "fix itrace_se_mask on gfx12" This reverts commit 0644adbcd1e84b7a617b34b436f9dd63d8270305. --- tinygrad/runtime/ops_amd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index af239b8948..d35539b495 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -819,7 +819,7 @@ class AMDDevice(HCQCompiled): f"ppfeaturemask={(ppfeaturemask&~0x8000):#x} (current {ppfeaturemask=:#x} & ~PP_GFXOFF_MASK) to amdgpu module parameters\n" "For more information read https://github.com/tinygrad/tinygrad/blob/master/extra/sqtt/README.md") SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine - self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(cpu_access=True, nolru=True)) for _ in range(self.se_cnt)] + self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE*1024*1024, BufferSpec(nolru=True)) for _ in range(self.se_cnt)] self.sqtt_itrace_se_mask = getenv("SQTT_ITRACE_SE_MASK", 2) # -1 enable all, 0 disable all, >0 bitmask for where to enable instruction tracing self.sqtt_next_cmd_id = itertools.count(0) cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_start(self.sqtt_buffers, self.sqtt_itrace_se_mask).submit(self) From 0c9d47deab501d5e3bf062f79ce62bf837159436 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Tue, 14 Oct 2025 18:33:12 +0800 Subject: [PATCH 4/6] hcq: add alignment to kernargs (#12669) --- tinygrad/runtime/support/hcq.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index 44592409b1..b7dcf12167 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -310,7 +310,7 @@ class HCQProgram(Generic[HCQDeviceType]): Returns: Arguments state with the given buffers and values set for the program. """ - argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size), + argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size, 8), size=self.kernargs_alloc_size) return self.args_state_t(argsbuf, self, bufs, vals=vals) From 4918c827c282729e69b978fef38dee3146290c9e Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Tue, 14 Oct 2025 18:34:34 +0800 Subject: [PATCH 5/6] amd: lib_gpu does not need cpu_access (#12670) --- tinygrad/runtime/ops_amd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index d35539b495..2c2736b340 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -458,7 +458,7 @@ class AMDProgram(HCQProgram): if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack(' Date: Tue, 14 Oct 2025 19:13:55 +0800 Subject: [PATCH 6/6] fix up some slow tests that launch python (#12672) * fix up some slow tests that launch python * svd nonfull in parallel * split test_advancedindex --- test/external/external_test_dev_var.py | 39 + test/test_tensor.py | 27 - test/unit/test_device.py | 9 +- test/unit/test_indexing.py | 937 +++++++++++++------------ test/unit/test_linalg.py | 28 +- tinygrad/device.py | 9 +- 6 files changed, 533 insertions(+), 516 deletions(-) create mode 100644 test/external/external_test_dev_var.py diff --git a/test/external/external_test_dev_var.py b/test/external/external_test_dev_var.py new file mode 100644 index 0000000000..41abbe8e79 --- /dev/null +++ b/test/external/external_test_dev_var.py @@ -0,0 +1,39 @@ +import subprocess, unittest, os, sys +from tinygrad.device import Device + +class TestTinygradSlow(unittest.TestCase): + def test_env_overwrite_default_device(self): + subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + + if Device.DEFAULT != "CPU": + # setting multiple devices fail + with self.assertRaises(subprocess.CalledProcessError): + subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + + # setting device via DEV + subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + + with self.assertRaises(subprocess.CalledProcessError): + subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], + shell=True, check=True) + +class TestRunAsModule(unittest.TestCase): + def test_module_runs(self): + p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE, + env={**os.environ, "DEBUG": "1"}, timeout=40,) + out = (p.stdout + p.stderr).decode() + self.assertEqual(p.returncode, 0, msg=out) + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_tensor.py b/test/test_tensor.py index 43b8202dc4..617eb242a3 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -1,4 +1,3 @@ -import subprocess import numpy as np import torch import unittest, copy, mmap, random, math, array @@ -515,32 +514,6 @@ class TestTinygrad(unittest.TestCase): print(a) print(c) - def test_env_overwrite_default_device(self): - subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - - if Device.DEFAULT != "CPU": - # setting multiple devices fail - with self.assertRaises(subprocess.CalledProcessError): - subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - - # setting device via DEV - subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - - with self.assertRaises(subprocess.CalledProcessError): - subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'], - shell=True, check=True) - def test_no_attributeerror_after_apply_uop_exception(self): try: Tensor.arange(4).reshape(3,2) diff --git a/test/unit/test_device.py b/test/unit/test_device.py index 984feaf751..e1eaaa1314 100644 --- a/test/unit/test_device.py +++ b/test/unit/test_device.py @@ -1,7 +1,7 @@ #!/usr/bin/env python -import unittest, os, subprocess, sys +import unittest, os, subprocess from tinygrad import Tensor -from tinygrad.device import Device, Compiler +from tinygrad.device import Device, Compiler, enumerate_devices_str from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, WIN, CI class TestDevice(unittest.TestCase): @@ -100,10 +100,7 @@ class TestCompiler(unittest.TestCase): class TestRunAsModule(unittest.TestCase): def test_module_runs(self): - p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE, - env={**os.environ, "DEBUG": "1"}, timeout=40,) - out = (p.stdout + p.stderr).decode() - self.assertEqual(p.returncode, 0, msg=out) + out = '\n'.join(enumerate_devices_str()) self.assertIn("CPU", out) # for sanity check if __name__ == "__main__": diff --git a/test/unit/test_indexing.py b/test/unit/test_indexing.py index 36a9885aa8..7d6240db6a 100644 --- a/test/unit/test_indexing.py +++ b/test/unit/test_indexing.py @@ -180,474 +180,6 @@ class TestIndexing(unittest.TestCase): # def delitem(): del reference[0] # self.assertRaises(TypeError, delitem) - # TODO: LLVM is quite fast, why are other compiled backends slow? - @unittest.skipIf(CI and Device.DEFAULT in ["CPU", "CL", "METAL", "NV", "AMD"], "slow") - def test_advancedindex(self): - # integer array indexing - - # pick a random valid indexer type - def ri(indices): - choice = random.randint(0, 2) - if choice == 0: return Tensor(indices) - if choice == 1: return list(indices) - return tuple(indices) - - def validate_indexing(x): - numpy_testing_assert_equal_helper(x[[0]], consec((1,))) - numpy_testing_assert_equal_helper(x[ri([0]),], consec((1,))) - numpy_testing_assert_equal_helper(x[ri([3]),], consec((1,), 4)) - numpy_testing_assert_equal_helper(x[[2, 3, 4]], consec((3,), 3)) - numpy_testing_assert_equal_helper(x[ri([2, 3, 4]),], consec((3,), 3)) - numpy_testing_assert_equal_helper(x[ri([0, 2, 4]),], np.array([1, 3, 5])) - - def validate_setting(x): - x[[0]] = -2 - numpy_testing_assert_equal_helper(x[[0]], np.array([-2])) - x[[0]] = -1 - numpy_testing_assert_equal_helper(x[ri([0]), ], np.array([-1])) - x[[2, 3, 4]] = 4 - numpy_testing_assert_equal_helper(x[[2, 3, 4]], np.array([4, 4, 4])) - x[ri([2, 3, 4]), ] = 3 - numpy_testing_assert_equal_helper(x[ri([2, 3, 4]), ], np.array([3, 3, 3])) - x[ri([0, 2, 4]), ] = Tensor([5, 4, 3]) - numpy_testing_assert_equal_helper(x[ri([0, 2, 4]), ], np.array([5, 4, 3])) - - # Case 1: Purely Integer Array Indexing - reference = consec((10,)) - validate_indexing(reference) - # setting values - validate_setting(reference) - - # Tensor with stride != 1 - # strided is [1, 3, 5, 7] - - # # TODO: set stride - # reference = consec((10,)) - # strided = set_(reference, (4,), (2,), 0) - - # numpy_testing_assert_equal_helper(strided[[0]], np.array([1])) - # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([1])) - # numpy_testing_assert_equal_helper(strided[ri([3]), ], np.array([7])) - # numpy_testing_assert_equal_helper(strided[[1, 2]], np.array([3, 5])) - # numpy_testing_assert_equal_helper(strided[ri([1, 2]), ], np.array([3, 5])) - # numpy_testing_assert_equal_helper(strided[ri([[2, 1], [0, 3]]), ], - # np.array([[5, 3], [1, 7]])) - - # stride is [4, 8] - - # strided = set_(reference, (2,), (4,), offset=4) - - # numpy_testing_assert_equal_helper(strided[[0]], np.array([5])) - # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([5])) - # numpy_testing_assert_equal_helper(strided[ri([1]), ], np.array([9])) - # numpy_testing_assert_equal_helper(strided[[0, 1]], np.array([5, 9])) - # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ], np.array([5, 9])) - # numpy_testing_assert_equal_helper(strided[ri([[0, 1], [1, 0]]), ], - # np.array([[5, 9], [9, 5]])) - - # reference is 1 2 - # 3 4 - # 5 6 - reference = consec((3, 2)) - numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([1, 3, 5])) - numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([2, 4, 6])) - numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], consec((1,))) - numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], consec((1,), 6)) - numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([1, 2])) - numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 2]), ri([1])]], np.array([2, 4, 4, 2, 6])) - numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 2, 3, 3])) - - rows = ri([[0, 0], - [1, 2]]) - columns = [0], - numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 1], - [3, 5]])) - - rows = ri([[0, 0], - [1, 2]]) - columns = ri([1, 0]) - numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[2, 1], - [4, 5]])) - rows = ri([[0, 0], - [1, 2]]) - columns = ri([[0, 1], - [1, 0]]) - numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 2], - [4, 5]])) - - # setting values - reference[ri([0]), ri([1])] = -1 - numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])], np.array([-1])) - reference[ri([0, 1, 2]), ri([0])] = Tensor([-1, 2, -4]) - numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], - np.array([-1, 2, -4])) - reference[rows, columns] = Tensor([[4, 6], [2, 3]]) - numpy_testing_assert_equal_helper(reference[rows, columns], - np.array([[4, 6], [2, 3]])) - - # Verify still works with Transposed (i.e. non-contiguous) Tensors - reference = Tensor([[0, 1, 2, 3], - [4, 5, 6, 7], - [8, 9, 10, 11]]).T - - # Transposed: [[0, 4, 8], - # [1, 5, 9], - # [2, 6, 10], - # [3, 7, 11]] - - numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([0, 1, 2])) - numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([4, 5, 6])) - numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], np.array([0])) - numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], np.array([6])) - numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([0, 4])) - numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 3]), ri([1])]], np.array([4, 5, 5, 4, 7])) - numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([0, 4, 1, 1])) - - rows = ri([[0, 0], - [1, 2]]) - columns = [0], - numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 0], [1, 2]])) - - rows = ri([[0, 0], - [1, 2]]) - columns = ri([1, 0]) - numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[4, 0], [5, 2]])) - rows = ri([[0, 0], - [1, 3]]) - columns = ri([[0, 1], - [1, 2]]) - numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 4], [5, 11]])) - - # TODO: non contiguous setitem - ''' - # setting values - reference[ri([0]), ri([1])] = -1 - numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])], - np.array([-1])) - reference[ri([0, 1, 2]), ri([0])] = np.array([-1, 2, -4]) - numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], - np.array([-1, 2, -4])) - reference[rows, columns] = np.array([[4, 6], [2, 3]]) - numpy_testing_assert_equal_helper(reference[rows, columns], - np.array([[4, 6], [2, 3]])) - ''' - - # stride != 1 - - # strided is [[1 3 5 7], - # [9 11 13 15]] - - # # TODO: set stride - # reference = Tensor.arange(0., 24).reshape(3, 8) - # strided = set_(reference, (2,4), (8,2), 1) - - # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([0])], np.array([1, 9])) - # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1])], np.array([3, 11])) - # numpy_testing_assert_equal_helper(strided[ri([0]), ri([0])], np.array([1])) - # numpy_testing_assert_equal_helper(strided[ri([1]), ri([3])], np.array([15])) - # numpy_testing_assert_equal_helper(strided[[ri([0, 0]), ri([0, 3])]], np.array([1, 7])) - # numpy_testing_assert_equal_helper(strided[[ri([1]), ri([0, 1, 1, 0, 3])]], np.array([9, 11, 11, 9, 15])) - # numpy_testing_assert_equal_helper(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 3, 9, 9])) - - # rows = ri([[0, 0], - # [1, 1]]) - # columns = [0], - # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 1], [9, 9]])) - - # rows = ri([[0, 1], - # [1, 0]]) - # columns = ri([1, 2]) - # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[3, 13], [11, 5]])) - # rows = ri([[0, 0], - # [1, 1]]) - # columns = ri([[0, 1], - # [1, 2]]) - # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 3], [11, 13]])) - - # setting values - - # strided is [[10, 11], - # [17, 18]] - - # # TODO: set stride - # reference = Tensor.arange(0., 24).reshape(3, 8) - # strided = set_(reference, (2,2), (7,1), 10) - - # numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], np.array([11])) - - # TODO non contiguous setitem - ''' - strided[ri([0]), ri([1])] = -1 - numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], - Tensor([-1])) - ''' - # # TODO: set stride - # reference = Tensor.arange(0., 24).reshape(3, 8) - # strided = set_(reference, (2,2), (7,1), 10) - - # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], np.array([11, 17])) - - # TODO non contiguous setitem - ''' - strided[ri([0, 1]), ri([1, 0])] = Tensor([-1, 2]) - numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], - Tensor([-1, 2])) - ''' - - # # TODO: set stride - # reference = Tensor.arange(0., 24).realize().reshape(3, 8) - # strided = set_(reference, (2,2), (7,1), 10) - - # rows = ri([[0], - # [1]]) - # columns = ri([[0, 1], - # [0, 1]]) - # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[10, 11], [17, 18]])) - - # TODO non contiguous setitem - ''' - strided[rows, columns] = Tensor([[4, 6], [2, 3]]) - numpy_testing_assert_equal_helper(strided[rows, columns], - Tensor([[4, 6], [2, 3]])) - ''' - - # Tests using less than the number of dims, and ellipsis - - # reference is 1 2 - # 3 4 - # 5 6 - reference = consec((3, 2)) - numpy_testing_assert_equal_helper(reference[ri([0, 2]),], np.array([[1, 2], [5, 6]])) - numpy_testing_assert_equal_helper(reference[ri([1]), ...], np.array([[3, 4]])) - numpy_testing_assert_equal_helper(reference[..., ri([1])], np.array([[2], [4], [6]])) - - # verify too many indices fails - with self.assertRaises(IndexError): reference[ri([1]), ri([0, 2]), ri([3])] - - # test invalid index fails - reference = Tensor.empty(10) - for err_idx in (10, -11): - with self.assertRaises(IndexError): - reference[err_idx] - # NOTE cannot check for out of bounds with Tensor indexing - # see tensor.py: __getitem__ (Tiny Things) - ''' - with self.assertRaises(IndexError): - reference[Tensor([err_idx], dtype=dtypes.int64)] - with self.assertRaises(IndexError): - reference[[err_idx]] - ''' - - def tensor_indices_to_np(tensor: Tensor, indices): - npt = tensor.numpy() - idxs = tuple(i.numpy().tolist() if isinstance(i, Tensor) and i.dtype == dtypes.int64 else - i for i in indices) - return npt, idxs - - def get_numpy(tensor, indices): - npt, idxs = tensor_indices_to_np(tensor, indices) - return Tensor(npt[idxs]) - - def set_numpy(tensor:Tensor, indices, value): - if not isinstance(value, int): - value = value.numpy() - npt, idxs = tensor_indices_to_np(tensor, indices) - npt[idxs] = value - return npt - - def assert_get_eq(tensor, indexer): - numpy_testing_assert_equal_helper(tensor[indexer], get_numpy(tensor, indexer)) - - def assert_set_eq(tensor: Tensor, indexer, val): - pyt = clone(tensor) - numt = clone(tensor) - pyt[indexer] = val - numt = set_numpy(numt, indexer, val) - numpy_testing_assert_equal_helper(pyt, numt) - - # NOTE: torch initiates the gradients using g0cpu (rand as gradients) - def assert_backward_eq(tensor: Tensor, indexer): - cpu = clone(tensor.float()) - cpu.requires_grad = True - outcpu = cpu[indexer].sum() - outcpu.backward() - dev = cpu.detach() - dev.requires_grad = True - outdev = dev[indexer].sum() - outdev.backward() - numpy_testing_assert_equal_helper(cpu.grad, dev.grad) - - def get_set_tensor(indexed: Tensor, indexer): - set_size = indexed[indexer].shape - set_count = indexed[indexer].numel() - set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size) #.cast(dtypes.float64) - return set_tensor - - # Tensor is 0 1 2 3 4 - # 5 6 7 8 9 - # 10 11 12 13 14 - # 15 16 17 18 19 - reference = Tensor.arange(0., 20).reshape(4, 5) - - indices_to_test = [ - # grab the second, fourth columns - [slice(None), [1, 3]], - - # first, third rows, - [[0, 2], slice(None)], - - # weird shape - [slice(None), [[0, 1], - [2, 3]]], - # negatives - [[-1], [0]], - [[0, 2], [-1]], - [slice(None), [-1]], - ] - - # only test dupes on gets - get_indices_to_test = indices_to_test + [[slice(None), [0, 1, 1, 2, 2]]] - - for indexer in get_indices_to_test: - assert_get_eq(reference, indexer) - assert_backward_eq(reference, indexer) - - for indexer in indices_to_test: - assert_set_eq(reference, indexer, 44) - assert_set_eq(reference, indexer, get_set_tensor(reference, indexer)) - - reference = Tensor.arange(0., 160).reshape(4, 8, 5) - - indices_to_test = [ - [slice(None), slice(None), [0, 3, 4]], - [slice(None), [2, 4, 5, 7], slice(None)], - [[2, 3], slice(None), slice(None)], - [slice(None), [0, 2, 3], [1, 3, 4]], - [slice(None), [0], [1, 2, 4]], - [slice(None), [0, 1, 3], [4]], - [slice(None), [[0, 1], [1, 0]], [[2, 3]]], - [slice(None), [[0, 1], [2, 3]], [[0]]], - [slice(None), [[5, 6]], [[0, 3], [4, 4]]], - [[0, 2, 3], [1, 3, 4], slice(None)], - [[0], [1, 2, 4], slice(None)], - [[0, 1, 3], [4], slice(None)], - [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)], - [[[0, 1], [1, 0]], [[2, 3]], slice(None)], - [[[0, 1], [2, 3]], [[0]], slice(None)], - [[[2, 1]], [[0, 3], [4, 4]], slice(None)], - [[[2]], [[0, 3], [4, 1]], slice(None)], - # non-contiguous indexing subspace - [[0, 2, 3], slice(None), [1, 3, 4]], - - # less dim, ellipsis - [[0, 2], ], - [[0, 2], slice(None)], - [[0, 2], Ellipsis], - [[0, 2], slice(None), Ellipsis], - [[0, 2], Ellipsis, slice(None)], - [[0, 2], [1, 3]], - [[0, 2], [1, 3], Ellipsis], - [Ellipsis, [1, 3], [2, 3]], - [Ellipsis, [2, 3, 4]], - [Ellipsis, slice(None), [2, 3, 4]], - [slice(None), Ellipsis, [2, 3, 4]], - - # ellipsis counts for nothing - [Ellipsis, slice(None), slice(None), [0, 3, 4]], - [slice(None), Ellipsis, slice(None), [0, 3, 4]], - [slice(None), slice(None), Ellipsis, [0, 3, 4]], - [slice(None), slice(None), [0, 3, 4], Ellipsis], - [Ellipsis, [[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)], - [[[0, 1], [1, 0]], [[2, 1], [3, 5]], Ellipsis, slice(None)], - [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None), Ellipsis], - ] - - for indexer in indices_to_test: - assert_get_eq(reference, indexer) - - assert_set_eq(reference, indexer, 212) - assert_set_eq(reference, indexer, get_set_tensor(reference, indexer)) - assert_backward_eq(reference, indexer) - - reference = Tensor.arange(0., 1296).reshape(3, 9, 8, 6) - - indices_to_test = [ - [slice(None), slice(None), slice(None), [0, 3, 4]], - [slice(None), slice(None), [2, 4, 5, 7], slice(None)], - [slice(None), [2, 3], slice(None), slice(None)], - [[1, 2], slice(None), slice(None), slice(None)], - [slice(None), slice(None), [0, 2, 3], [1, 3, 4]], - [slice(None), slice(None), [0], [1, 2, 4]], - [slice(None), slice(None), [0, 1, 3], [4]], - [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3]]], - [slice(None), slice(None), [[0, 1], [2, 3]], [[0]]], - [slice(None), slice(None), [[5, 6]], [[0, 3], [4, 4]]], - [slice(None), [0, 2, 3], [1, 3, 4], slice(None)], - [slice(None), [0], [1, 2, 4], slice(None)], - [slice(None), [0, 1, 3], [4], slice(None)], - [slice(None), [[0, 1], [3, 4]], [[2, 3], [0, 1]], slice(None)], - [slice(None), [[0, 1], [3, 4]], [[2, 3]], slice(None)], - [slice(None), [[0, 1], [3, 2]], [[0]], slice(None)], - [slice(None), [[2, 1]], [[0, 3], [6, 4]], slice(None)], - [slice(None), [[2]], [[0, 3], [4, 2]], slice(None)], - [[0, 1, 2], [1, 3, 4], slice(None), slice(None)], - [[0], [1, 2, 4], slice(None), slice(None)], - [[0, 1, 2], [4], slice(None), slice(None)], - [[[0, 1], [0, 2]], [[2, 4], [1, 5]], slice(None), slice(None)], - [[[0, 1], [1, 2]], [[2, 0]], slice(None), slice(None)], - [[[2, 2]], [[0, 3], [4, 5]], slice(None), slice(None)], - [[[2]], [[0, 3], [4, 5]], slice(None), slice(None)], - [slice(None), [3, 4, 6], [0, 2, 3], [1, 3, 4]], - [slice(None), [2, 3, 4], [1, 3, 4], [4]], - [slice(None), [0, 1, 3], [4], [1, 3, 4]], - [slice(None), [6], [0, 2, 3], [1, 3, 4]], - [slice(None), [2, 3, 5], [3], [4]], - [slice(None), [0], [4], [1, 3, 4]], - [slice(None), [6], [0, 2, 3], [1]], - [slice(None), [[0, 3], [3, 6]], [[0, 1], [1, 3]], [[5, 3], [1, 2]]], - [[2, 2, 1], [0, 2, 3], [1, 3, 4], slice(None)], - [[2, 0, 1], [1, 2, 3], [4], slice(None)], - [[0, 1, 2], [4], [1, 3, 4], slice(None)], - [[0], [0, 2, 3], [1, 3, 4], slice(None)], - [[0, 2, 1], [3], [4], slice(None)], - [[0], [4], [1, 3, 4], slice(None)], - [[1], [0, 2, 3], [1], slice(None)], - [[[1, 2], [1, 2]], [[0, 1], [2, 3]], [[2, 3], [3, 5]], slice(None)], - - # less dim, ellipsis - [Ellipsis, [0, 3, 4]], - [Ellipsis, slice(None), [0, 3, 4]], - [Ellipsis, slice(None), slice(None), [0, 3, 4]], - [slice(None), Ellipsis, [0, 3, 4]], - [slice(None), slice(None), Ellipsis, [0, 3, 4]], - [slice(None), [0, 2, 3], [1, 3, 4]], - [slice(None), [0, 2, 3], [1, 3, 4], Ellipsis], - [Ellipsis, [0, 2, 3], [1, 3, 4], slice(None)], - [[0], [1, 2, 4]], - [[0], [1, 2, 4], slice(None)], - [[0], [1, 2, 4], Ellipsis], - [[0], [1, 2, 4], Ellipsis, slice(None)], - [[1], ], - [[0, 2, 1], [3], [4]], - [[0, 2, 1], [3], [4], slice(None)], - [[0, 2, 1], [3], [4], Ellipsis], - [Ellipsis, [0, 2, 1], [3], [4]], - ] - - for indexer in indices_to_test: - assert_get_eq(reference, indexer) - assert_set_eq(reference, indexer, 1333) - assert_set_eq(reference, indexer, get_set_tensor(reference, indexer)) - - indices_to_test += [ - [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3], [3, 0]]], - [slice(None), slice(None), [[2]], [[0, 3], [4, 4]]], - ] - for indexer in indices_to_test: - assert_get_eq(reference, indexer) - assert_set_eq(reference, indexer, 1333) - assert_backward_eq(reference, indexer) - # TODO setitem backward ''' def test_set_item_to_scalar_tensor(self): @@ -1568,5 +1100,474 @@ class TestNumpy(unittest.TestCase): numpy_testing_assert_equal_helper(kernel, kernel2) ''' +def tensor_indices_to_np(tensor: Tensor, indices): + npt = tensor.numpy() + idxs = tuple(i.numpy().tolist() if isinstance(i, Tensor) and i.dtype == dtypes.int64 else + i for i in indices) + return npt, idxs + +def get_numpy(tensor, indices): + npt, idxs = tensor_indices_to_np(tensor, indices) + return Tensor(npt[idxs]) + +def set_numpy(tensor:Tensor, indices, value): + if not isinstance(value, int): + value = value.numpy() + npt, idxs = tensor_indices_to_np(tensor, indices) + npt[idxs] = value + return npt + +def assert_get_eq(tensor, indexer): + numpy_testing_assert_equal_helper(tensor[indexer], get_numpy(tensor, indexer)) + +def assert_set_eq(tensor: Tensor, indexer, val): + pyt = clone(tensor) + numt = clone(tensor) + pyt[indexer] = val + numt = set_numpy(numt, indexer, val) + numpy_testing_assert_equal_helper(pyt, numt) + +# NOTE: torch initiates the gradients using g0cpu (rand as gradients) +def assert_backward_eq(tensor: Tensor, indexer): + cpu = clone(tensor.float()) + cpu.requires_grad = True + outcpu = cpu[indexer].sum() + outcpu.backward() + dev = cpu.detach() + dev.requires_grad = True + outdev = dev[indexer].sum() + outdev.backward() + numpy_testing_assert_equal_helper(cpu.grad, dev.grad) + +def get_set_tensor(indexed: Tensor, indexer): + set_size = indexed[indexer].shape + set_count = indexed[indexer].numel() + set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size) #.cast(dtypes.float64) + return set_tensor + +@unittest.skipIf(CI and Device.DEFAULT in ["CPU", "CL", "METAL", "NV", "AMD"], "slow") +class TestAdvancedIndexing(unittest.TestCase): + def test_integer_array_indexing(self): + # pick a random valid indexer type + def ri(indices): + choice = random.randint(0, 2) + if choice == 0: return Tensor(indices) + if choice == 1: return list(indices) + return tuple(indices) + + def validate_indexing(x): + numpy_testing_assert_equal_helper(x[[0]], consec((1,))) + numpy_testing_assert_equal_helper(x[ri([0]),], consec((1,))) + numpy_testing_assert_equal_helper(x[ri([3]),], consec((1,), 4)) + numpy_testing_assert_equal_helper(x[[2, 3, 4]], consec((3,), 3)) + numpy_testing_assert_equal_helper(x[ri([2, 3, 4]),], consec((3,), 3)) + numpy_testing_assert_equal_helper(x[ri([0, 2, 4]),], np.array([1, 3, 5])) + + def validate_setting(x): + x[[0]] = -2 + numpy_testing_assert_equal_helper(x[[0]], np.array([-2])) + x[[0]] = -1 + numpy_testing_assert_equal_helper(x[ri([0]), ], np.array([-1])) + x[[2, 3, 4]] = 4 + numpy_testing_assert_equal_helper(x[[2, 3, 4]], np.array([4, 4, 4])) + x[ri([2, 3, 4]), ] = 3 + numpy_testing_assert_equal_helper(x[ri([2, 3, 4]), ], np.array([3, 3, 3])) + x[ri([0, 2, 4]), ] = Tensor([5, 4, 3]) + numpy_testing_assert_equal_helper(x[ri([0, 2, 4]), ], np.array([5, 4, 3])) + + # Case 1: Purely Integer Array Indexing + reference = consec((10,)) + validate_indexing(reference) + # setting values + validate_setting(reference) + + # Tensor with stride != 1 + # strided is [1, 3, 5, 7] + + # # TODO: set stride + # reference = consec((10,)) + # strided = set_(reference, (4,), (2,), 0) + + # numpy_testing_assert_equal_helper(strided[[0]], np.array([1])) + # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([1])) + # numpy_testing_assert_equal_helper(strided[ri([3]), ], np.array([7])) + # numpy_testing_assert_equal_helper(strided[[1, 2]], np.array([3, 5])) + # numpy_testing_assert_equal_helper(strided[ri([1, 2]), ], np.array([3, 5])) + # numpy_testing_assert_equal_helper(strided[ri([[2, 1], [0, 3]]), ], + # np.array([[5, 3], [1, 7]])) + + # stride is [4, 8] + + # strided = set_(reference, (2,), (4,), offset=4) + + # numpy_testing_assert_equal_helper(strided[[0]], np.array([5])) + # numpy_testing_assert_equal_helper(strided[ri([0]), ], np.array([5])) + # numpy_testing_assert_equal_helper(strided[ri([1]), ], np.array([9])) + # numpy_testing_assert_equal_helper(strided[[0, 1]], np.array([5, 9])) + # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ], np.array([5, 9])) + # numpy_testing_assert_equal_helper(strided[ri([[0, 1], [1, 0]]), ], + # np.array([[5, 9], [9, 5]])) + + # reference is 1 2 + # 3 4 + # 5 6 + reference = consec((3, 2)) + numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([1, 3, 5])) + numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([2, 4, 6])) + numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], consec((1,))) + numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], consec((1,), 6)) + numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([1, 2])) + numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 2]), ri([1])]], np.array([2, 4, 4, 2, 6])) + numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 2, 3, 3])) + + rows = ri([[0, 0], + [1, 2]]) + columns = [0], + numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 1], + [3, 5]])) + + rows = ri([[0, 0], + [1, 2]]) + columns = ri([1, 0]) + numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[2, 1], + [4, 5]])) + rows = ri([[0, 0], + [1, 2]]) + columns = ri([[0, 1], + [1, 0]]) + numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[1, 2], + [4, 5]])) + + # setting values + reference[ri([0]), ri([1])] = -1 + numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])], np.array([-1])) + reference[ri([0, 1, 2]), ri([0])] = Tensor([-1, 2, -4]) + numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], + np.array([-1, 2, -4])) + reference[rows, columns] = Tensor([[4, 6], [2, 3]]) + numpy_testing_assert_equal_helper(reference[rows, columns], + np.array([[4, 6], [2, 3]])) + + # Verify still works with Transposed (i.e. non-contiguous) Tensors + reference = Tensor([[0, 1, 2, 3], + [4, 5, 6, 7], + [8, 9, 10, 11]]).T + + # Transposed: [[0, 4, 8], + # [1, 5, 9], + # [2, 6, 10], + # [3, 7, 11]] + + numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], np.array([0, 1, 2])) + numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([1])], np.array([4, 5, 6])) + numpy_testing_assert_equal_helper(reference[ri([0]), ri([0])], np.array([0])) + numpy_testing_assert_equal_helper(reference[ri([2]), ri([1])], np.array([6])) + numpy_testing_assert_equal_helper(reference[[ri([0, 0]), ri([0, 1])]], np.array([0, 4])) + numpy_testing_assert_equal_helper(reference[[ri([0, 1, 1, 0, 3]), ri([1])]], np.array([4, 5, 5, 4, 7])) + numpy_testing_assert_equal_helper(reference[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([0, 4, 1, 1])) + + rows = ri([[0, 0], + [1, 2]]) + columns = [0], + numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 0], [1, 2]])) + + rows = ri([[0, 0], + [1, 2]]) + columns = ri([1, 0]) + numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[4, 0], [5, 2]])) + rows = ri([[0, 0], + [1, 3]]) + columns = ri([[0, 1], + [1, 2]]) + numpy_testing_assert_equal_helper(reference[rows, columns], np.array([[0, 4], [5, 11]])) + + # TODO: non contiguous setitem + ''' + # setting values + reference[ri([0]), ri([1])] = -1 + numpy_testing_assert_equal_helper(reference[ri([0]), ri([1])], + np.array([-1])) + reference[ri([0, 1, 2]), ri([0])] = np.array([-1, 2, -4]) + numpy_testing_assert_equal_helper(reference[ri([0, 1, 2]), ri([0])], + np.array([-1, 2, -4])) + reference[rows, columns] = np.array([[4, 6], [2, 3]]) + numpy_testing_assert_equal_helper(reference[rows, columns], + np.array([[4, 6], [2, 3]])) + ''' + + # stride != 1 + + # strided is [[1 3 5 7], + # [9 11 13 15]] + + # # TODO: set stride + # reference = Tensor.arange(0., 24).reshape(3, 8) + # strided = set_(reference, (2,4), (8,2), 1) + + # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([0])], np.array([1, 9])) + # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1])], np.array([3, 11])) + # numpy_testing_assert_equal_helper(strided[ri([0]), ri([0])], np.array([1])) + # numpy_testing_assert_equal_helper(strided[ri([1]), ri([3])], np.array([15])) + # numpy_testing_assert_equal_helper(strided[[ri([0, 0]), ri([0, 3])]], np.array([1, 7])) + # numpy_testing_assert_equal_helper(strided[[ri([1]), ri([0, 1, 1, 0, 3])]], np.array([9, 11, 11, 9, 15])) + # numpy_testing_assert_equal_helper(strided[[ri([0, 0, 1, 1]), ri([0, 1, 0, 0])]], np.array([1, 3, 9, 9])) + + # rows = ri([[0, 0], + # [1, 1]]) + # columns = [0], + # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 1], [9, 9]])) + + # rows = ri([[0, 1], + # [1, 0]]) + # columns = ri([1, 2]) + # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[3, 13], [11, 5]])) + # rows = ri([[0, 0], + # [1, 1]]) + # columns = ri([[0, 1], + # [1, 2]]) + # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[1, 3], [11, 13]])) + + # setting values + + # strided is [[10, 11], + # [17, 18]] + + # # TODO: set stride + # reference = Tensor.arange(0., 24).reshape(3, 8) + # strided = set_(reference, (2,2), (7,1), 10) + + # numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], np.array([11])) + + # TODO non contiguous setitem + ''' + strided[ri([0]), ri([1])] = -1 + numpy_testing_assert_equal_helper(strided[ri([0]), ri([1])], + Tensor([-1])) + ''' + # # TODO: set stride + # reference = Tensor.arange(0., 24).reshape(3, 8) + # strided = set_(reference, (2,2), (7,1), 10) + + # numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], np.array([11, 17])) + + # TODO non contiguous setitem + ''' + strided[ri([0, 1]), ri([1, 0])] = Tensor([-1, 2]) + numpy_testing_assert_equal_helper(strided[ri([0, 1]), ri([1, 0])], + Tensor([-1, 2])) + ''' + + # # TODO: set stride + # reference = Tensor.arange(0., 24).realize().reshape(3, 8) + # strided = set_(reference, (2,2), (7,1), 10) + + # rows = ri([[0], + # [1]]) + # columns = ri([[0, 1], + # [0, 1]]) + # numpy_testing_assert_equal_helper(strided[rows, columns], np.array([[10, 11], [17, 18]])) + + # TODO non contiguous setitem + ''' + strided[rows, columns] = Tensor([[4, 6], [2, 3]]) + numpy_testing_assert_equal_helper(strided[rows, columns], + Tensor([[4, 6], [2, 3]])) + ''' + + # Tests using less than the number of dims, and ellipsis + + # reference is 1 2 + # 3 4 + # 5 6 + reference = consec((3, 2)) + numpy_testing_assert_equal_helper(reference[ri([0, 2]),], np.array([[1, 2], [5, 6]])) + numpy_testing_assert_equal_helper(reference[ri([1]), ...], np.array([[3, 4]])) + numpy_testing_assert_equal_helper(reference[..., ri([1])], np.array([[2], [4], [6]])) + + # verify too many indices fails + with self.assertRaises(IndexError): reference[ri([1]), ri([0, 2]), ri([3])] + + # test invalid index fails + reference = Tensor.empty(10) + for err_idx in (10, -11): + with self.assertRaises(IndexError): + reference[err_idx] + # NOTE cannot check for out of bounds with Tensor indexing + # see tensor.py: __getitem__ (Tiny Things) + ''' + with self.assertRaises(IndexError): + reference[Tensor([err_idx], dtype=dtypes.int64)] + with self.assertRaises(IndexError): + reference[[err_idx]] + ''' + + def test_numpy_parity_and_backward_2d(self): + # Tensor is 0 1 2 3 4 + # 5 6 7 8 9 + # 10 11 12 13 14 + # 15 16 17 18 19 + reference = Tensor.arange(0., 20).reshape(4, 5) + + indices_to_test = [ + # grab the second, fourth columns + [slice(None), [1, 3]], + + # first, third rows, + [[0, 2], slice(None)], + + # weird shape + [slice(None), [[0, 1], + [2, 3]]], + # negatives + [[-1], [0]], + [[0, 2], [-1]], + [slice(None), [-1]], + ] + + # only test dupes on gets + get_indices_to_test = indices_to_test + [[slice(None), [0, 1, 1, 2, 2]]] + + for indexer in get_indices_to_test: + assert_get_eq(reference, indexer) + assert_backward_eq(reference, indexer) + + for indexer in indices_to_test: + assert_set_eq(reference, indexer, 44) + assert_set_eq(reference, indexer, get_set_tensor(reference, indexer)) + + def test_numpy_parity_and_backward_3d(self): + reference = Tensor.arange(0., 160).reshape(4, 8, 5) + + indices_to_test = [ + [slice(None), slice(None), [0, 3, 4]], + [slice(None), [2, 4, 5, 7], slice(None)], + [[2, 3], slice(None), slice(None)], + [slice(None), [0, 2, 3], [1, 3, 4]], + [slice(None), [0], [1, 2, 4]], + [slice(None), [0, 1, 3], [4]], + [slice(None), [[0, 1], [1, 0]], [[2, 3]]], + [slice(None), [[0, 1], [2, 3]], [[0]]], + [slice(None), [[5, 6]], [[0, 3], [4, 4]]], + [[0, 2, 3], [1, 3, 4], slice(None)], + [[0], [1, 2, 4], slice(None)], + [[0, 1, 3], [4], slice(None)], + [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)], + [[[0, 1], [1, 0]], [[2, 3]], slice(None)], + [[[0, 1], [2, 3]], [[0]], slice(None)], + [[[2, 1]], [[0, 3], [4, 4]], slice(None)], + [[[2]], [[0, 3], [4, 1]], slice(None)], + # non-contiguous indexing subspace + [[0, 2, 3], slice(None), [1, 3, 4]], + + # less dim, ellipsis + [[0, 2], ], + [[0, 2], slice(None)], + [[0, 2], Ellipsis], + [[0, 2], slice(None), Ellipsis], + [[0, 2], Ellipsis, slice(None)], + [[0, 2], [1, 3]], + [[0, 2], [1, 3], Ellipsis], + [Ellipsis, [1, 3], [2, 3]], + [Ellipsis, [2, 3, 4]], + [Ellipsis, slice(None), [2, 3, 4]], + [slice(None), Ellipsis, [2, 3, 4]], + + # ellipsis counts for nothing + [Ellipsis, slice(None), slice(None), [0, 3, 4]], + [slice(None), Ellipsis, slice(None), [0, 3, 4]], + [slice(None), slice(None), Ellipsis, [0, 3, 4]], + [slice(None), slice(None), [0, 3, 4], Ellipsis], + [Ellipsis, [[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None)], + [[[0, 1], [1, 0]], [[2, 1], [3, 5]], Ellipsis, slice(None)], + [[[0, 1], [1, 0]], [[2, 1], [3, 5]], slice(None), Ellipsis], + ] + + for indexer in indices_to_test: + assert_get_eq(reference, indexer) + + assert_set_eq(reference, indexer, 212) + assert_set_eq(reference, indexer, get_set_tensor(reference, indexer)) + assert_backward_eq(reference, indexer) + + def test_numpy_parity_and_backward_4d(self): + reference = Tensor.arange(0., 1296).reshape(3, 9, 8, 6) + + indices_to_test = [ + [slice(None), slice(None), slice(None), [0, 3, 4]], + [slice(None), slice(None), [2, 4, 5, 7], slice(None)], + [slice(None), [2, 3], slice(None), slice(None)], + [[1, 2], slice(None), slice(None), slice(None)], + [slice(None), slice(None), [0, 2, 3], [1, 3, 4]], + [slice(None), slice(None), [0], [1, 2, 4]], + [slice(None), slice(None), [0, 1, 3], [4]], + [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3]]], + [slice(None), slice(None), [[0, 1], [2, 3]], [[0]]], + [slice(None), slice(None), [[5, 6]], [[0, 3], [4, 4]]], + [slice(None), [0, 2, 3], [1, 3, 4], slice(None)], + [slice(None), [0], [1, 2, 4], slice(None)], + [slice(None), [0, 1, 3], [4], slice(None)], + [slice(None), [[0, 1], [3, 4]], [[2, 3], [0, 1]], slice(None)], + [slice(None), [[0, 1], [3, 4]], [[2, 3]], slice(None)], + [slice(None), [[0, 1], [3, 2]], [[0]], slice(None)], + [slice(None), [[2, 1]], [[0, 3], [6, 4]], slice(None)], + [slice(None), [[2]], [[0, 3], [4, 2]], slice(None)], + [[0, 1, 2], [1, 3, 4], slice(None), slice(None)], + [[0], [1, 2, 4], slice(None), slice(None)], + [[0, 1, 2], [4], slice(None), slice(None)], + [[[0, 1], [0, 2]], [[2, 4], [1, 5]], slice(None), slice(None)], + [[[0, 1], [1, 2]], [[2, 0]], slice(None), slice(None)], + [[[2, 2]], [[0, 3], [4, 5]], slice(None), slice(None)], + [[[2]], [[0, 3], [4, 5]], slice(None), slice(None)], + [slice(None), [3, 4, 6], [0, 2, 3], [1, 3, 4]], + [slice(None), [2, 3, 4], [1, 3, 4], [4]], + [slice(None), [0, 1, 3], [4], [1, 3, 4]], + [slice(None), [6], [0, 2, 3], [1, 3, 4]], + [slice(None), [2, 3, 5], [3], [4]], + [slice(None), [0], [4], [1, 3, 4]], + [slice(None), [6], [0, 2, 3], [1]], + [slice(None), [[0, 3], [3, 6]], [[0, 1], [1, 3]], [[5, 3], [1, 2]]], + [[2, 2, 1], [0, 2, 3], [1, 3, 4], slice(None)], + [[2, 0, 1], [1, 2, 3], [4], slice(None)], + [[0, 1, 2], [4], [1, 3, 4], slice(None)], + [[0], [0, 2, 3], [1, 3, 4], slice(None)], + [[0, 2, 1], [3], [4], slice(None)], + [[0], [4], [1, 3, 4], slice(None)], + [[1], [0, 2, 3], [1], slice(None)], + [[[1, 2], [1, 2]], [[0, 1], [2, 3]], [[2, 3], [3, 5]], slice(None)], + + # less dim, ellipsis + [Ellipsis, [0, 3, 4]], + [Ellipsis, slice(None), [0, 3, 4]], + [Ellipsis, slice(None), slice(None), [0, 3, 4]], + [slice(None), Ellipsis, [0, 3, 4]], + [slice(None), slice(None), Ellipsis, [0, 3, 4]], + [slice(None), [0, 2, 3], [1, 3, 4]], + [slice(None), [0, 2, 3], [1, 3, 4], Ellipsis], + [Ellipsis, [0, 2, 3], [1, 3, 4], slice(None)], + [[0], [1, 2, 4]], + [[0], [1, 2, 4], slice(None)], + [[0], [1, 2, 4], Ellipsis], + [[0], [1, 2, 4], Ellipsis, slice(None)], + [[1], ], + [[0, 2, 1], [3], [4]], + [[0, 2, 1], [3], [4], slice(None)], + [[0, 2, 1], [3], [4], Ellipsis], + [Ellipsis, [0, 2, 1], [3], [4]], + ] + + for indexer in indices_to_test: + assert_get_eq(reference, indexer) + assert_set_eq(reference, indexer, 1333) + assert_set_eq(reference, indexer, get_set_tensor(reference, indexer)) + + indices_to_test += [ + [slice(None), slice(None), [[0, 1], [1, 0]], [[2, 3], [3, 0]]], + [slice(None), slice(None), [[2]], [[0, 3], [4, 4]]], + ] + for indexer in indices_to_test: + assert_get_eq(reference, indexer) + assert_set_eq(reference, indexer, 1333) + assert_backward_eq(reference, indexer) + if __name__ == '__main__': unittest.main() diff --git a/test/unit/test_linalg.py b/test/unit/test_linalg.py index a54418b162..5647e4faa6 100644 --- a/test/unit/test_linalg.py +++ b/test/unit/test_linalg.py @@ -26,18 +26,22 @@ class TestLinAlg(unittest.TestCase): orthogonality_helper(V) reconstruction_helper([U,s_diag,V],a) - def test_svd_nonfull(self): - sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)] - for size in sizes: - a = Tensor.randn(size).realize() - U,S,V = a.svd(full_matrices=False) - b_shape,m,n = size[0:-2],size[-2],size[-1] - k = min(m,n) - s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k))) - #reduced U,V is only orthogonal along smaller dim - if (m < n): orthogonality_helper(U),orthogonality_helper(V) - else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1)) - reconstruction_helper([U,s_diag,V],a) + def _test_svd_nonfull(self, size): + a = Tensor.randn(size).realize() + U,S,V = a.svd(full_matrices=False) + b_shape,m,n = size[0:-2],size[-2],size[-1] + k = min(m,n) + s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k))) + #reduced U,V is only orthogonal along smaller dim + if (m < n): orthogonality_helper(U),orthogonality_helper(V) + else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1)) + reconstruction_helper([U,s_diag,V],a) + + # faster for parallel pytest + def test_svd_nonfull_2_2(self): self._test_svd_nonfull((2,2)) + def test_svd_nonfull_5_3(self): self._test_svd_nonfull((5,3)) + def test_svd_nonfull_3_5(self): self._test_svd_nonfull((3,5)) + def test_svd_nonfull_2_2_2_2_3(self): self._test_svd_nonfull((2,2,2,2,3)) @unittest.skip("very big. recommend wrapping with TinyJit around inner function") def test_svd_large(self): diff --git a/tinygrad/device.py b/tinygrad/device.py index 3e1788c946..bd021ebea1 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, replace from collections import defaultdict -from typing import Any, Generic, TypeVar, Iterator, Sequence, cast +from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup @@ -357,7 +357,7 @@ if PROFILE: from tinygrad.uop.ops import launch_viz launch_viz("PROFILE", fn) -if __name__ == "__main__": +def enumerate_devices_str() -> Generator[str, None, None]: from tinygrad import Tensor, Device for device in ALL_DEVICES: @@ -376,4 +376,7 @@ if __name__ == "__main__": result = (colored('PASS', 'green') if any_works else f"{colored('FAIL', 'yellow')}") + ''.join([f'\n{" "*16} {x}' for x in compilers_results]) except Exception as e: result = f"{colored('FAIL', 'red')} {e}" - print(f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}") + yield f"{'*' if device == Device.DEFAULT else ' '} {device:10s}: {result}" + +if __name__ == "__main__": + for s in enumerate_devices_str(): print(s)