forked from tinygrad/tinygrad
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 geohot
|
2c90f3ea76 | ||
|
geohot
|
d99457657b | ||
|
geohot
|
8a34a4e2c7 |
+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):
|
||||
|
||||
+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)
|
||||
|
||||
Reference in New Issue
Block a user