diff --git a/test/backend/test_const_folding.py b/test/backend/test_const_folding.py index 0e4239ecfe..0c2c801560 100644 --- a/test/backend/test_const_folding.py +++ b/test/backend/test_const_folding.py @@ -188,7 +188,6 @@ class TestTautologicalCompare(unittest.TestCase): np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False) np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False) - @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly") def test_a_eq_a(self): # self eq is always true for int or bool a = Tensor([1, 2, 3]) diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index 4cecf8116e..83a9355e6b 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -50,8 +50,9 @@ wgsl_matcher = PatternMatcher([ (UPat.store(UPat.var("b"), UPat.var("var"), name="s"), lambda b,var,s: packed_store(b,var) if is_packed(s) else None), (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None), - # fix nan check: 'a != a -> is_nan()' - (UPat.var("a") != UPat.var("a"), is_nan), + # fix nan check: 'a != a -> is_nan()'. the decomp rewrites (a != a).logical_not() to CMPEQ, so match both forms + (UPat.var("a", dtypes.floats) != UPat.var("a"), is_nan), + (UPat.var("a", dtypes.floats).alu(Ops.CMPEQ, UPat.var("a")), lambda a: is_nan(a).ne(True)), ]) class WGSLRenderer(CStyleLanguage):