diff --git a/test/test_ops.py b/test/test_ops.py index 599e03d476..d342a71df4 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -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): diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 62d884bd98..e9cc0e08dd 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -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: