From 1eaeb2a9853f2f48186224def1a0b8663fbc568f Mon Sep 17 00:00:00 2001 From: Dhruv Maniya Date: Wed, 15 Jul 2026 20:19:09 +0530 Subject: [PATCH] ENH: add edge and wrap modes to pad --- src/array_api_extra/_delegation.py | 22 ++++++----- src/array_api_extra/_lib/_funcs.py | 61 ++++++++++++++++++++++++++++-- tests/test_funcs.py | 40 +++++++++++++++++++- 3 files changed, 109 insertions(+), 14 deletions(-) diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index 0edcc7cd..6ae00bc3 100644 --- a/src/array_api_extra/_delegation.py +++ b/src/array_api_extra/_delegation.py @@ -809,7 +809,7 @@ def one_hot( def pad( x: Array, pad_width: int | tuple[int, int] | Sequence[tuple[int, int]], - mode: Literal["constant"] = "constant", + mode: Literal["constant", "edge", "wrap"] = "constant", *, constant_values: complex = 0, xp: ModuleType | None = None, @@ -828,8 +828,9 @@ def pad( A single tuple, ``(before, after)``, is equivalent to a list of ``x.ndim`` copies of this tuple. mode : str, optional - Only "constant" mode is currently supported, which pads with - the value passed to `constant_values`. + Padding mode. "constant" pads with the value passed to + `constant_values`, "edge" pads with the edge values of the array, and + "wrap" pads by wrapping values from the opposite edge. constant_values : python scalar, optional Use this value to pad the input. Default is zero. xp : array_namespace, optional @@ -839,12 +840,12 @@ def pad( ------- array The input array, - padded with ``pad_width`` elements equal to ``constant_values``. + padded according to ``mode``. """ xp = array_namespace(x) if xp is None else xp - if mode != "constant": - msg = "Only `'constant'` mode is currently supported" + if mode not in {"constant", "edge", "wrap"}: + msg = f"Unsupported padding mode {mode!r}" raise NotImplementedError(msg) if ( @@ -853,9 +854,12 @@ def pad( or is_jax_namespace(xp) or is_pydata_sparse_namespace(xp) ): - return xp.pad(x, pad_width, mode, constant_values=constant_values) + if mode == "constant": + return xp.pad(x, pad_width, mode, constant_values=constant_values) + if not is_pydata_sparse_namespace(xp): + return xp.pad(x, pad_width, mode) - if is_torch_namespace(xp): + if mode == "constant" and is_torch_namespace(xp): # normalize `pad_width` on the host rather than through a tensor as done in # `torch/_numpy`'s implementation (avoids device transfers) pad_width_seq = normalize_pad_width(pad_width, x.ndim) @@ -863,7 +867,7 @@ def pad( flat_pad_width = [w for pair in reversed(pad_width_seq) for w in pair] return xp.nn.functional.pad(x, tuple(flat_pad_width), value=constant_values) - return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp) + return _funcs.pad(x, pad_width, mode=mode, constant_values=constant_values, xp=xp) def searchsorted( diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py index 7f28cbd2..8f50f0ed 100644 --- a/src/array_api_extra/_lib/_funcs.py +++ b/src/array_api_extra/_lib/_funcs.py @@ -609,19 +609,74 @@ def pad( x: Array, pad_width: int | tuple[int, int] | Sequence[tuple[int, int]], *, + mode: Literal["constant", "edge", "wrap"] = "constant", constant_values: complex = 0, xp: ModuleType, ) -> Array: # numpydoc ignore=PR01,RT01 """See docstring in `array_api_extra._delegation.py`.""" pad_width_seq = normalize_pad_width(pad_width, x.ndim) - slices: list[slice] = [] - newshape: list[int] = [] - for ax, w_tpl in enumerate(pad_width_seq): + if len(pad_width_seq) != x.ndim: + msg = f"expected {x.ndim} pairs of pad widths, got {len(pad_width_seq)}" + raise ValueError(msg) + + for w_tpl in pad_width_seq: if len(w_tpl) != 2: msg = f"expect a 2-tuple (before, after), got {w_tpl}." raise ValueError(msg) + if w_tpl[0] < 0 or w_tpl[1] < 0: + msg = "index can't contain negative values" + raise ValueError(msg) + if mode != "constant": + for axis, (before, after) in enumerate(pad_width_seq): + if before == 0 and after == 0: + continue + + axis_size = eager_shape(x)[axis] + if axis_size == 0: + msg = f"can't extend empty axis {axis} using mode {mode!r}" + raise ValueError(msg) + + parts: list[Array] = [] + if mode == "edge": + shape = list(eager_shape(x)) + if before: + before_slice = [slice(None)] * x.ndim + before_slice[axis] = slice(0, 1) + shape[axis] = before + parts.append(xp.broadcast_to(x[tuple(before_slice)], tuple(shape))) + + parts.append(x) + + if after: + after_slice = [slice(None)] * x.ndim + after_slice[axis] = slice(-1, None) + shape[axis] = after + parts.append(xp.broadcast_to(x[tuple(after_slice)], tuple(shape))) + else: + before_repeats, before_remainder = divmod(before, axis_size) + after_repeats, after_remainder = divmod(after, axis_size) + + if before_remainder: + before_slice = [slice(None)] * x.ndim + before_slice[axis] = slice(axis_size - before_remainder, None) + parts.append(x[tuple(before_slice)]) + parts.extend([x] * before_repeats) + parts.append(x) + parts.extend([x] * after_repeats) + if after_remainder: + after_slice = [slice(None)] * x.ndim + after_slice[axis] = slice(0, after_remainder) + parts.append(x[tuple(after_slice)]) + + x = xp.concat(parts, axis=axis) + + return x + + slices: list[slice] = [] + newshape: list[int] = [] + for ax, w_tpl in enumerate(pad_width_seq): sh = eager_shape(x)[ax] if w_tpl[0] == 0 and w_tpl[1] == 0: diff --git a/tests/test_funcs.py b/tests/test_funcs.py index c4e05683..cd03ab00 100644 --- a/tests/test_funcs.py +++ b/tests/test_funcs.py @@ -1496,10 +1496,46 @@ def test_ndim(self, xp: ModuleType): padded = pad(a, 2) assert padded.shape == (6, 7, 8) + def test_edge(self, xp: ModuleType): + a = xp.asarray([1, 2, 3]) + padded = pad(a, (2, 1), mode="edge") + assert_equal(padded, xp.asarray([1, 1, 1, 2, 3, 3])) + + def test_edge_ndim(self, xp: ModuleType): + a = xp.asarray([[1, 2], [3, 4]]) + padded = pad(a, ((1, 2), (2, 1)), mode="edge") + expected = xp.asarray( + [ + [1, 1, 1, 2, 2], + [1, 1, 1, 2, 2], + [3, 3, 3, 4, 4], + [3, 3, 3, 4, 4], + [3, 3, 3, 4, 4], + ] + ) + assert_equal(padded, expected) + + def test_wrap(self, xp: ModuleType): + a = xp.asarray([1, 2, 3]) + padded = pad(a, (5, 4), mode="wrap") + assert_equal(padded, xp.asarray([2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1])) + + def test_wrap_ndim(self, xp: ModuleType): + a = xp.asarray([[1, 2], [3, 4]]) + padded = pad(a, ((1, 1), (1, 1)), mode="wrap") + expected = xp.asarray([[4, 3, 4, 3], [2, 1, 2, 1], [4, 3, 4, 3], [2, 1, 2, 1]]) + assert_equal(padded, expected) + + @pytest.mark.parametrize("mode", ["edge", "wrap"]) + def test_empty_axis(self, xp: ModuleType, mode: str): + a = xp.asarray([]) + with pytest.raises(ValueError, match="can't extend empty axis"): + _ = pad(a, 1, mode=mode) # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + def test_mode_not_implemented(self, xp: ModuleType): a = xp.asarray([1, 2, 3]) - with pytest.raises(NotImplementedError, match="Only `'constant'`"): - _ = pad(a, 2, mode="edge") # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + with pytest.raises(NotImplementedError, match="Unsupported padding mode"): + _ = pad(a, 2, mode="reflect") # type: ignore[arg-type] # pyright: ignore[reportArgumentType] def test_device(self, xp: ModuleType, device: Device): a = xp.asarray(0.0, device=device)