ENH: pad: add delegation

Author: lucascolley

docs/api-reference.md

@@ -11,6 +11,7 @@
     create_diagonal
     expand_dims
     kron
+    pad
     setdiff1d
     sinc
 ```

src/array_api_extra/__init__.py

@@ -1,6 +1,6 @@
 """Extra array functions built on top of the array API standard."""
 
-from ._funcs import (
+from ._lib._funcs import (
     atleast_nd,
     cov,
     create_diagonal,

src/array_api_extra/_delegators.py

@@ -0,0 +1,59 @@
+"""Delegators to existing implementations for Public API Functions."""
+
+from ._lib import _funcs
+from ._lib._utils._compat import (
+    array_namespace,
+    is_cupy_namespace,
+    is_jax_namespace,
+    is_numpy_namespace,
+    is_torch_namespace,
+)
+from ._lib._utils._typing import Array, ModuleType
+
+
+def pad(
+    x: Array,
+    pad_width: int,
+    mode: str = "constant",
+    *,
+    constant_values: bool | int | float | complex = 0,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Pad the input array.
+
+    Parameters
+    ----------
+    x : array
+        Input array.
+    pad_width : int
+        Pad the input array with this many elements from each side.
+    mode : str, optional
+        Only "constant" mode is currently supported, which pads with
+        the value passed to `constant_values`.
+    constant_values : python scalar, optional
+        Use this value to pad the input. Default is zero.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        The input array,
+        padded with ``pad_width`` elements equal to ``constant_values``.
+    """
+    xp = array_namespace(x) if xp is None else xp
+
+    value = 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, (pad_width,), value=value)
+
+    if is_numpy_namespace(x) or is_jax_namespace(xp) or is_cupy_namespace(xp):
+        return xp.pad(x, pad_width, mode, constant_values=value)
+
+    return _funcs.pad(x, pad_width, mode, constant_values=value, xp=xp)

src/array_api_extra/_lib/__init__.py

@@ -1 +1 @@
-"""Modules housing private functions."""
+"""Array-agnostic implementations for the public API."""

src/array_api_extra/_lib/_compat.py

@@ -2,9 +2,9 @@
 
 import warnings
 
-from ._lib import _compat, _utils
-from ._lib._compat import array_namespace
-from ._lib._typing import Array, ModuleType
+from ._utils import _compat, _helpers
+from ._utils._compat import array_namespace
+from ._utils._typing import Array, ModuleType
 
 __all__ = [
     "atleast_nd",
@@ -136,7 +136,7 @@ def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
     m = atleast_nd(m, ndim=2, xp=xp)
     m = xp.astype(m, dtype)
 
-    avg = _utils.mean(m, axis=1, xp=xp)
+    avg = _helpers.mean(m, axis=1, xp=xp)
     fact = m.shape[1] - 1
 
     if fact <= 0:
@@ -449,7 +449,7 @@ def setdiff1d(
     else:
         x1 = xp.unique_values(x1)
         x2 = xp.unique_values(x2)
-    return x1[_utils.in1d(x1, x2, assume_unique=True, invert=True, xp=xp)]
+    return x1[_helpers.in1d(x1, x2, assume_unique=True, invert=True, xp=xp)]
 
 
 def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:
@@ -546,8 +546,8 @@ def pad(
     pad_width: int,
     mode: str = "constant",
     *,
-    xp: ModuleType | None = None,
     constant_values: bool | int | float | complex = 0,
+    xp: ModuleType | None = None,
 ) -> Array:
     """
     Pad the input array.
@@ -561,10 +561,10 @@ def pad(
     mode : str, optional
         Only "constant" mode is currently supported, which pads with
         the value passed to `constant_values`.
-    xp : array_namespace, optional
-        The standard-compatible namespace for `x`. Default: infer.
     constant_values : python scalar, optional
         Use this value to pad the input. Default is zero.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
 
     Returns
     -------

src/array_api_extra/_funcs.py renamed to src/array_api_extra/_lib/_funcs.py

@@ -0,0 +1 @@
+"""Modules housing private utility functions."""

src/array_api_extra/_lib/_utils/__init__.py

@@ -0,0 +1,31 @@
+"""Acquire helpers from array-api-compat."""
+# Allow packages that vendor both `array-api-extra` and
+# `array-api-compat` to override the import location
+
+try:
+    from ..._array_api_compat_vendor import (  # pyright: ignore[reportMissingImports]
+        array_namespace,  # pyright: ignore[reportUnknownVariableType]
+        device,  # pyright: ignore[reportUnknownVariableType]
+        is_cupy_namespace,  # pyright: ignore[reportUnknownVariableType]
+        is_jax_namespace,  # pyright: ignore[reportUnknownVariableType]
+        is_numpy_namespace,  # pyright: ignore[reportUnknownVariableType]
+        is_torch_namespace,  # pyright: ignore[reportUnknownVariableType]
+    )
+except ImportError:
+    from array_api_compat import (  # pyright: ignore[reportMissingTypeStubs]
+        array_namespace,  # pyright: ignore[reportUnknownVariableType]
+        device,
+        is_cupy_namespace,  # pyright: ignore[reportUnknownVariableType]
+        is_jax_namespace,  # pyright: ignore[reportUnknownVariableType]
+        is_numpy_namespace,  # pyright: ignore[reportUnknownVariableType]
+        is_torch_namespace,  # pyright: ignore[reportUnknownVariableType]
+    )
+
+__all__ = [
+    "array_namespace",
+    "device",
+    "is_cupy_namespace",
+    "is_jax_namespace",
+    "is_numpy_namespace",
+    "is_torch_namespace",
+]

src/array_api_extra/_lib/_utils/_compat.py

@@ -15,3 +15,7 @@ def array_namespace(
     use_compat: bool | None = None,
 ) -> ArrayModule: ...  # numpydoc ignore=GL08
 def device(x: Array, /) -> Device: ...  # numpydoc ignore=GL08
+def is_cupy_namespace(xp: ModuleType, /) -> bool: ...
+def is_jax_namespace(xp: ModuleType, /) -> bool: ...
+def is_numpy_namespace(xp: ModuleType, /) -> bool: ...
+def is_torch_namespace(xp: ModuleType, /) -> bool: ...

src/array_api_extra/_lib/_compat.pyi renamed to src/array_api_extra/_lib/_utils/_compat.pyi

@@ -1,4 +1,4 @@
-"""Utility functions used by `array_api_extra/_funcs.py`."""
+"""Helper functions used by `array_api_extra/_funcs.py`."""
 
 from . import _compat
 from ._typing import Array, ModuleType