fix weak const promote dropping the movement ops (#17846)

Co-authored-by: chenyu <[email protected]>

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

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)