github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

move UOp.pad to mixin (#15657)

Browse Source 
the same arg works for Tensor.pad
This commit is contained in:
chenyu
2026-04-08 13:15:19 -04:00
committed by GitHub
parent 28b14b0e38
commit cb681da840
2 changed files with 6 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tinygrad/mixin/movement.py
View File

@@ -107,6 +107,12 @@ class MovementMixin:
ret = self._mop(Ops.RESHAPE, arg=new_shape)
return self if ret.shape == self.shape else ret
def pad(self, arg:tuple[tuple[sint, sint] | None, ...]) -> Self:
if self.ndim != len(arg):
raise ValueError(f"{self.ndim=} != {len(arg)=}")
ret = self._mop(Ops.PAD, tuple(x if x is not None else (0, 0) for x in arg))
return self if ret.shape == self.shape else ret
def shrink(self, arg: tuple[tuple[sint, sint] | None, ...]) -> Self:
"""
Returns a tensor that shrinks the each axis based on input arg.

10
tinygrad/uop/ops.py
View File

@@ -650,16 +650,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
if ret.shape == self.shape and same_shape_noop: return self
return ret
# in these four, if the shape doesn't change we can return self
#def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
#def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
#def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, same_shape_noop=True)
def pad(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.PAD, arg, same_shape_noop=True)
# in these two, we have custom logic to check if they are a no-op
#def permute(self, arg:tuple[int, ...]): return self._mop(Ops.PERMUTE, arg, same_shape_noop=False) if arg != tuple(range(len(self.shape))) else self
#def flip(self, arg:tuple[bool, ...]): return self._mop(Ops.FLIP, arg, same_shape_noop=False) if any(arg) and len(arg) == len(self.shape) else self
# *** uop UNIQUE ***
# TODO: use this in Buffer