nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

fix sym_infer for CAST (#17338)

Browse Source
This commit is contained in:
chenyu
2026-07-31 14:38:38 -04:00
committed by GitHub
parent 8dc225e28e
commit 277433259e
2 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_uop_symbolic.py
+4
View File
@@ -1202,6 +1202,10 @@ class TestSymInfer(unittest.TestCase):
# floor: 1 % -1000 = -999, 1 // -1000 = -1
assert sym_infer(a%b, var_vals) == -999
assert sym_infer(a//b, var_vals) == -1
def test_sym_infer_with_cast(self):
a = Variable("a", 0, 100, dtypes.int)
assert sym_infer(a.cast(dtypes.long) + 1, {a.expr: 5}) == 6
assert sym_infer(a.cast(dtypes.float) * 0.5, {a.expr: 5}) == 2.5
def test_sym_infer_with_bitcast(self):
a = Variable("a", 1, 10, dtypes.int)
expr = ((a.bitcast(dtypes.uint) << UOp.const(1)).bitcast(dtypes.int) + 2)
tinygrad/uop/render.py
+2
View File
@@ -58,6 +58,8 @@ renderer_infer = PatternMatcher([
(UPat(Ops.CDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
(UPat(Ops.FLOORMOD, name="x"), lambda ctx,x: f"floormod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
(UPat(Ops.FLOORDIV, name="x"), lambda ctx,x: f"floordiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
(UPat(Ops.CAST, name="x"),
lambda ctx,x: f"{'float' if dtypes.is_float(x.dtype) else 'bool' if x.dtype is dtypes.bool else 'int'}({ctx[x.src[0]]})"),
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}, {x.src[0].dtype!r}, {x.dtype!r})"),
]) + renderer