ENH: pad: normalize pad_width on the host for torch

src/array_api_extra/_delegation.py

@@ -15,7 +15,12 @@
     is_torch_namespace,
 )
 from ._lib._utils._compat import device as get_device
-from ._lib._utils._helpers import asarrays, deprecated, eager_shape
+from ._lib._utils._helpers import (
+    asarrays,
+    deprecated,
+    eager_shape,
+    normalize_pad_width,
+)
 from ._lib._utils._typing import Array, DType
 
 __all__ = [
@@ -791,12 +796,17 @@ def pad(
     ):
         return xp.pad(x, pad_width, mode, constant_values=constant_values)
 
-    # https://github.com/pytorch/pytorch/blob/cf76c05b4dc629ac989d1fb8e789d4fac04a095a/torch/_numpy/_funcs_impl.py#L2045-L2056
     if is_torch_namespace(xp):
-        pad_width = xp.asarray(pad_width)
-        pad_width = xp.broadcast_to(pad_width, (x.ndim, 2))
-        pad_width = xp.flip(pad_width, axis=(0,)).flatten()
-        return xp.nn.functional.pad(x, tuple(pad_width), value=constant_values)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
+        # 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)
 

src/array_api_extra/_lib/_funcs.py

@@ -19,6 +19,7 @@
     eager_shape,
     meta_namespace,
     ndindex,
+    normalize_pad_width,
 )
 from ._utils._typing import Array, Device, DType
 
@@ -546,17 +547,7 @@ def pad(
     xp: ModuleType,
 ) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in `array_api_extra._delegation.py`."""
-    # make pad_width a list of length-2 tuples of ints
-    if isinstance(pad_width, int):
-        pad_width_seq = [(pad_width, pad_width)] * x.ndim
-    elif (
-        isinstance(pad_width, tuple)
-        and len(pad_width) == 2
-        and all(isinstance(i, int) for i in pad_width)
-    ):
-        pad_width_seq = [cast(tuple[int, int], pad_width)] * x.ndim
-    else:
-        pad_width_seq = cast(list[tuple[int, int]], list(pad_width))
+    pad_width_seq = normalize_pad_width(pad_width, x.ndim)
 
     slices: list[slice] = []
     newshape: list[int] = []

src/array_api_extra/_lib/_utils/_helpers.py

@@ -18,6 +18,7 @@
     Generic,
     Literal,
     ParamSpec,
+    Sequence,
     TypeAlias,
     TypeVar,
     cast,
@@ -54,6 +55,7 @@ def override(func):
     "eager_shape",
     "in1d",
     "is_python_scalar",
+    "normalize_pad_width",
     "jax_autojit",
     "meta_namespace",
     "pickle_flatten",
@@ -608,3 +610,20 @@ def outer(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
         return inner(wargs).obj
 
     return outer
+
+
+
+def normalize_pad_width(
+    pad_width: int | tuple[int, int] | Sequence[tuple[int, int]],
+    ndim: int,
+) -> list[tuple[int, int]]:  # numpydoc ignore=PR01,RT01
+    """Normalize `pad_width` to a list of `ndim` (before, after) pairs of ints."""
+    if isinstance(pad_width, int):
+        return [(pad_width, pad_width)] * ndim
+    if (
+        isinstance(pad_width, tuple)
+        and len(pad_width) == 2
+        and all(isinstance(i, int) for i in pad_width)
+    ):
+        return [cast(tuple[int, int], pad_width)] * ndim
+    return cast(list[tuple[int, int]], list(pad_width))