From 18bd27961e4a7d3ec9b6db905e98a72d28b9a3a5 Mon Sep 17 00:00:00 2001 From: Ralf Gommers Date: Sun, 5 Jul 2026 16:40:05 +0000 Subject: [PATCH 1/4] ENH: pad: normalize `pad_width` on the host for torch `pad_width` is host metadata (an int or (before, after) pairs), but the torch delegation branch normalized it by round-tripping through a tensor: `xp.asarray(pad_width)` places it on the *default* device, and the `tuple(...)` call to build the python ints that `torch.nn.functional.pad` expects then performs one device-to-host transfer per element. With a CUDA default device that is a needless host-to-device copy plus 2*ndim synchronizing `.item()` reads per call; with a data-free `meta` default device it raises outright. Factor the pure-python normalization out of `_funcs.pad` into `normalize_pad_width` and reuse it in the torch branch, reversing the axis order on the host. This also aligns the torch branch with the documented `pad_width` types (the tensor broadcast incidentally accepted undocumented forms such as shape ``(n, 1)``). Co-Authored-By: Claude Fable 5 --- src/array_api_extra/_delegation.py | 11 ++++++----- src/array_api_extra/_lib/_funcs.py | 28 +++++++++++++++++----------- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index 6a91bdb2..5c3294d8 100644 --- a/src/array_api_extra/_delegation.py +++ b/src/array_api_extra/_delegation.py @@ -791,12 +791,13 @@ 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 = _funcs.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..f81e3f6e 100644 --- a/src/array_api_extra/_lib/_funcs.py +++ b/src/array_api_extra/_lib/_funcs.py @@ -538,6 +538,22 @@ def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array: ) +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)) + + def pad( x: Array, pad_width: int | tuple[int, int] | Sequence[tuple[int, int]], @@ -546,17 +562,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] = [] From 57d719cad25939b9a9f829130ae160aaf4bfbb38 Mon Sep 17 00:00:00 2001 From: Ralf Gommers Date: Sun, 5 Jul 2026 20:56:12 +0300 Subject: [PATCH 2/4] Update src/array_api_extra/_delegation.py Co-authored-by: Lucas Colley --- src/array_api_extra/_delegation.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index 5c3294d8..b0f0ff65 100644 --- a/src/array_api_extra/_delegation.py +++ b/src/array_api_extra/_delegation.py @@ -796,7 +796,11 @@ def pad( # `torch/_numpy`'s implementation (avoids device transfers) pad_width_seq = _funcs.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] + 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) From 58404fb09601be0ab602547c1c6e8b9017aaf5b9 Mon Sep 17 00:00:00 2001 From: Ralf Gommers Date: Sun, 5 Jul 2026 21:54:17 +0300 Subject: [PATCH 3/4] MAINT: move `normalize_pad_width` to `_utils/_helpers.py` --- src/array_api_extra/_lib/_funcs.py | 17 +---------------- src/array_api_extra/_lib/_utils/_helpers.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py index f81e3f6e..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 @@ -538,22 +539,6 @@ def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array: ) -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)) - - def pad( x: Array, pad_width: int | tuple[int, int] | Sequence[tuple[int, 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)) From 69e2b70474039ebe11957b6134d2220bc1492cd9 Mon Sep 17 00:00:00 2001 From: Ralf Gommers Date: Sun, 5 Jul 2026 22:09:38 +0300 Subject: [PATCH 4/4] fix issue in previous commit that the linter caught --- src/array_api_extra/_delegation.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index b0f0ff65..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__ = [ @@ -794,7 +799,7 @@ def pad( if 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 = _funcs.normalize_pad_width(pad_width, x.ndim) + pad_width_seq = normalize_pad_width(pad_width, x.ndim) # torch.nn.functional.pad counts dimensions from the last one flat_pad_width = [ w