ENH: Adding apply_along_axis

src/array_api_extra/__init__.py

@@ -2,6 +2,7 @@
 
 from . import testing
 from ._delegation import (
+    apply_along_axis,
     argpartition,
     atleast_nd,
     broadcast_shapes,
@@ -41,6 +42,7 @@
 __all__ = [
     "__version__",
     "angle",
+    "apply_along_axis",
     "apply_where",
     "argpartition",
     "at",

src/array_api_extra/_delegation.py

@@ -1,7 +1,7 @@
 """Delegation to existing implementations for Public API Functions."""
 
-from collections.abc import Sequence
-from typing import Literal
+from collections.abc import Callable, Sequence
+from typing import Any, Literal
 
 from ._lib import _funcs
 from ._lib._utils._compat import (
@@ -26,6 +26,7 @@
 from ._lib._utils._typing import Array, ArrayNamespace, Device, DType
 
 __all__ = [
+    "apply_along_axis",
     "argpartition",
     "atleast_nd",
     "broadcast_shapes",
@@ -1732,3 +1733,47 @@ def nanmax(
         return xp.nanmax(a, axis=axis)
 
     return _funcs.nanmax(a, axis=axis, xp=xp)
+
+
+def apply_along_axis(
+    func1d: Callable[..., Array],
+    axis: int,
+    arr: Array,
+    *args: Any,
+    xp: ArrayNamespace | None = None,
+    **kwargs: Any,
+) -> Array:
+    """
+    Apply a function to 1-D slices along a given axis.
+
+    Parameters
+    ----------
+    func1d : callable
+        Function to apply to 1-D slices of `arr`.
+    axis : int
+        Axis along which `func1d` is applied.
+    arr : Array
+        Input array.
+    *args : Any
+        Additional positional arguments to pass to `func1d`.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `arr`. Default: infer.
+    **kwargs : Any
+        Additional keyword arguments to pass to `func1d`.
+
+    Returns
+    -------
+    Array
+        The output array formed by applying `func1d` to each 1-D slice.
+    """
+    if xp is None:
+        xp = array_namespace(arr)
+    if (
+        is_numpy_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_dask_namespace(xp)
+        or is_jax_namespace(xp)
+    ):
+        return xp.apply_along_axis(func1d, axis, arr, *args, **kwargs)
+
+    return _funcs.apply_along_axis(func1d, axis, arr, *args, **kwargs)

src/array_api_extra/_lib/_funcs.py

@@ -4,7 +4,7 @@
 import warnings
 from collections.abc import Callable, Sequence
 from types import NoneType
-from typing import Literal, cast, overload
+from typing import Any, Literal, cast, overload
 
 from ._at import at
 from ._utils import _compat, _helpers
@@ -24,6 +24,7 @@
 
 __all__ = [
     "angle",
+    "apply_along_axis",
     "apply_where",
     "argpartition",
     "atleast_nd",
@@ -233,6 +234,42 @@ def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace) -> Array:
     return x
 
 
+def apply_along_axis(
+    func1d: Callable[..., Array],
+    axis: int,
+    arr: Array,
+    /,
+    *args: Any,
+    **kwargs: Any,
+) -> Array:
+    # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
+    if axis < 0:
+        axis += arr.ndim
+
+    if not 0 <= axis < arr.ndim:
+        msg = f"axis {axis} is out of bounds for array of dimension {arr.ndim}"
+        raise ValueError(msg)
+
+    xp = array_namespace(arr)
+    perm = (*(i for i in range(arr.ndim) if i != axis), axis)
+    arr = xp.permute_dims(arr, perm)
+    leading_shape = arr.shape[:-1]
+    arr = xp.reshape(arr, (-1, arr.shape[-1]))
+    results = [func1d(slice, *args, **kwargs) for slice in xp.unstack(arr, axis=0)]
+
+    if not all(hasattr(result, "dtype") for result in results):
+        results = [xp.asarray(result, dtype=arr.dtype) for result in results]
+
+    out = xp.stack(results, axis=0)
+    out = xp.reshape(out, leading_shape + out.shape[1:])
+    before = tuple(range(axis))
+    after = tuple(range(axis, len(leading_shape)))
+    res = tuple(range(len(leading_shape), out.ndim))
+    perm = before + res + after
+    return xp.permute_dims(out, perm)
+
+
 # `float` in signature to accept `math.nan` for Dask.
 # `int`s are still accepted as `float` is a superclass of `int` in typing
 def broadcast_shapes(  # numpydoc ignore=PR01,RT01

tests/test_funcs.py

@@ -15,6 +15,7 @@
 import array_api_extra._delegation as delegated_func
 from array_api_extra import (
     angle,
+    apply_along_axis,
     apply_where,
     argpartition,
     at,
@@ -107,6 +108,95 @@ def test_all_contains_all_public_functions():
     )
 
 
+class TestApplyAlongAxis:
+    @pytest.mark.xfail_xp_backend(
+        Backend.DASK,
+        reason="Dask apply_along_axis infers output shape differently",
+        strict=False,
+    )
+    def test_simple(self, xp: ArrayNamespace):
+        x = xp.asarray([[1, 2, 3], [4, 5, 6]])
+        actual = apply_along_axis(
+            lambda row: row + 1,
+            axis=1,
+            arr=x,
+        )
+        expected = xp.asarray([[2, 3, 4], [5, 6, 7]])
+        assert_equal(actual, expected)
+
+    @pytest.mark.xfail_xp_backend(
+        Backend.DASK,
+        reason="Dask implementation computes eagerly",
+        strict=False,
+    )
+    @pytest.mark.xfail_xp_backend(
+        Backend.SPARSE,
+        reason="Sparse cannot stack scalar arrays with different fill values",
+        strict=False,
+    )
+    def test_scalar_output(self, xp: ArrayNamespace):
+        x = xp.asarray([[1, 2, 3], [4, 5, 6]])
+        actual = apply_along_axis(
+            xp.sum,
+            axis=1,
+            arr=x,
+        )
+        expected = xp.asarray([6, 15])
+        assert_equal(actual, expected)
+
+    @pytest.mark.xfail_xp_backend(
+        Backend.DASK,
+        reason="Dask implementation computes eagerly",
+        strict=False,
+    )
+    def test_negative_axis(self, xp: ArrayNamespace):
+        x = xp.asarray([[1, 2, 3], [4, 5, 6]])
+        actual = apply_along_axis(
+            lambda row: xp.flip(row, axis=0),
+            axis=-1,
+            arr=x,
+        )
+        expected = xp.asarray([[3, 2, 1], [6, 5, 4]])
+        assert_equal(actual, expected)
+
+    @pytest.mark.xfail_xp_backend(
+        Backend.SPARSE,
+        reason="Sparse cannot stack scalar arrays with different fill values",
+        strict=False,
+    )
+    def test_3d_scalar_output(self, xp: ArrayNamespace):
+        arr_np = np.arange(24).reshape(2, 3, 4)
+        x = xp.asarray(arr_np)
+        expected = np.apply_along_axis(lambda a: np.sum(a), 1, arr_np)  # noqa: PLW0108
+        actual = apply_along_axis(
+            xp.sum,
+            axis=1,
+            arr=x,
+        )
+        assert_equal(actual, xp.asarray(expected))
+
+    def test_axis_out_of_bounds(self, xp: ArrayNamespace):
+        x = xp.asarray([[1, 2, 3], [4, 5, 6]])
+        with pytest.raises(ValueError, match="out of bounds"):
+            _ = functions.apply_along_axis(lambda row: row, 2, x)
+        with pytest.raises(ValueError, match="out of bounds"):
+            _ = functions.apply_along_axis(lambda row: row, -3, x)
+
+    @pytest.mark.xfail_xp_backend(
+        Backend.SPARSE,
+        reason="Sparse cannot stack scalar arrays with different fill values",
+        strict=False,
+    )
+    def test_non_array_output(self, xp: ArrayNamespace):
+        x = xp.asarray([[1, 2, 3], [4, 5, 6]])
+        actual = functions.apply_along_axis(
+            lambda row: int(row[0]) + 1,  # type: ignore[arg-type,return-value]  # pyright: ignore[reportArgumentType]
+            1,
+            x,
+        )
+        assert_equal(actual, xp.asarray([2, 5]))
+
+
 class TestApplyWhere:
     @staticmethod
     def f1(x: Array, y: Array | int = 10) -> Array: