From edb76e84a28838fe690a2d06fd32bac944eae614 Mon Sep 17 00:00:00 2001 From: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Date: Tue, 28 Jul 2026 03:09:15 +0530 Subject: [PATCH] Adding apply_along_axis Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> --- src/array_api_extra/__init__.py | 2 + src/array_api_extra/_delegation.py | 47 ++++++++++++++++++- src/array_api_extra/_lib/_funcs.py | 44 +++++++++++++++++- tests/test_funcs.py | 74 ++++++++++++++++++++++++++++++ 4 files changed, 164 insertions(+), 3 deletions(-) diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index 64adcbc8..13610cd5 100644 --- a/src/array_api_extra/__init__.py +++ b/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", diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index c5c16a37..5f03f6b8 100644 --- a/src/array_api_extra/_delegation.py +++ b/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,45 @@ 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. + xp : array_namespace, optional + The standard-compatible namespace for `arr`. Default: infer. + + 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 + ) diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py index 3f1ef511..a26482a3 100644 --- a/src/array_api_extra/_lib/_funcs.py +++ b/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,47 @@ 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: + """Apply a function to 1-D slices along a given axis.""" + 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) + + if kwargs is None: + kwargs = {} + + if not isinstance(args, tuple): + args = (args,) + + xp = array_namespace(arr) + perm = tuple(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 diff --git a/tests/test_funcs.py b/tests/test_funcs.py index 8b639564..dfbdc4b3 100644 --- a/tests/test_funcs.py +++ b/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,79 @@ 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( + np.sum, + 1, + arr_np, + None, + ) + actual = apply_along_axis( + xp.sum, + axis=1, + arr=x, + ) + assert_equal(actual, xp.asarray(expected)) + + class TestApplyWhere: @staticmethod def f1(x: Array, y: Array | int = 10) -> Array: