diff --git a/test/test_ops.py b/test/test_ops.py index 9e2437d09f..64eda92e37 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -851,6 +851,8 @@ class TestOps(unittest.TestCase): helper_test_op([(10,20)], lambda x: x.argmax(0, False).type(torch.int32), lambda x: x.argmax(0, False), forward_only=True) helper_test_op([(10,20)], lambda x: x.argmax(1, False).type(torch.int32), lambda x: x.argmax(1, False), forward_only=True) helper_test_op([(10,20)], lambda x: x.argmax(1, True).type(torch.int32), lambda x: x.argmax(1, True), forward_only=True) + # regression test for bitwise_not then argmax + helper_test_op(None, lambda x: (~x).argmax().type(torch.int32), lambda x: (~x).argmax(), forward_only=True, vals=[[2, 2]]) def test_argmin(self): # check if it returns the first index for multiple occurences diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 50a15e8427..2067000469 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -3152,7 +3152,7 @@ class Tensor(SimpleMathTrait): ``` """ if self.dtype != dtypes.bool and not dtypes.is_int(self.dtype): raise RuntimeError(f"{self.dtype} is not supported") - return self.logical_not() if self.dtype == dtypes.bool else self ^ ((1<<8*self.dtype.itemsize)-1) + return self.logical_not() if self.dtype == dtypes.bool else self ^ -1 def lshift(self, x:int): """