Tensor.max_shape on mixin, pad_to value like pad

test/backend/test_tensor.py

@@ -653,9 +653,19 @@ class TestZeroShapeTensor(unittest.TestCase):
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
+    np.testing.assert_equal(Tensor([1, 2]).pad_to(4, value=2).numpy(), [1, 2, 2, 2])
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3, value=-1).numpy(), [[1, 2, -1], [-1, -1, -1]])
+    np.testing.assert_equal(Tensor([1, 2]).pad_to(None, value=5).numpy(), [1, 2])  # no-op pad ignores the fill
     with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
     with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)
 
+  def test_max_shape(self):
+    from tinygrad import UOp
+    t = Tensor.empty(2, UOp.variable('v', 1, 32), 4)
+    self.assertEqual(t.max_shape, (2, 32, 4))
+    self.assertEqual(t.max_numel(), 2*32*4)
+    self.assertEqual(Tensor.empty(2, 3).max_shape, (2, 3))
+
   def test_shrink_into_zero(self):
     t = Tensor.rand(3, 4).realize()
     assert t.shrink((None, (2, 2))).realize().shape == (3, 0)

tinygrad/llm/model.py

@@ -296,9 +296,8 @@ def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor,
   out_shape = v.shape
   if not static:
     # pad the variable-length sequence to its max size with no-op steps: beta=0 and alpha=1 leave the state untouched
-    # pad_to goes through the uop: Tensor has no max_shape
-    q, k, v, beta = (x.pad_to(x.uop.max_shape) for x in (q, k, v, beta))
-    alpha = (alpha - 1).pad_to(alpha.uop.max_shape) + 1
+    q, k, v, beta = (x.pad_to(x.max_shape) for x in (q, k, v, beta))
+    alpha = alpha.pad_to(alpha.max_shape, value=1)
     tokens = q.shape[2]
   core, kq = Tensor.empty(batch, heads, tokens, value_dim), (q*k).sum(-1).contiguous()
   state = state if state.uop.op is Ops.AFTER else state.contiguous()  # keep the AFTER chain of in-place state updates
@@ -334,7 +333,7 @@ class GatedDeltaNetBlock(FFNBlock):
     initial = Tensor(start_pos).eq(0)
     is_kda = hasattr(self, "ssm_g_a")
     symbolic = isinstance(T, UOp)
-    T_pad = int(cast(UOp, T).vmax) if symbolic else T  # symbolic chunks are padded to their max size: one graph serves every size
+    T_pad = x.max_shape[1]  # symbolic chunks are padded to their max size: one graph serves every size
 
     # input processing
     x = x.half()

tinygrad/mixin/movement.py

@@ -46,6 +46,16 @@ class MovementMixin:
     """
     return prod(self.shape)
 
+  @property
+  def max_shape(self) -> tuple[int, ...]:
+    """The shape with every symbolic dimension replaced by its maximum."""
+    from tinygrad.uop.ops import to_max_shape  # deferred: ops.py imports the mixins
+    return to_max_shape(self.shape)
+
+  def max_numel(self) -> int:
+    """The number of elements in `max_shape`."""
+    return prod(self.max_shape)
+
   def size(self, dim:int|None=None) -> sint|tuple[sint, ...]:
     """
     Returns the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.

tinygrad/mixin/op.py

@@ -289,6 +289,12 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if value == 0: return base
     return MovementMixin.pad(X.const_like(True, dtypes.bool), pads).where(base, value)
 
+  def pad_to(self, shape, *args, value:ConstType=0) -> Self:
+    # same mask trick as _pad_constant so the fill survives backends that realize PAD as 0-fill
+    ret = MovementMixin.pad_to(self, shape, *args)
+    if value == 0 or ret is self: return ret
+    return MovementMixin.pad_to(self.const_like(True, dtypes.bool), shape, *args).where(ret, value)
+
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
     # shrink first for negative pads, then wrap the non-negative remainder
     X = self.shrink(tuple((-smin(pB,0), smin(pA+sh,sh)) for (pB,pA),sh in zip(pX, self.shape)))

tinygrad/uop/ops.py

@@ -470,10 +470,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if (ret:=self._shape) is None: raise RuntimeError(f"shape requested, but {self.op} doesn't have a shape")
     return ret
 
-  @property
-  def max_shape(self) -> tuple[int, ...]: return to_max_shape(self.shape)
-  def max_numel(self) -> int: return prod(self.max_shape)
-
   @property
   def shard_shape(self) -> tuple[sint, ...]:
     if not isinstance(self.device, tuple) or self.axis is None: return self.shape