ENH: allow tuple pad_width in pad

Author: ev-br

src/array_api_extra/_funcs.py

@@ -555,7 +555,7 @@ def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:
 
 def pad(
     x: Array,
-    pad_width: int,
+    pad_width: int | tuple,
     mode: str = "constant",
     *,
     xp: ModuleType | None = None,
@@ -568,8 +568,9 @@ def pad(
     ----------
     x : array
         Input array.
-    pad_width : int
+    pad_width : int or tuple of ints
         Pad the input array with this many elements from each side.
+        Ifa tuple, each element applies to the corresponding axis of `x`.
     mode : str, optional
         Only "constant" mode is currently supported, which pads with
         the value passed to `constant_values`.
@@ -590,16 +591,23 @@ def pad(
 
     value = constant_values
 
+    if isinstance(pad_width, int):
+        pad_width = (pad_width,) * x.ndim
+
     if xp is None:
         xp = array_namespace(x)
 
     padded = xp.full(
-        tuple(x + 2 * pad_width for x in x.shape),
+        tuple(x + 2 * w for (x, w) in zip(x.shape, pad_width)),
         fill_value=value,
         dtype=x.dtype,
         device=_compat.device(x),
     )
-    padded[(slice(pad_width, -pad_width, None),) * x.ndim] = x
+    sl = tuple(
+        slice(w, -w, None) if w > 0 else slice(None, None, None)
+        for w in pad_width
+    )
+    padded[sl] = x
     return padded
 
 

tests/test_funcs.py

@@ -416,3 +416,11 @@ def test_device(self):
 
     def test_xp(self):
         assert_array_equal(pad(xp.asarray(0), 1, xp=xp), xp.zeros(3))
+
+    def test_tuple_width(self):
+        a = xp.reshape(xp.arange(12), (3, 4))
+        padded = pad(a, (1, 0))
+        assert padded.shape == (5, 4)
+
+        padded = pad(a, (1, 2))
+        assert padded.shape == (5, 8)