CONST related cleanups [pr] (#17352)

ConstFloat(nan) != nan should be False, and some Invalid bool cleanups

test/null/test_uops.py

@@ -5,7 +5,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
 from test.helpers import eval_uop, to_uops_list
@@ -45,6 +45,12 @@ class TestDTypeFromUOp(unittest.TestCase):
       with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program)
       type_verify(UOp.const(value, concrete).sink(), spec_program)
 
+  def test_invalid_stated_dtype(self):
+    # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
+    # and the spec is what rejects a non-bool Invalid
+    self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid())
+    with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared)
+
   def test_invalid_dtype_and_consumers(self):
     invalid = UOp.invalid()
     self.assertIs(invalid.dtype, dtypes.bool)
@@ -112,6 +118,20 @@ class TestSafeCast(unittest.TestCase):
     self.assertEqual(a.cast(dtypes.int8).cast(dtypes.int64).simplify(), a.cast(dtypes.int64))
     self.assertEqual(a.cast(dtypes.int8).cast(dtypes.float).simplify(), a.cast(dtypes.float))
 
+class TestConstFloatEq(unittest.TestCase):
+  def test_nan_eq_ne_agree(self):
+    nan = dtypes.float32.const(math.nan)
+    self.assertTrue(nan == math.nan)
+    self.assertFalse(nan != math.nan)  # float.__ne__ would say True here
+    self.assertFalse(nan == Invalid)
+    self.assertTrue(nan != Invalid)    # __ne__ must defer to the reflected eq, not swallow NotImplemented
+
+  def test_matchers_agree_on_nan(self):
+    n = UOp.const(math.nan, dtypes.float32)
+    for compiled in (False, True):
+      pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled)
+      self.assertTrue(pm.rewrite(n), f"{compiled=}")
+
 class TestExecALU(unittest.TestCase):
   def test_sqrt(self):
     self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)

tinygrad/dtype.py

@@ -17,6 +17,7 @@ class ConstFloat(float):
     if self is other: return True
     if isinstance(other, float) and math.isnan(self) and math.isnan(other): return True
     return float.__eq__(self, other)
+  def __ne__(self, other): return res if (res:=self.__eq__(other)) is NotImplemented else not res  # float.__ne__ disagrees with __eq__ on nan
   def __hash__(self): return hash(self.bits)
   def __repr__(self): return f"ConstFloat({float.__repr__(self)})"
   def __str__(self): return float.__repr__(self)
@@ -74,8 +75,7 @@ class DType(metaclass=DTypeMetaClass):
   def max(self):
     if dtypes.is_int(self): return 2**(self.bitsize)-1+self.min
     return float("inf") if dtypes.is_float(self) else True
-  def const(self, val: tuple[ConstType, ...]|ConstType):
-    if isinstance(val, tuple): return tuple(map(self.const, val))
+  def const(self, val: ConstType):
     if isinstance(val, InvalidType): return val
     # NOTE: float('nan') != float('nan'), so we canonicalize here
     if isinstance(val, float) and math.isnan(val): val = math.nan

tinygrad/uop/ops.py

@@ -189,7 +189,6 @@ class UOpMetaClass(type):
   ucache:dict[tuple, weakref.ReferenceType[UOp]] = {}
   def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
                metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None):
-    if op is Ops.CONST and arg is Invalid: dtype = dtypes.bool
     if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void
     # CONST derives its dtype by value only when the constructor omits one
     # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py
@@ -618,7 +617,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def alu(self, op, *src:UOp, **kwargs): return UOp(op, src=(self, *src), **kwargs)
   @staticmethod
   def const(b:ConstLike, dtype:DType|None=None):
-    if dtype is None: dtype = dtypes.from_py(b)
+    if dtype is None or b is Invalid: dtype = dtypes.from_py(b)
     if isinstance(b, UOp): return b.cast(dtype)
     # NOTE: it always has to be STACK now, even if they are all the same
     if isinstance(b, tuple): return UOp.stack(*[UOp.const(c, dtype) for c in b])

tinygrad/uop/spec.py

@@ -53,8 +53,8 @@ spec_shared = PatternMatcher([
   # NOOP. TODO: remove this
   (UPat(Ops.NOOP), lambda: True),
 
-  # CONST is everywhere
-  (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.val) is type(x.dtype.const(x.val))),
+  # CONST is everywhere; Invalid is a bool const
+  (UPat(Ops.CONST, src=(), name="x"), lambda x: x.dtype is dtypes.bool if x.is_invalid else type(x.val) is type(x.dtype.const(x.val))),
 
   # STACK is everywhere too
   (UPat(Ops.STACK, dtype=dtypes.void, src=()), lambda: True),