From f5b00de3199f2d7fc0ce2d568dd9103aa200cc85 Mon Sep 17 00:00:00 2001 From: Teddy Tennant Date: Sun, 30 Aug 2026 13:17:26 -0400 Subject: [PATCH] fix weak const promote dropping the movement ops (#17846) Co-authored-by: chenyu --- test/unit/test_dtype_weak.py | 1 + tinygrad/mixin/elementwise.py | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 6495ed5b7b..c247392956 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -240,6 +240,7 @@ class TestWeakBounds(unittest.TestCase): def test_padded_weak_const_keeps_its_zeros(self): self.assertEqual(Tensor(1).expand(1).cat(Tensor(2).expand(2), Tensor(3).expand(3)).tolist(), [1, 2, 2, 3, 3, 3]) self.assertEqual((Tensor(5).reshape(1).pad((1, 1)) == 5).tolist(), [False, True, False]) + self.assertEqual((Tensor(5).reshape(1,1).expand(1,2).pad(((0,2),(0,0))) + Tensor([[1],[2],[3]])).tolist(), [[6,6],[2,2],[3,3]]) class TestWeakStorageBoundary(unittest.TestCase): # weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index a640356be2..326cc3d64d 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -1,7 +1,7 @@ import math, functools, operator from typing import TYPE_CHECKING, Literal, Self from tinygrad.uop import Ops -from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype, least_upper_float, weak_dtype +from tinygrad.dtype import dtypes, ConstType, DType, PyConst, least_upper_dtype, least_upper_float, weak_dtype from tinygrad.helpers import argfix, polyN from tinygrad.mixin.creation import CreationMixin @@ -9,6 +9,9 @@ if TYPE_CHECKING: from tinygrad.uop.ops import UOp, sint +def remint(u:'UOp', dt:DType) -> 'UOp': + return u.const_like(u.val, dt) if u.op is Ops.CONST else u.replace(src=(remint(u.src[0], dt),)+u.src[1:]) + class ElementwiseMixin(CreationMixin): # required to implement def alu(self, op: Ops, *src: Self) -> Self: @@ -26,7 +29,7 @@ class ElementwiseMixin(CreationMixin): 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 and t._uop.vmin == t._uop.vmax: - return t._wrap_uop(t._uop.const_like(t._uop.base.val, weak_dtype(out_dtype))) + return t._wrap_uop(remint(t._uop, weak_dtype(out_dtype))) return t.cast(out_dtype) return promote(x), promote(y)