From e688e077586ff0f6b554b864a928ec295e1e7e9a Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Sat, 15 Aug 2026 20:05:10 -0700
Subject: [PATCH] add max_shape/max_numel to mixins + pad_to (#17553)

---
 test/backend/test_tensor.py | 10 ++++++++++
 tinygrad/mixin/movement.py  | 10 ++++++++++
 tinygrad/mixin/op.py        |  6 ++++++
 tinygrad/uop/ops.py         |  4 ----
 4 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/test/backend/test_tensor.py b/test/backend/test_tensor.py
index c44a2b9e08..c77b8d649b 100644
--- a/test/backend/test_tensor.py
+++ b/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)
diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py
index b916f8af9d..c900993019 100644
--- a/tinygrad/mixin/movement.py
+++ b/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.
diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py
index 94b1dbe4ef..9f65649042 100644
--- a/tinygrad/mixin/op.py
+++ b/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)))
diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index 81f347524f..40363b0350 100644
--- a/tinygrad/uop/ops.py
+++ b/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