diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 247c611250..94569476ce 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -55,11 +55,11 @@ class TestDTypeFromUOp(unittest.TestCase): invalid = UOp.invalid() self.assertIs(invalid.dtype, dtypes.bool) self.assertIs(UOp.const(Invalid, dtypes.float32), invalid) - self.assertIs((moved:=invalid.reshape((1,))).cast(dtypes.float32), moved) scratch = Tensor.invalids(4, dtype=dtypes.float32) self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort() if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool)) invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32) + for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid) for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)), UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)), UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared) diff --git a/tinygrad/mixin/dtype.py b/tinygrad/mixin/dtype.py index 08d112da74..b46bac675b 100644 --- a/tinygrad/mixin/dtype.py +++ b/tinygrad/mixin/dtype.py @@ -30,7 +30,7 @@ class DTypeMixin: print(t.dtype, t.numpy()) ``` """ - return self if self.dtype == (dt:=to_dtype(dtype)) or self._uop.base.is_invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt)) + return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt)) def bitcast(self, dtype:DTypeLike) -> Self: """ diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index b6f19facfd..aa746bb454 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -24,6 +24,7 @@ class ElementwiseMixin(CreationMixin): out_dtype = least_upper_dtype(x.dtype, y.dtype) # keep weak CONST weak, might lift weakint -> weakfloat def promote(t): + if t._uop.base.is_invalid: return t # invalid bool is weak const if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.val, weak_dtype(out_dtype))) return t.cast(out_dtype) return promote(x), promote(y) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 4f8eabd954..ccf39be4dc 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -801,7 +801,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # arg is the other srcs; all are cast to the promoted dtype, spec requires STACK srcs to match its dtype srcs = (self,)+tuple(arg) dtype = cast(DType, dtype_from_uop(Ops.STACK, srcs, None)) - return UOp(Ops.STACK, dtype, tuple(u.cast(dtype) for u in srcs)) + # TODO: why cast here? + return UOp(Ops.STACK, dtype, tuple(u if u.base.is_invalid else u.cast(dtype) for u in srcs)) case _: raise RuntimeError(f"{op} is not a MovementOp") usrcs = [shape_to_shape_arg(arg) for arg in src_args] if len(usrcs) == 0: return UOp(op, src=(self,), arg=arg) @@ -1762,7 +1763,7 @@ def lower_weak_node(u:UOp) -> UOp|None: if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary else unwrap(dtype_from_uop(u.op, src, u.arg))) - return u.replace(dtype=None, src=src[:start]+tuple(s.cast(dt) for s in src[start:])).cast(u.dtype) + return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else s.cast(dt) for s in src[start:])).cast(u.dtype) pm_lower_weak = PatternMatcher([ (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)), # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default.