ENH: add `apply_along_axis`

[paddyroddy]

Are there plans to support apply_along_axis?

• numpy
• cupy
• torch
• dask
• jax

[lucascolley]

@mdhaber I think I remember some past discussion on this? Maybe I'm getting mixed up with something else.

[mdhaber]

I don't remember specifically - I'd guess it was about how stats functions often rely on apply_along_axis to provide an axis argument or how the resampling tests use apply_along_axis if the user-specified statistic doesn't support an axis argument.
I can imagine this providing a speedup vs looping on some backends (even though it doesn't really with NumPy). If so, I think it would be valuable, and we'd use it.
Something that would make it much more useful is if the arr argument would support a tuple of arrays. Currently, we concatenate arrays along axis and then split each slice within func1d, but that's quite inefficient. (Dealing with higher dimensional core shapes is not as bad because we can typically reshape arr then undo in func1d.)

[rgommers]

This seems useful indeed, and was also discussed in numpy/numpy#30726 (tl;dr it's needed here, nothing to be done in numpy).

There indeed isn't a PyTorch equivalent to call. It was discussed a number of times, e.g. in https://discuss.pytorch.org/t/apply-a-function-along-an-axis/130440. That implementation kinda works, but has some mismatches in signature and things like negative axis handling. A few quick iterations with Deepseek combining the suggestion from the PyTorch forum with the NumPy implementation gives (untested!):

def apply_along_axis(
    func1d: Callable,
    axis: int,
    arr: torch.Tensor,
    *args: Any,
    **kwargs: Any
) -> torch.Tensor:

    # Handle negative axis
    if axis < 0:
        axis = arr.dim() + axis
    
    if axis >= arr.dim():
        raise ValueError(f"axis {axis} is out of bounds for tensor of dimension {arr.dim()}")
    
    # Get slices along the axis
    slices = torch.unbind(arr, dim=axis)
    
    # Apply function to each slice
    results = [func1d(slice, *args, **kwargs) for slice in slices]
    
    # Check if all results are tensors
    if not all(isinstance(r, torch.Tensor) for r in results):
        # Convert scalars to tensors
        results = [r if isinstance(r, torch.Tensor) else torch.tensor(r, dtype=arr.dtype) 
                  for r in results]
    
    # Try to stack results (works if all results have same shape)
    try:
        return torch.stack(results, dim=axis)
    except RuntimeError:
        # Stacking failed - results have different shapes
        # Use concatenation with unsqueeze to add the axis dimension back
        # Check if we need to unsqueeze based on the dimension of results
        return torch.cat([
            r.unsqueeze(axis) if r.dim() == arr.dim() - 1 else r 
            for r in results
        ], dim=axis)

That would need tests and docs, but is probably not too far off. The PyTorch logic can be much simpler than the NumPy one - no need for __array_wrap__/__array_function__ handling (nor __torch_function__ I'd say), and no array scalars either.

[lucascolley]

how close is that to being definable in terms of the standard? Looks pretty close, can we do unbind already?

[rgommers]

Ah good point. unstack is very close to unbind:

• https://data-apis.org/array-api/latest/API_specification/generated/array_api.unstack.html
• https://github.com/data-apis/array-api-compat/blob/d4dae3234048ebc5730f83943910e7c7dae0b33f/array_api_compat/common/_aliases.py#L622

They seem to raise different errors for axis out of bounds, but that should not matter too much.

[lucascolley]

@paddyroddy might you be interested in contributing an array-agnostic version here?

[paddyroddy]

@lucascolley possibly yes. Although I've just paused work on my current project @glass-dev for the next 3 months, so might be a while before I'm back on it

[prady0t]

A few quick iterations with Deepseek combining the suggestion from the PyTorch forum with the NumPy implementation gives (untested!):

I tried testing this implementation (using xp.unstack instead of unbind), and It's slightly incorrect. It works well with some func1d like lambda x: x + 1 but fails for something like func1d = xp.sum The reason for that is that xp.unstack(arr, axis=axis) splits the array into (N-1) dimensional slices by removing the specified axis. However, apply_along_axis is supposed to apply func1d to every 1-D slice obtained by fixing all the other indices and varying only along the chosen axis. As a result, the implementation processes the wrong slices for arrays with more than two dimensions, leading to incorrect output shapes and values.

For example this test fails:

    import numpy as np
    import torch

    arr_np = np.arange(24).reshape(2, 3, 4)
    arr_torch = torch.arange(24).reshape(2, 3, 4)

    expected = np.apply_along_axis(np.sum, axis=1, arr=arr_np)

    actual = apply_along_axis(
        torch.sum,
        axis=1,
        arr=arr_torch,
    )

    np.testing.assert_array_equal(actual.numpy(), expected)

E       AssertionError: 
E       Arrays are not equal
E       
E       (shapes (3,), (2, 4) mismatch)
E        ACTUAL: array([ 60,  92, 124])
E        DESIRED: array([[12, 15, 18, 21],
E              [48, 51, 54, 57]])

As a fix, what I did is:

• move the axis to the end
• flatten all the remaining dimensions so each row is a 1-D slice
• apply the function to every 1-D slice
• stack the results, reshape, and permute the dimensions back to their original order

All while keeping it Array API compatible. See:

def apply_along_axis(
    func1d: Callable[..., Array],
    axis: int,
    arr: Array,
    *args,
    **kwargs
) -> 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)

    # Move the target axis to the end.
    perm = tuple(i for i in range(arr.ndim) if i != axis) + (axis,)
    arr = xp.permute_dims(arr, perm)

    leading_shape = arr.shape[:-1]
    axis_len = arr.shape[-1]

    arr = xp.reshape(arr, (-1, axis_len))
    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)

    # Restore the leading dimensions.
    out = xp.reshape(out, leading_shape + out.shape[1:])
    # Move the result dimensions back to the position of the original axis.

    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)

The above test passes with the changes. This needs some cleaning (chopping out too many variables, blank lines, and comments that AI added), but I guess this logic is a bit more robust.

I'll open a PR to discuss this further.