mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 11:56:08 +00:00
Merge branch 'master' into new_shape
This commit is contained in:
George Hotz
+39
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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__":
|
||||
|
||||
+469
-468
@@ -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()
|
||||
|
||||
+16
-12
@@ -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):
|
||||
|
||||
@@ -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)):
|
||||
|
||||
+6
-3
@@ -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)
|
||||
|
||||
@@ -458,7 +458,7 @@ class AMDProgram(HCQProgram):
|
||||
if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack('<q', rel_sym_offset - apply_image_offset + addent) # R_AMDGPU_REL64
|
||||
else: raise RuntimeError(f"unknown AMD reloc {typ}")
|
||||
|
||||
self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(cpu_access=True, nolru=True))
|
||||
self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(nolru=True))
|
||||
self.dev.allocator._copyin(self.lib_gpu, image)
|
||||
self.dev.synchronize()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
+1
-1
@@ -926,7 +926,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] = {}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user