sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

simplify equivalent const max (#17505)

Browse Source 
* add const max folds

* add regression test

* move
This commit is contained in:
Raine
2026-08-12 08:39:29 -07:00
committed by GitHub
parent 4f106ebe87
commit de04781b36
2 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_uop_symbolic.py
+5
View File
@@ -948,6 +948,11 @@ class TestSymbolic(unittest.TestCase):
self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")
def test_equivalent_const_max(self):
x = Variable("x", -10, 10)
self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)")
self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)")
def test_where_combine(self):
cond = Variable("x", 0, 3) < 2
a = Variable("a", 0, 3)
tinygrad/uop/symbolic.py
+2
View File
@@ -249,6 +249,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
(UPat(Ops.RANGE, src=(UPat(Ops.CONST,)), name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
# max folding
((UPat.cvar("a") < UPat.var("b")).where(UPat.var("b"), UPat.cvar("c")), lambda a,b,c: UOp.maximum(a,b) if a.val == c.val else None),
((UPat.var("a") < UPat.cvar("b")).where(UPat.cvar("c"), UPat.var("a")), lambda a,b,c: UOp.maximum(a,b) if b.val == c.val else None),
(UPat.maximum(UPat.var("x"), UPat.var("y")), lambda x,y: x if x.vmin >= y.vmax else y if x.vmax <= y.vmin else None),
# TODO: why does this rule break beautiful_mnist?
#((UPat.var("x")+UPat.var("z")).maximum(UPat.var("y")+UPat.var("z")), lambda x,y,z: x.maximum(y) + z),