support Invalid in _pad_constant [PR] (#16733)

and simplified image_conv2d
This commit is contained in:
chenyu
2026-06-24 21:03:36 -04:00
committed by GitHub
parent 9d5d2253b3
commit 65dd099b63
2 changed files with 3 additions and 4 deletions
+2 -2
View File
@@ -6,7 +6,7 @@ from tinygrad.mixin.movement import MovementMixin
from tinygrad.mixin.reduce import ReduceMixin
from tinygrad.uop import Ops
from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element
from tinygrad.dtype import ConstType, DTypeLike, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
from tinygrad.dtype import ConstType, DTypeLike, Invalid, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
from tinygrad.helpers import all_int, argfix, argsort, ceildiv, flatten, flat_to_grouped, fully_flatten, get_shape, make_tuple, merge_dicts, prod
from tinygrad.helpers import resolve_pool_pads, round_up
@@ -285,7 +285,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
pads = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) if has_neg else pX
base = MovementMixin.pad(X, pads)
if value == 0: return base
base = base.cast(least_upper_dtype(base.dtype, dtypes.from_py(value)))
if value is not Invalid: base = base.cast(least_upper_dtype(base.dtype, dtypes.from_py(value)))
return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, base.const_like(value))
def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
+1 -2
View File
@@ -859,8 +859,7 @@ class Tensor(RandMixin, metaclass=TensorMeta):
def is_pow2(v): return v > 0 and v & (v - 1) == 0
# pad dimension i to amt with invalids
def ipad(t, i, amt):
shape = (None,)*i + (amt,) + (None,)*(t.ndim-i-1)
return Tensor(True, device=t.device).expand(t.shape).pad_to(shape).where(t.pad_to(shape), Invalid) if amt != t.shape[i] else t
return t.pad(tuple(None if d != i else (0, amt-s) for d,s in enumerate(t.shape)), value=Invalid) if amt != t.shape[i] else t
# align a dimension, use at to specify the dimension to pad in, defaults to first
def pad_align(t, dim, at=None, force=False):
# align to 64 pixels when height is real, otherwise 64 bytes is sufficient