diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 50da9ac84c..f06c1d11ff 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -221,7 +221,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): float_one_bits = uint_bits.ones_like(dtype=dtype).bitcast(uint_dtype) return uint_bits.rshift(dtype.bitsize - nmant).bitwise_or(float_one_bits).bitcast(dtype)[:prod(shape)].sub(1).reshape(shape) - def _pad_constant(self, pX, value:float) -> Self: + def _pad_constant(self, pX, value:ConstType) -> Self: # shrink first for negative pads, then pad with only non-negative values pX = tuple((0, 0) if p is None else p for p in pX) has_neg = not all(resolve(p >= 0) for p in flatten(pX)) @@ -230,7 +230,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): base = MovementMixin.pad(X, pads) if value == 0: return base base = base.cast(least_upper_dtype(base.dtype, dtypes.from_py(value))) - return base + MovementMixin.pad(X.ones_like(), pads).cast(dtypes.bool).where(base.zeros_like(), base.full_like(value)) + return MovementMixin.pad(X.ones_like(dtype=dtypes.bool), pads).where(base, base.full_like(value)) def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self: if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, self.shape)): raise ValueError('Padding value causes wrapping around more than once.') @@ -243,7 +243,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): for d,(pB,pA) in enumerate(pads): if mode == "reflect": if pB >= (s:=X.shape[d]) or pA>=s: raise ValueError(f"Padding ({pB}, {pA}) should be less than the input size={s} for dim={d}.") - slcB, slcA = slice(pB,0,-1), slice(s-2 if s-2>=0 else None, s-2-pA if s-2-pA>=0 else None, -1) + slcB, slcA = slice(pB,0,-1), slice(s-2, s-2-pA if s-2-pA>=0 else None, -1) xB, xA = (X[[slc if i == d else slice(None) for i in range(X.ndim)]] if p > 0 else None for slc, p in ((slcB, pB), (slcA, pA))) else: shrB, shrA = tuple((0,1) if i==d else None for i in range(X.ndim)), tuple((X.shape[i]-1,X.shape[i]) if i==d else None for i in range(X.ndim)) @@ -253,7 +253,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): # shrink after for negative pads (reflection/replication must see full data first) return X.shrink(tuple((-min(pB,0), min(pA+s,s)) for (pB,pA),s in zip(pX, X.shape))) - def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str="constant", value:float=0.0) -> Self: + def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str="constant", value:ConstType=0.0) -> Self: """ Returns a tensor with padding applied based on the input `padding`. @@ -1144,16 +1144,16 @@ class OpMixin(ElementwiseMixin, ReduceMixin): ``` """ axis = tuple(range(-len(k_ := make_tuple(kernel_size, 2)), 0)) + s_ = stride if stride is not None else k_ def pool(x:Self, padding_:Sequence[int]) -> Self: - return x._pad_constant(((0,0),)*(x.ndim-len(k_)) + flat_to_grouped(padding_), 0.0)._pool(k_, stride if stride is not None else k_, dilation) + return x._pad_constant(((0,0),)*(x.ndim-len(k_)) + flat_to_grouped(padding_), 0.0)._pool(k_, s_, dilation) reg_pads = resolve_pool_pads(padding, len(k_)) - ceil_pads = self._apply_ceil_mode(reg_pads, k_, stride if stride is not None else k_, dilation) + pads = self._apply_ceil_mode(reg_pads, k_, s_, dilation) if ceil_mode else reg_pads if not count_include_pad: - pads = ceil_pads if ceil_mode else reg_pads return pool(self, pads).sum(axis) / pool(self.ones_like(), pads).sum(axis) - if not ceil_mode: return pool(self, reg_pads).mean(axis) - return pool(self, ceil_pads).sum(axis) / pool(self._pad_constant(((0,0),)*(self.ndim-len(k_)) + flat_to_grouped(reg_pads), 0.0).ones_like(), - tuple(cp-rp for cp,rp in zip(ceil_pads, reg_pads))).sum(axis) + if not ceil_mode: return pool(self, pads).mean(axis) + return pool(self, pads).sum(axis) / pool(self._pad_constant(((0,0),)*(self.ndim-len(k_)) + flat_to_grouped(reg_pads), 0.0).ones_like(), + tuple(cp-rp for cp,rp in zip(pads, reg_pads))).sum(axis) def max_pool2d(self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0, ceil_mode=False, return_indices=False) -> Self | tuple[Self, Self]: