fix parse_valid for float uop (#8681)

x < c -> X <= c-1 only works for int
This commit is contained in:
chenyu
2025-01-19 18:15:49 -05:00
committed by GitHub
parent 168c16646a
commit 2d0842386d
2 changed files with 3 additions and 1 deletions
+2
View File
@@ -662,6 +662,8 @@ class TestOps(unittest.TestCase):
ten0, ten1 = Tensor(data[0], dtype=dtypes.bool), Tensor(data[1], dtype=dtypes.bool)
helper_test_op([], lambda: tor0&tor1, lambda: ten0&ten1, forward_only=True)
helper_test_op(None, lambda x: (1 < x) & (x < 2), forward_only=True, vals=[[1.2, 1.2, 1.2, 3.2]])
self.helper_test_exception([(4), (4)], torch.bitwise_and, Tensor.bitwise_and, expected=RuntimeError)
def test_or(self):
+1 -1
View File
@@ -1083,7 +1083,7 @@ def parse_valid(valid:UOp) -> tuple[UOp, bool, int]:
if valid.op is Ops.CMPNE and valid.src[1].op is Ops.CONST and valid.src[1].arg == 1 and \
(s0:=valid.src[0]).op is Ops.CMPLT and s0.src[1].op is Ops.CONST: return s0.src[0], False, s0.src[1].arg
# X < c -> X <= c-1
if valid.op is Ops.CMPLT and valid.src[1].op is Ops.CONST: return valid.src[0], True, valid.src[1].arg-1
if valid.op is Ops.CMPLT and valid.src[1].op is Ops.CONST and dtypes.is_int(valid.src[0].dtype): return valid.src[0], True, valid.src[1].arg-1
raise ValueError(f"not able to parse {valid=}")
def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None: