mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-06-15 17:40:13 +08:00
d59e6e7a37
* move more tests to test/null, split some existing ones * null work * null work * move more * fixes * move PIL * PIL in CLIP * don't move that
167 lines
6.3 KiB
Python
167 lines
6.3 KiB
Python
# modified from
|
|
# https://github.com/arogozhnikov/einops/blob/master/tests/test_examples.py
|
|
# https://github.com/arogozhnikov/einops/blob/master/tests/test_ops.py
|
|
# https://github.com/arogozhnikov/einops/blob/master/tests/test_parsing.py
|
|
|
|
import numpy as np
|
|
import unittest
|
|
from tinygrad import Tensor
|
|
|
|
|
|
class test_rearrange_examples(unittest.TestCase):
|
|
def test_tensor_train_example_numpy(self):
|
|
# kept here just for a collection, only tested for numpy
|
|
# https://arxiv.org/pdf/1509.06569.pdf, (5)
|
|
x = Tensor.ones([3, 4, 5, 6])
|
|
rank = 4
|
|
|
|
# creating appropriate Gs
|
|
Gs = [Tensor.ones([d, d, rank, rank]) for d in x.shape]
|
|
Gs[0] = Gs[0][:, :, :1, :]
|
|
Gs[-1] = Gs[-1][:, :, :, :1]
|
|
|
|
# einsum way
|
|
y = x.reshape((1,) + x.shape)
|
|
for G in Gs:
|
|
# taking partial results left-to-right
|
|
# y = numpy.einsum('i j alpha beta, alpha i ... -> beta ... j', G, y)
|
|
y = Tensor(np.einsum("i j a b, a i ... -> b ... j", G.numpy(), y.numpy()))
|
|
y1 = y.reshape(-1)
|
|
|
|
# alternative way
|
|
y = x.reshape(-1)
|
|
for G in Gs:
|
|
i, j, alpha, beta = G.shape
|
|
y = y.rearrange("(i rest alpha) -> rest (alpha i)", alpha=alpha, i=i)
|
|
y = y @ G.rearrange("i j alpha beta -> (alpha i) (j beta)")
|
|
y = y.rearrange("rest (beta j) -> (beta rest j)", beta=beta, j=j)
|
|
y2 = y
|
|
assert np.allclose(y1.numpy(), y2.numpy())
|
|
|
|
# yet another way
|
|
y = x
|
|
for G in Gs:
|
|
i, j, alpha, beta = G.shape
|
|
y = y.rearrange("i ... (j alpha) -> ... j (alpha i)", alpha=alpha, i=i)
|
|
y = y @ G.rearrange("i j alpha beta -> (alpha i) (j beta)")
|
|
y3 = y.reshape(-1)
|
|
assert np.allclose(y1.numpy(), y3.numpy())
|
|
|
|
|
|
class test_rearrange_ops(unittest.TestCase):
|
|
def test_rearrange_ellipsis_ops(self):
|
|
identity_patterns = [
|
|
"...->...",
|
|
"a b c d e-> a b c d e",
|
|
"a b c d e ...-> ... a b c d e",
|
|
"a b c d e ...-> a ... b c d e",
|
|
"... a b c d e -> ... a b c d e",
|
|
"a ... e-> a ... e",
|
|
"a ... -> a ... ",
|
|
"a ... c d e -> a (...) c d e",
|
|
]
|
|
|
|
equivalent_rearrange_patterns = [
|
|
("a b c d e -> (a b) c d e", "a b ... -> (a b) ... "),
|
|
("a b c d e -> a b (c d) e", "... c d e -> ... (c d) e"),
|
|
("a b c d e -> a b c d e", "... -> ... "),
|
|
("a b c d e -> (a b c d e)", "... -> (...)"),
|
|
("a b c d e -> b (c d e) a", "a b ... -> b (...) a"),
|
|
("a b c d e -> b (a c d) e", "a b ... e -> b (a ...) e"),
|
|
]
|
|
|
|
xnp = np.arange(2 * 3 * 4 * 5 * 6, dtype=np.int32).reshape([2, 3, 4, 5, 6])
|
|
x = Tensor(xnp)
|
|
for pattern in identity_patterns:
|
|
assert np.array_equal(xnp, x.rearrange(pattern).numpy()), pattern
|
|
|
|
for pattern1, pattern2 in equivalent_rearrange_patterns:
|
|
assert np.array_equal(x.rearrange(pattern1).numpy(), x.rearrange(pattern2).numpy())
|
|
|
|
def test_rearrange_consistency(self):
|
|
shape = [1, 2, 3, 5, 7, 11]
|
|
xnp = np.arange(np.prod(shape), dtype=np.int32).reshape(shape)
|
|
x = Tensor(xnp)
|
|
for pattern in [
|
|
"a b c d e f -> a b c d e f",
|
|
"b a c d e f -> a b d e f c",
|
|
"a b c d e f -> f e d c b a",
|
|
"a b c d e f -> (f e) d (c b a)",
|
|
"a b c d e f -> (f e d c b a)",
|
|
]:
|
|
result = x.rearrange(pattern).numpy()
|
|
assert len(np.setdiff1d(xnp, result)) == 0
|
|
assert result.dtype == xnp.dtype
|
|
|
|
result = x.rearrange("a b c d e f -> a (b) (c d e) f").numpy()
|
|
assert np.array_equal(xnp.flatten(), result.flatten())
|
|
|
|
result = x.rearrange("a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11").numpy()
|
|
assert np.array_equal(xnp, result)
|
|
|
|
result1 = x.rearrange("a b c d e f -> f e d c b a").numpy()
|
|
result2 = x.rearrange("f e d c b a -> a b c d e f").numpy()
|
|
assert np.array_equal(result1, result2)
|
|
|
|
result = x.rearrange("a b c d e f -> (f d) c (e b) a").rearrange("(f d) c (e b) a -> a b c d e f", b=2, d=5).numpy()
|
|
assert np.array_equal(xnp, result)
|
|
|
|
sizes = dict(zip("abcdef", shape))
|
|
temp = x.rearrange("a b c d e f -> (f d) c (e b) a", **sizes)
|
|
result = temp.rearrange("(f d) c (e b) a -> a b c d e f", **sizes).numpy()
|
|
assert np.array_equal(xnp, result)
|
|
|
|
x2 = np.arange(2 * 3 * 4, dtype=np.int32).reshape([2, 3, 4])
|
|
result = Tensor(x2).rearrange("a b c -> b c a").numpy()
|
|
assert x2[1, 2, 3] == result[2, 3, 1]
|
|
assert x2[0, 1, 2] == result[1, 2, 0]
|
|
|
|
def test_rearrange_permutations(self):
|
|
# tests random permutation of axes against two independent numpy ways
|
|
for n_axes in range(1, 10):
|
|
x = np.arange(2**n_axes, dtype=np.int32).reshape([2] * n_axes)
|
|
permutation = np.random.permutation(n_axes)
|
|
left_expression = " ".join("i" + str(axis) for axis in range(n_axes))
|
|
right_expression = " ".join("i" + str(axis) for axis in permutation)
|
|
expression = left_expression + " -> " + right_expression
|
|
result = Tensor(x).rearrange(expression).numpy()
|
|
|
|
for pick in np.random.randint(0, 2, [10, n_axes]):
|
|
assert x[tuple(pick)] == result[tuple(pick[permutation])]
|
|
|
|
for n_axes in range(1, 10):
|
|
x = np.arange(2**n_axes, dtype=np.int32).reshape([2] * n_axes)
|
|
permutation = np.random.permutation(n_axes)
|
|
left_expression = " ".join("i" + str(axis) for axis in range(n_axes)[::-1])
|
|
right_expression = " ".join("i" + str(axis) for axis in permutation[::-1])
|
|
expression = left_expression + " -> " + right_expression
|
|
result = Tensor(x).rearrange(expression).numpy()
|
|
assert result.shape == x.shape
|
|
expected_result = np.zeros_like(x)
|
|
for original_axis, result_axis in enumerate(permutation):
|
|
expected_result |= ((x >> original_axis) & 1) << result_axis
|
|
|
|
assert np.array_equal(result, expected_result)
|
|
|
|
|
|
class test_rearrange_parsing(unittest.TestCase):
|
|
def test_unicode_ellipsis(self):
|
|
equivalent_rearrange_patterns = [
|
|
("a b … -> (a b) … ", "a b ... -> (a b) ... "),
|
|
("… c d e -> … (c d) e", "... c d e -> ... (c d) e"),
|
|
("… -> … ", "... -> ... "),
|
|
("… -> (…)", "... -> (...)"),
|
|
("a b … -> b (…) a", "a b ... -> b (...) a"),
|
|
("a b … e -> b (a …) e", "a b ... e -> b (a ...) e"),
|
|
]
|
|
|
|
xnp = np.arange(2 * 3 * 4 * 5 * 6, dtype=np.int32).reshape([2, 3, 4, 5, 6])
|
|
x = Tensor(xnp)
|
|
|
|
for pattern1, pattern2 in equivalent_rearrange_patterns:
|
|
assert np.array_equal(x.rearrange(pattern1).numpy(), x.rearrange(pattern2).numpy())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|