add UOp.const_with_shape [pr] (#7825)

* add UOp.const_with_shape [pr]

* lines

test/test_schedule.py

@@ -1873,7 +1873,7 @@ class TestView(unittest.TestCase):
     b = a.pad(((0, 10), None))[10:]
     sched = check_schedule(b.contiguous(), 1)
     # TODO: this VALID can clean up, where do we need st?
-    self.assertIs(store_val(sched[-1]), UOp(Ops.VALID, dtypes.bool, (b.lazydata.st.to_uop(),)).where(x:=UOp.const(b.dtype, 0), x))
+    self.assertIs(store_val(sched[-1]), UOp.const_with_shape(b.dtype, 0, b.lazydata.st.shape))
     run_schedule(sched)
     np.testing.assert_equal(b.numpy(), 0)
 
@@ -1884,7 +1884,7 @@ class TestView(unittest.TestCase):
     assert b.shape == (10, 10)
     sched = check_schedule(b.contiguous(), 1)
     self.assertEqual(sched[-1].ast.full_shape, (10, 10))
-    self.assertIs(store_val(sched[-1]), UOp(Ops.VALID, dtypes.bool, (b.lazydata.st.to_uop(),)).where(x:=UOp.const(b.dtype, 0), x))
+    self.assertIs(store_val(sched[-1]), UOp.const_with_shape(b.dtype, 0, b.lazydata.st.shape))
     run_schedule(sched)
     np.testing.assert_equal(b.numpy(), 0)
 

tinygrad/engine/schedule.py

@@ -369,7 +369,7 @@ do_realize = PatternMatcher([
 
 def generate_valid(ctx:ScheduleContext, b:UOp, to_store:UOp, base:UOp) -> UOp:
   if isinstance((val:=to_store.arg), UOp): ctx.var_vals.update([val.unbind()])
-  return UOp(Ops.VALID, dtypes.bool, (unwrap(base.st).to_uop(),)).where(UOp.const(base.dtype, val), 0)
+  return UOp.const_with_shape(base.dtype, val, unwrap(base.st).shape)
 
 break_sched = PatternMatcher([
   # consts are always fused and generated

tinygrad/ops.py

@@ -340,6 +340,13 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   @property
   def is_contiguous_base(self): return self.op is Ops.CONTIGUOUS and not (self.src[0].base.op is Ops.VIEW and len(self.src[0].base.src) == 2)
 
+  # *** from LazyBuffer ***
+
+  @staticmethod
+  def const_with_shape(dtype:DType, val:ConstLike, shape:Tuple[sint,...]) -> UOp:
+    from tinygrad.shape.shapetracker import ShapeTracker
+    return UOp(Ops.VALID, dtypes.bool, (ShapeTracker.from_shape(()).reshape((1,)*len(shape)).expand(shape).to_uop(),)).where(UOp.const(dtype, val), 0)
+
   # *** uop movement ops ***
 
   @property