symbolic new test
MODIFY
test/test_symbolic.py
test/test_symbolic.py
import unittest
from tinygrad.shape.symbolic import Variable
-def helper_test_variable(v, n, m, s):
- assert v.min == n
- assert v.max == m
- assert str(v) == s
class TestSymbolic(unittest.TestCase):
+ def helper_test_variable(self, v, n, m, s):
+ self.assertEqual(v.min, n)
+ self.assertEqual(v.max, m)
+ self.assertEqual(str(v), s)
+
def test_mul_0(self):
- helper_test_variable(Variable("a", 0, 8)*0, 0, 0, "0")
+ self.helper_test_variable(Variable("a", 0, 8)*0, 0, 0, "0")
def test_mul_1(self):
- helper_test_variable(Variable("a", 0, 8)*1, 0, 8, "a")
+ self.helper_test_variable(Variable("a", 0, 8)*1, 0, 8, "a")
def test_mul_2(self):
- helper_test_variable(Variable("a", 0, 8)*2, 0, 16, "(a*2)")
+ self.helper_test_variable(Variable("a", 0, 8)*2, 0, 16, "(a*2)")
def test_div_1(self):
- helper_test_variable(Variable("a", 0, 8)//1, 0, 8, "a")
+ self.helper_test_variable(Variable("a", 0, 8)//1, 0, 8, "a")
def test_mod_1(self):
- helper_test_variable(Variable("a", 0, 8)%1, 0, 0, "0")
+ self.helper_test_variable(Variable("a", 0, 8)%1, 0, 0, "0")
def test_add_min_max(self):
- helper_test_variable(Variable("a", 0, 8) * 2 + 12, 12, 16+12, "((a*2)+12)")
+ self.helper_test_variable(Variable("a", 0, 8) * 2 + 12, 12, 16+12, "((a*2)+12)")
def test_div_min_max(self):
- helper_test_variable(Variable("a", 0, 7) // 2, 0, 3, "(a//2)")
+ self.helper_test_variable(Variable("a", 0, 7) // 2, 0, 3, "(a//2)")
def test_sum_div_min_max(self):
- helper_test_variable(Variable.sum([Variable("a", 0, 7), Variable("b", 0, 3)]) // 2, 0, 5, "((a+b)//2)")
+ self.helper_test_variable(Variable.sum([Variable("a", 0, 7), Variable("b", 0, 3)]) // 2, 0, 5, "((a+b)//2)")
def test_sum_div_factor(self):
- helper_test_variable(Variable.sum([Variable("a", 0, 7)*4, Variable("b", 0, 3)*4]) // 2, 0, 20, "((a*2)+(b*2))")
+ self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*4, Variable("b", 0, 3)*4]) // 2, 0, 20, "((a*2)+(b*2))")
def test_sum_div_some_factor(self):
- helper_test_variable(Variable.sum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, "((b*2)+((a*5)//2))")
+ self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, "((b*2)+((a*5)//2))")
def test_sum_div_no_factor(self):
- helper_test_variable(Variable.sum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*5]) // 2, 0, 25, "(((a*5)+(b*5))//2)")
+ self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*5]) // 2, 0, 25, "(((a*5)+(b*5))//2)")
def test_mod_factor(self):
- helper_test_variable(Variable.sum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 50, "(((a*100)+(b*50))%100)")
+ self.helper_test_variable(Variable.sum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 50, "(((a*100)+(b*50))%100)")
def test_sum_0(self):
- helper_test_variable(Variable.sum([Variable("a", 0, 7)]), 0, 7, "a")
+ self.helper_test_variable(Variable.sum([Variable("a", 0, 7)]), 0, 7, "a")
def test_mod_remove(self):
- helper_test_variable(Variable("a", 0, 6)%100, 0, 6, "a")
+ self.helper_test_variable(Variable("a", 0, 6)%100, 0, 6, "a")
def test_gt_remove(self):
- helper_test_variable(Variable("a", 0, 6) >= 25, 0, 0, "0")
+ self.helper_test_variable(Variable("a", 0, 6) >= 25, 0, 0, "0")
def test_lt_remove(self):
- helper_test_variable(Variable("a", 0, 6) < -3, 0, 0, "0")
- helper_test_variable(Variable("a", 0, 6) < 3, 0, 1, "(a<3)")
- helper_test_variable(Variable("a", 0, 6) < 8, 1, 1, "1")
+ self.helper_test_variable(Variable("a", 0, 6) < -3, 0, 0, "0")
+ self.helper_test_variable(Variable("a", 0, 6) < 3, 0, 1, "(a<3)")
+ self.helper_test_variable(Variable("a", 0, 6) < 8, 1, 1, "1")
def test_and_fold(self):
- helper_test_variable(Variable.ands([Variable.num(0), Variable("a", 0, 1)]), 0, 0, "0")
+ self.helper_test_variable(Variable.ands([Variable.num(0), Variable("a", 0, 1)]), 0, 0, "0")
def test_and_remove(self):
- helper_test_variable(Variable.ands([Variable.num(1), Variable("a", 0, 1)]), 0, 1, "a")
+ self.helper_test_variable(Variable.ands([Variable.num(1), Variable("a", 0, 1)]), 0, 1, "a")
def test_mod_factor(self):
# this is technically wrong, if b is 0 the output will be negative
- helper_test_variable(Variable.sum([Variable.num(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, -1, 27, "((-1+a)%28)")
+ self.helper_test_variable(Variable.sum([Variable.num(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, -1, 27, "((-1+a)%28)")
def test_div_factor(self):
# TODO: this isn't right
- helper_test_variable(Variable.sum([Variable.num(-44), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) // 40, -1, 9, "(b+-1)")
+ self.helper_test_variable(Variable.sum([Variable.num(-44), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) // 40, -1, 9, "(b+-1)")
def test_mul_div(self):
- helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")
+ self.helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")
def test_mul_div_factor_mul(self):
- helper_test_variable((Variable("a", 0, 10)*8)//4, 0, 20, "(a*2)")
+ self.helper_test_variable((Variable("a", 0, 10)*8)//4, 0, 20, "(a*2)")
def test_mul_div_factor_div(self):
- helper_test_variable((Variable("a", 0, 10)*4)//8, 0, 5, "(a//2)")
+ self.helper_test_variable((Variable("a", 0, 10)*4)//8, 0, 5, "(a//2)")
def test_div_remove(self):
- helper_test_variable(Variable.sum([Variable("idx0", 0, 127)*4, Variable("idx2", 0, 3)])//4, 0, 127, "idx0")
+ self.helper_test_variable(Variable.sum([Variable("idx0", 0, 127)*4, Variable("idx2", 0, 3)])//4, 0, 127, "idx0")
+
+class TestSymbolicNumeric(unittest.TestCase):
+ def helper_test_numeric(self, f):
+ for i in range(10):
+ v = f(Variable.num(i))
+ #print(i, f(i), v.min, v.max)
+ self.assertEqual(v.min, v.max)
+ self.assertEqual(v.min, f(i))
+
+ def test_times_2(self): self.helper_test_numeric(lambda x: x*2)
+ def test_times_2_plus_3(self): self.helper_test_numeric(lambda x: x*2 + 3)
+ def test_times_2_plus_3_mod_4(self): self.helper_test_numeric(lambda x: (x*2 + 3)%4)
if __name__ == '__main__':
unittest.main()

MODIFY
tinygrad/shape/symbolic.py
tinygrad/shape/symbolic.py
@@ -4,7 +4,9 @@ from typing import List
from tinygrad.helpers import partition, modn, all_same
class Node:
- b, min, max = 0, -math.inf, math.inf # make mypy happy
+ b: int
+ min: int
+ max: int
expr: str
def __str__(self):
if self.min == self.max: return str(self.min) # this is universal
@@ -49,6 +51,7 @@ class Node:
else:
a = self
if a.min >= 0 and a.max < b: return a
+ if a.min == a.max: return Variable.num(modn(a.min, b))
return ModNode(a, b)
def __ge__(self, b:int):
if self.max < b: return Variable.num(0)
@@ -96,8 +99,7 @@ class OpNode(Node):
def __init__(self, a:Node, b:int):
self.a, self.b = a,b
self.min, self.max = self.minmax(a,b)
- @staticmethod
- def minmax(a, b): raise NotImplementedError()
+ minmax = staticmethod(lambda a,b: (1//0, 1//0))
@property
def expr(self):
return f"({self.a}{self.op}{self.b})"
@@ -106,8 +108,7 @@ class RedNode(Node):
def __init__(self, nodes:List[Node]):
self.nodes = nodes
self.min, self.max = self.minmax(nodes)
- @staticmethod
- def minmax(nodes): raise NotImplementedError()
+ minmax = staticmethod(lambda nodes: (1//0, 1//0))
@property
def expr(self):
return f"({self.op.join([str(x) for x in self.nodes])})"
@@ -116,6 +117,7 @@ class RedNode(Node):
class MulNode(OpNode): op, minmax = "*", staticmethod(lambda a,b: (a.min*b, a.max*b))
class DivNode(OpNode): op, minmax = "//", staticmethod(lambda a,b: (int(a.min/b), int(a.max/b)))
+# TODO: this is wrong! should be flipped
class ModNode(OpNode): op, minmax = "%", staticmethod(lambda a,b: (min(a.min, 0), max(a.max, b-1)))
class GeNode(OpNode): op, minmax = ">=", staticmethod(lambda a,b: (0,1))
class LtNode(OpNode): op, minmax = "<", staticmethod(lambda a,b: (0,1))

symbolic new test