IQ.Pilot Release Commit @ f2a861c

artifacts/package_sources/tinygrad/test/unit/test_rearrange_einops.py

@@ -0,0 +1,166 @@
+# 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()