ENH: add edge and wrap modes to pad

Author: iamdhrv

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,17 +854,20 @@ 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)
         # torch.nn.functional.pad counts dimensions from the last one
         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(

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:

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)