IQ.Pilot Release Commit @ 0798119

tinygrad_repo/test/null/test_symbolic_tensor.py

@@ -0,0 +1,112 @@
+import unittest
+from tinygrad import Variable
+from tinygrad.tensor import Tensor
+
+class TestSymbolic(unittest.TestCase):
+  def assert_tuple_equal(self, x, y):
+    for a,b in zip(x,y): self.assertFalse(a != b)
+
+  def test_cat_dim0_is_expanded(self):
+    i = Variable("i", 1, 5).bind(3)
+    j = Variable("j", 1, 5).bind(3)
+    k = Variable("k", 1, 5).bind(3)
+    t = Tensor.rand(5, 4)[:i].cat(Tensor.rand(5, 4)[:j], dim=0).cat(Tensor.rand(5, 4)[:k], dim=0)
+    self.assert_tuple_equal(t.shape, (i+j+k, 4))
+    t = Tensor.rand(5, 3)[:i].cat(Tensor.rand(5, 3)[:i], dim=0).cat(Tensor.rand(3, 3), dim=0)
+    self.assert_tuple_equal(t.shape, (2*i+3, 3))
+
+  def test_cat_dim1_strides(self):
+    i = Variable("i", 1, 5).bind(4)
+    j = Variable("j", 1, 5).bind(4)
+    k = Variable("k", 1, 5).bind(4)
+    t = Tensor.rand(3, 5)[:, :i].cat(Tensor.rand(3, 5)[:, :j], dim=1).cat(Tensor.rand(3, 5)[:, :k], dim=1)
+    self.assert_tuple_equal(t.shape, (3, i+j+k))
+
+class TestSymbolicVarVals(unittest.TestCase):
+  def assert_equal(self, x, y): self.assertFalse(x != y)
+
+  def test_shrink_unbind(self):
+    v = Variable("v", 1, 100)
+    bv = Variable("v", 1, 100).bind(2)
+    t = Tensor.rand(3, 4).shrink(((0,bv),(0,4)))
+    unbound_st, var_val = t.uop.unbind_all()
+    assert var_val == {v: 2}
+    t = Tensor.rand(3, 4).shrink(((bv, bv+1), (0, 4)))
+    unbound_st, var_val = t.uop.unbind_all()
+    assert var_val == {v: 2}
+
+class TestSymbolicReshape(unittest.TestCase):
+  def test_reshape(self):
+    a = Tensor.rand(5, 4)
+    b = Tensor.rand(5, 6)
+    for i in range(1, 6):
+      vi = Variable("i", 1, 5).bind(i)
+      ret = a[:vi]
+      ret = ret.reshape((vi, 4))
+      assert ret.shape == (vi, 4)
+      ret = b[:vi]
+      ret = ret.reshape((vi, 2, 3))
+      assert ret.shape == (vi, 2, 3)
+
+  def test_two_symbol_reshape(self):
+    t = Tensor.rand(5, 5)
+    for i in range(1, 6):
+      for j in range(1, 6):
+        vi = Variable("i", 1, 5).bind(i)
+        vj = Variable("j", 1, 5).bind(j)
+        ret = t[:vi, :vj]
+        ret = ret.reshape(vj, vi)
+        assert ret.shape == (vj, vi)
+        ret = ret.reshape(vi, vj)
+        assert ret.shape == (vi, vj)
+        ret = ret.reshape(1, vi*vj)
+        assert ret.shape == (1, vi*vj)
+
+class TestSymbolicExpand(unittest.TestCase):
+  def test_expand_into_symbols(self):
+    vi = Variable("i", 1, 5).bind(3)
+    vj = Variable("j", 1, 5).bind(3)
+    a = Tensor([[1], [2], [3]]).expand((3, vi))
+    assert a.shape == (3, vi)
+    a = a.reshape(3, vi, 1).expand((3, vi, vj))
+    assert a.shape == (3, vi, vj)
+
+  def test_plus_expands_constant(self):
+    a = Tensor.rand(3, 5)
+    for i in range(1, 6):
+      vi = Variable("i", 1, 5).bind(i)
+      ret = a[:, :vi]
+      ret = ret + 1
+      self.assertTupleEqual(ret.shape, (3, vi))
+
+  def test_pad_then_expand_into_symbols(self):
+    vi = Variable("i", 1, 10).bind(3)
+    a = Tensor(1).unsqueeze(0).pad((0, 24)).unsqueeze(0).expand((vi, 25))
+    self.assertEqual(a.shape, (vi, 25))
+    self.assertEqual(a.reshape(25*vi).shape, (vi*25,))
+    self.assertEqual(a.reshape(vi*25).shape, (vi*25,))
+
+class TestSymbolicShrink(unittest.TestCase):
+  def test_shrink_symbols_simple(self):
+    vi = Variable("i", 1, 5)
+    t = Tensor.rand(5, 5).shrink(((0, 5),(0,vi)))
+    assert t.shape == (5, vi)
+
+  def test_shrink_symbols(self):
+    vi = Variable("i", 1, 5)
+    t = Tensor.rand(3, 5).shrink(((0, 2), (vi, vi+1)))
+    assert t.shape == (2, 1)
+
+class TestSymbolicContiguousViewOffset(unittest.TestCase):
+  def test_shrink_from_start(self):
+    v = Variable("v", 1, 10).bind(5)
+    t = Tensor.rand(10).realize().shrink(((0, v),))
+    self.assertEqual(t.uop.contiguous_view_offset(), 0)
+
+  def test_shrink_with_offset(self):
+    v = Variable("v", 1, 7).bind(4)
+    t = Tensor.rand(10).realize().shrink(((3, 3+v),))
+    self.assertEqual(t.uop.contiguous_view_offset(), 3)
+
+if __name__ == '__main__':
+  unittest.main()