From d81b03cff4064edfc4fccb9d82f31927e774b944 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 19 Mar 2026 05:02:01 -0400 Subject: [PATCH] pad_to to mixin [pr] (#15365) --- tinygrad/mixin/movement.py | 3 +++ tinygrad/tensor.py | 5 ----- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index aafb9d67b0..a7b0b0bc1e 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -174,6 +174,9 @@ class MovementMixin: def shrink_to(self, shape, *args) -> Self: return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)])) + def pad_to(self, shape, *args) -> Self: + return self._mop(Ops.PAD, tuple([(0, 0 if ns is None else ns-s) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)])) + def view(self, shape, *args) -> Self: """`.view` is an alias for `.reshape`.""" return self.reshape(shape, *args) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index c5f191e3e0..9d5080f1ab 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1168,11 +1168,6 @@ class Tensor(OpMixin): if mode in {"reflect", "replicate"}: return self._pad_reflect_replicate(pX, mode) raise NotImplementedError(f"{mode=} is not supported") - # convenience - def pad_to(self, shape, *args): - if len(new_shape := argfix(shape, *args)) != self.ndim: raise ValueError(f"dim mismatch, cannot pad {self.shape} to {new_shape}") - return self.pad(tuple([None if ns is None else (0, ns-s) for s,ns in zip(self.shape, new_shape)])) - # ***** movement high level ops ***** def _getitem(self, indices, v: Tensor|None = None) -> Tensor: