fixes

Author: lucascolley

src/array_api_extra/__init__.py

@@ -1,12 +1,12 @@
 """Extra array functions built on top of the array API standard."""
 
+from ._delegators import pad
 from ._lib._funcs import (
     atleast_nd,
     cov,
     create_diagonal,
     expand_dims,
     kron,
-    pad,
     setdiff1d,
     sinc,
 )

src/array_api_extra/_delegators.py

@@ -44,16 +44,18 @@ def pad(
     """
     xp = array_namespace(x) if xp is None else xp
 
-    value = constant_values
+    if mode != "constant":
+        msg = "Only `'constant'` mode is currently supported"
+        raise NotImplementedError(msg)
 
     # 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)
+        return xp.nn.functional.pad(x, (pad_width,), value=constant_values)
 
-    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)
+    if is_numpy_namespace(xp) or is_jax_namespace(xp) or is_cupy_namespace(xp):
+        return xp.pad(x, pad_width, mode, constant_values=constant_values)
 
-    return _funcs.pad(x, pad_width, mode, constant_values=value, xp=xp)
+    return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp)

src/array_api_extra/_lib/_funcs.py

@@ -544,46 +544,14 @@ def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:
 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``.
-    """
-    if mode != "constant":
-        msg = "Only `'constant'` mode is currently supported"
-        raise NotImplementedError(msg)
-
-    value = constant_values
-
-    if xp is None:
-        xp = array_namespace(x)
-
+    xp: ModuleType,
+) -> Array:  # numpydoc ignore=PR01,RT01
+    """See docstring in `_delegators.py`."""
     padded = xp.full(
         tuple(x + 2 * pad_width for x in x.shape),
-        fill_value=value,
+        fill_value=constant_values,
         dtype=x.dtype,
         device=_compat.device(x),
     )

src/array_api_extra/_lib/_utils/_compat.py

@@ -3,7 +3,7 @@
 # `array-api-compat` to override the import location
 
 try:
-    from ..._array_api_compat_vendor import (  # pyright: ignore[reportMissingImports]
+    from ...._array_api_compat_vendor import (  # pyright: ignore[reportMissingImports]
         array_namespace,  # pyright: ignore[reportUnknownVariableType]
         device,  # pyright: ignore[reportUnknownVariableType]
         is_cupy_namespace,  # pyright: ignore[reportUnknownVariableType]