diff --git a/test/null/test_uops.py b/test/null/test_uops.py index ea701a30ca..09f6e9a55a 100644 --- a/test/null/test_uops.py +++ b/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) diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index af8f7c7267..9272b7e617 100644 --- a/tinygrad/dtype.py +++ b/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 diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 8341eca529..4168089812 100644 --- a/tinygrad/uop/ops.py +++ b/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]) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 69b13657d0..a6af821b61 100644 --- a/tinygrad/uop/spec.py +++ 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),