IQ.Pilot Release Commit @ 0798119

tinygrad_repo/test/null/test_rewrite_bottom_up_gate.py

@@ -0,0 +1,28 @@
+import unittest
+from tinygrad.uop.ops import PatternMatcher, UOp, graph_rewrite, Ops, UPat, BottomUpGate
+
+def assert_not_reached(): assert False, "This function should not be reached"
+def gate(): raise BottomUpGate
+
+class TestBottomUpGate(unittest.TestCase):
+  def test_basic_bottom_up_gate(self):
+    """Test that BottomUpGate stops bottom-up"""
+    pm = PatternMatcher([
+      (UPat(Ops.ADD), gate),
+      (UPat(Ops.MUL), assert_not_reached)
+    ])
+
+    a,b,c = UOp.variable("a",0,10), UOp.variable("b",0,10), UOp.variable("c",0,10)
+    graph_rewrite((a*a)+(b*c), pm, bottom_up=True)
+
+  def test_bottom_up_gate_with_rewriting(self):
+    pm = PatternMatcher([
+      (UPat.var("a")+UPat.var("a"), lambda a: 2*a),
+      (UPat(Ops.MUL), gate),
+      (UPat(Ops.CONST), assert_not_reached)
+    ])
+    a = UOp.variable("a",0,10)
+    graph_rewrite(a+a, pm, bottom_up=True)
+
+if __name__ == "__main__":
+  unittest.main()