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)