diff --git a/extra/onnx_ops.py b/extra/onnx_ops.py index e7ce7b7d7e..3ff3e42e49 100644 --- a/extra/onnx_ops.py +++ b/extra/onnx_ops.py @@ -209,34 +209,13 @@ def _auto_pad(X: Tensor, auto_pad, strides, kernel_shape, dilations): return pad_shape[::2] + pad_shape[1::2] if auto_pad == "SAME_UPPER" else pad_shape[1::2] + pad_shape[::2] raise NotImplementedError(f"auto_pad={auto_pad} not implemented") -def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Tensor=None, axes: Tensor=None, mode="constant", value: float=0.): - constant_value = value if constant_value is None else float(to_python_const(constant_value)) - seq_pads = list(pads) if isinstance(pads, tuple) else to_python_const(pads) - seq_pads = [math.ceil(i) for i in seq_pads] - seq_axes = to_python_const(axes) if axes is not None else None - base_shape = x.shape - pads = _format_padding(seq_pads, ndims=len(x.shape), axes=seq_axes) - if mode == "wrap": - repeat_args = [math.ceil(dim[0]/sh) + math.ceil(dim[1]/sh) + 1 for dim, sh in zip(pads, base_shape)] - new_shape = [s*r for s,r in zip(base_shape, repeat_args)] - shrink_args = [(sh-dim[0]%sh if dim[0]%sh != 0 else 0, nsh-(sh-dim[1]%sh if dim[1]%sh != 0 else 0)) for dim, sh, nsh in zip(pads, base_shape, new_shape)] - return x.repeat(tuple(repeat_args)).shrink(tuple(shrink_args)) - if mode == "reflect": - for i,s in enumerate(x.shape): - if pads[i] != (0,0): - xL = x.flip(i).shrink(tuple((s-pads[i][0]-1, s_-1) if i_ == i else None for i_,s_ in enumerate(x.shape))) - xR = x.flip(i).shrink(tuple((1, pads[i][1]+1) if i_ == i else None for i_ in range(x.ndim))) - x = xL.cat(x, xR, dim=i) - return x - if mode == "edge": - for i,s in enumerate(x.shape): - if pads[i] != (0,0): - xL = x.shrink(tuple((0,1) if i_ == i else None for i_ in range(x.ndim))).expand([pads[i][0] if i_ == i else None for i_ in range(x.ndim)]) - xR = x.shrink(tuple((s_-1, s_) if i_ == i else None for i_,s_ in enumerate(x.shape))).expand([pads[i][1] if i_ == i else None for i_ in range(x.ndim)]) - x = xL.cat(x, xR, dim=i) - return x - if mode == "constant": - return _padded(x, seq_pads, axes=seq_axes, constant_value=constant_value) +# (x1_begin, x2_begin, ..., x1_end, x2_end, ...) -> (..., x2_start, x2_end, x1_start, x1_end) +def _onnx_pads_to_pad2d_pads(pads): return flatten(reversed(list((pB, pE) for pB, pE in zip(pads, pads[len(pads)//2:])))) +def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Optional[Tensor]=None, axes: Optional[Tensor]=None, mode="constant", value=0): + pads, value, axes = to_python_const(pads), to_python_const(constant_value) or value or 0, to_python_const(axes) or list(range(x.ndim)) + real_pads = [0] * (x.ndim*2) + for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)] + return x.pad(padding=_onnx_pads_to_pad2d_pads(to_python_const(real_pads)), mode={"edge":"replicate", "wrap":"circular"}.get(mode, mode), value=value) def AveragePool(X: Tensor, kernel_shape, auto_pad="NOTSET", ceil_mode=0, count_include_pad=0, dilations=1, pads=None, strides=1): pixel_axes = tuple(range(2, X.ndim)) diff --git a/test/test_ops.py b/test/test_ops.py index 2b8d7e1d68..768f27efc7 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1468,6 +1468,18 @@ class TestOps(unittest.TestCase): # no max pad sizes for replicate helper_test_op([(1,1,5,5)], lambda x: torch.nn.functional.pad(x, (3,11,0,30), mode="replicate"), lambda x: x.pad((3,11,0,30), mode="replicate")) + def test_pad_circular_mode(self): + helper_test_op([(1,1,5,5)], lambda x: torch.nn.functional.pad(x, (0,2,3,2), mode="circular"), lambda x: x.pad((0,2,3,2), mode="circular")) + helper_test_op([(5,5,5)], lambda x: torch.nn.functional.pad(x, (0,2), mode="circular"), lambda x: x.pad((0,2), mode="circular")) + helper_test_op([(1,1,5,5,5)], lambda x: torch.nn.functional.pad(x, (1,2,3,5,1,2),mode="circular"),lambda x:x.pad((1,2,3,5,1,2),mode="circular")) + # circular pad cannot wrap around more than once + self.helper_test_exception([(1,1,5,5)], + lambda x: torch.nn.functional.pad(x, (3,6,0,0), mode="circular"), lambda x: x.pad((3,6,0,0), mode="circular"), + expected=(RuntimeError, ValueError)) + with self.assertRaises(NotImplementedError): + # negative pads with circular pads is not supported + Tensor.randn(1,1,5,5).pad((3,-5,1,-5), mode="circular") + def test_pad_reshape(self): helper_test_op([(1, 2)], lambda x: torch.nn.functional.pad(x, (0, 1, 1, 0)).reshape((3, 2)), diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index d84e6e52c0..d09539d498 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1031,7 +1031,7 @@ class Tensor(SimpleMathTrait): print(t.pad((1, 2, 0, -1), value=-float('inf')).numpy()) ``` """ - if mode not in {"constant", "reflect", "replicate"}: raise NotImplementedError(f"{mode=} is not supported") + if mode not in {"constant", "reflect", "replicate", "circular"}: raise NotImplementedError(f"{mode=} is not supported") if (flat:=all(isinstance(p, (int,UOp)) for p in padding)) and len(padding)%2 != 0: raise ValueError("Flat padding must have even number of pads") # turn flat padding into group padding pX = ((0,0),)*(self.ndim - len(padding)//2) + tuple(zip(padding[-2::-2], padding[::-2])) if flat else padding @@ -1043,6 +1043,11 @@ class Tensor(SimpleMathTrait): return _constant(X, pX, value) if all(resolve(p >= 0) for p in flatten(pX)) else \ _constant(X.shrink(tuple((-smin(pB,0),smin(pA+s,s)) for (pB,pA),s in zip(pX, X.shape))), pads, value) assert all_int(self.shape), f"does not support symbolic shape {self.shape}" + if mode == "circular": + if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, X.shape)): raise ValueError('Padding value causes wrapping around more than once.') + if any(pB<0 or pA<0 for pB,pA in pX): raise NotImplementedError("Negative pads with circular pads is not supported") + orig_shape, X = X.shape, X.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pads)) + return X.shrink(tuple((0 if pB == 0 else osh-pB, xsh if pA == 0 else xsh-osh+pA) for (pB,pA),osh,xsh in zip(pads, orig_shape, X.shape))) 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}.")