diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index 6a91bdb2..f2e9cd5e 100644 --- a/src/array_api_extra/_delegation.py +++ b/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) diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py index 4e3b8753..d0072bf7 100644 --- a/src/array_api_extra/_lib/_funcs.py +++ b/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] = [] diff --git a/src/array_api_extra/_lib/_utils/_helpers.py b/src/array_api_extra/_lib/_utils/_helpers.py index 6fb7c78b..2a44920b 100644 --- a/src/array_api_extra/_lib/_utils/_helpers.py +++ b/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))