Support for axis arguments on reduction functions

[seibert]

Now that we support array expressions, I'm finding that I really want to be able to pass the axis argument to reduction functions like np.sum(), np.mean(), etc. This raises the issue again of how to do dispatch based on argument values, but it is becoming much more critical now that I can so easily create arrays to reduce in nopython mode.

[pitrou]

Well as long as we support only a couple signatures (e.g. np.sum(a) and np.sum(a, axis=...)), then it shouldn't be a problem. It's when there's some combinatorial explosion in argument signatures that our current approach doesn't work.

Another question: can specifying the axis change the return type?

[seibert]

If you mean "can the dtype change if axis is specified?", I don't think so. There are still the handful of reduction functions (like np.mean) that can return a different dtype, but that is independent of whether axis is set.

One aspect of axis that I didn't appreciate is that it can take a tuple of integers to sum multiple axes at once (starting with Numpy 1.7). I don't think we need to support that in the first pass implementation of this feature since that usage is still pretty rare.

[jorgehatccrma]

Is there a workaround to pass axis argument to a reduction function like np.sum()? I'm currently getting the following error:

...
Failed at nopython (nopython frontend)
Internal error at <numba.typeinfer.CallConstraint object at 0x127adfad0>:
--%<-----------------------------------------------------------------
Traceback (most recent call last):
  File "/usr/local/lib/python2.7/site-packages/numba/typeinfer.py", line 128, in propagate
    constraint(typeinfer)
  File "/usr/local/lib/python2.7/site-packages/numba/typeinfer.py", line 373, in __call__
    self.resolve(typeinfer, typevars, fnty)
  File "/usr/local/Cellar/python/2.7.12_2/Frameworks/Python.framework/Versions/2.7/lib/python2.7/contextlib.py", line 35, in __exit__
    self.gen.throw(type, value, traceback)
  File "/usr/local/lib/python2.7/site-packages/numba/errors.py", line 249, in new_error_context
    six.reraise(type(newerr), newerr, sys.exc_info()[2])
  File "/usr/local/lib/python2.7/site-packages/numba/errors.py", line 243, in new_error_context
    yield
  File "/usr/local/lib/python2.7/site-packages/numba/typeinfer.py", line 373, in __call__
    self.resolve(typeinfer, typevars, fnty)
  File "/usr/local/lib/python2.7/site-packages/numba/typeinfer.py", line 386, in resolve
    sig = typeinfer.resolve_call(fnty, pos_args, kw_args)
  File "/usr/local/lib/python2.7/site-packages/numba/typeinfer.py", line 1021, in resolve_call
    return self.context.resolve_function_type(fnty, pos_args, kw_args)
  File "/usr/local/lib/python2.7/site-packages/numba/typing/context.py", line 202, in resolve_function_type
    return func.get_call_type(self, args, kws)
  File "/usr/local/lib/python2.7/site-packages/numba/types/functions.py", line 49, in get_call_type
    sig = temp.apply(args, kws)
  File "/usr/local/lib/python2.7/site-packages/numba/typing/templates.py", line 185, in apply
    sig = generic(args, kws)
  File "/usr/local/lib/python2.7/site-packages/numba/typing/npydecl.py", line 352, in generic
    assert not kws
InternalError:
[1] During: resolving callee type: Function(<unbound method Numpy_redirect_sum.sum>)

[fritzo]

Is it feasible to at least raise an error when summing over one axis? I just spent a long time debugging, and enventually found that my_2d_array.sum(0) was being silently interpreted as my_2d_array.sum().

[fritzo]

Re: combinatorial explosion (@pitrou) I often use the keepdims argument, as in my_array.sum(axis=n, keepdims=True). I believe the three important cases are .sum(), .sum(axis=n), .sum(axis=n, keepdims=True), and that there is virtually no need for .sum(keepdims=True) since the result is a scalar and broadcasts correctly.

Re: tuples: It may be easiest to transform

my_ndarray.sum(axis=(2,3,4)) = my_ndarray.sum(axis=2)     \
                                         .sum(axis=3-1)   \
                                         .sum(axis=4-1-1)

my_ndarray.sum(axis=(2,3,4), keepdims=True) = my_ndarray.sum(axis=2, keepdims=True) \
                                                        .sum(axis=3, keepdims=True) \
                                                        .sum(axis=4, keepdims=True)

This should incur very little overhead, since the array sizes are much smaller for the later .sum()s.

Could an owner please label this issue as feature request? Thanks!

[nunocalaim]

Is there any plans on implementing this? For me it would be convenient to have the axis argument for np.mean and np.sum.

[ghost]

I have built the following workaround for 2d arrays:

@nb.njit
def np_apply_along_axis(func1d, axis, arr):
  assert arr.ndim == 2
  assert axis in [0, 1]
  if axis == 0:
    result = np.empty(arr.shape[1])
    for i in range(len(result)):
      result[i] = func1d(arr[:, i])
  else:
    result = np.empty(arr.shape[0])
    for i in range(len(result)):
      result[i] = func1d(arr[i, :])
  return result

@nb.njit
def np_mean(array, axis):
  return np_apply_along_axis(np.mean, axis, array)

@nb.njit
def np_std(array, axis):
  return np_apply_along_axis(np.std, axis, array)

This allows to use np_mean/np_std instead of np.mean/np.std with axis support in numba.

[aldanor]

Since a bunch of reduction functions have been implemented lately (like np.ptp in the latest 0.43 release), would it make sense to enhance them so they actually support axis arguments? A few important examples would be:

• np.mean
• np.min
• np.max
• np.ptp
• (etc)

[greenstick]

Adding to the chorus here – the axis argument for numpy aggregation functions is crucial for so many applications. Numba is already an awesome project; would be great to have this added in

[m-beau]

Hi, any news on this side? The only thing which prevents me from decorating almost all of my utility functions with @njit is numba's lack of support for the axis parameter.

Thank you for all the hard work!

[larstokle]

Hi, I ended up here looking for solutions for error messages using keepdims in sum. To allow for axis and keepdims would be a great feature. 👍

I managed a workaround for keepdims using array.sum(axis=0).reshape((1, -1)) and array.sum(axis=1).reshape((-1, 1)) for 2d arrays. It is not the nicest solution, but numba compiled it at least, and I got some speedup for my function. I guess this should work for other numpy functions as well.

Although I have not tried, I guess this could extend to to nd array with axis=sometuple if one does something like @fritzo explains above, and handles the dimensions to reshape something like

arr_shape = np.array(my_array.shape)
arr_shape[list(sometuple)] = 1
my_arr_summed = my_array.sum(axis=sometuple).reshape(arr_shape)

Hope this can be of some help to someone :)

[7iW]

Inspired by the above comment by @joelrich, here is a way to apply numpy functions (e.g. np.mean or np.std) with axis=0 for arrays of arbitrary dimension.

import numpy as np
import numba as nb


@nb.njit
def apply_along_axis_0(func1d, arr):
    """Like calling func1d(arr, axis=0)"""
    if arr.size == 0:
        raise RuntimeError("Must have arr.size > 0")
    ndim = arr.ndim
    if ndim == 0:
        raise RuntimeError("Must have ndim > 0")
    elif 1 == ndim:
        return func1d(arr)
    else:
        result_shape = arr.shape[1:]
        out = np.empty(result_shape, arr.dtype)
        _apply_along_axis_0(func1d, arr, out)
        return out


@nb.njit
def _apply_along_axis_0(func1d, arr, out):
    """Like calling func1d(arr, axis=0, out=out). Require arr to be 2d or bigger."""
    ndim = arr.ndim
    if ndim < 2:
        raise RuntimeError("_apply_along_axis_0 requires 2d array or bigger")
    elif ndim == 2:  # 2-dimensional case
        for i in range(len(out)):
            out[i] = func1d(arr[:, i])
    else:  # higher dimensional case
        for i, out_slice in enumerate(out):
            _apply_along_axis_0(func1d, arr[:, i], out_slice)


@nb.njit
def nb_mean_axis_0(arr):
    return apply_along_axis_0(np.mean, arr)


@nb.njit
def nb_std_axis_0(arr):
    return apply_along_axis_0(np.std, arr)


@nb.njit
def nb_amax_axis_0(arr):
    return apply_along_axis_0(np.amax, arr)

import random
def test_apply_along_axis_0():
    for numpy_func1d, numba_func1d in [
        (np.mean, nb_mean_axis_0),
        (np.std, nb_std_axis_0),
        (np.amax, nb_amax_axis_0),
    ]:

        for shape_length in range(1, 4):  # number of dimensions
            shape = tuple(random.randrange(1, 4) for _ in range(shape_length))
            arr = np.array(np.random.randn(*shape))  # a random array of given shape
            native_result = numpy_func1d(arr, axis=0)
            our_result = numba_func1d(arr)
            assert np.all(native_result == our_result)

[disadone]

Any news? Is there a plan to add axis option into numba?

[esc]

@disadone thanks for asking about this. This is a long standing feature request which is both a lot of work to implement and also probably not too straightforward. As such, it is still awaiting a champion.

[silgon]

@joelrich comment is really neat. However I got a problem with certain kinds of functions like any, all, etc. These kind of functions have boolean inputs and outputs. So, the workaround is really small, you just need to add dtype=arr.dtype to the empty function as follows:

@nb.njit
def np_apply_along_axis(func1d, axis, arr):
    assert arr.ndim == 2
    assert axis in [0, 1]
    if axis == 0:
        result = np.empty(arr.shape[1], dtype=arr.dtype)
        for i in range(len(result)):
            result[i] = func1d(arr[:, i])
    else:
        result = np.empty(arr.shape[0], dtype=arr.dtype)
        for i in range(len(result)):
            result[i] = func1d(arr[i, :])
    return result

[jordiyapz]

no progress yet?

[FirefoxMetzger]

I am super new to numba but I know my way around numpy. The following is a pure python alternative to axis and keepdims:

def reduce(x:np.ndarray, reduce_op:np.ufunc, *, axis:Union[int, Tuple[int]]=None, keepdims:bool=False) -> np.ndarray:
    if axis is None:
        axis = tuple([x for x in range(x.ndim)])

    if isinstance(axis, int):
        axis = (axis,)

    # tuple to list for indexing semantics
    axis = [*axis]

    reduced_axes = np.zeros(x.ndim, dtype=bool)
    reduced_axes[axis] = True
    
    original_shape = np.array(x.shape)
    new_shape = original_shape[~reduced_axes]

    # this could be reversed, but we are calling a reduction op on it anyway
    new_axes = -np.arange(1, len(axis)+1)
    x = np.moveaxis(x, axis, new_axes)

    total_reduce = np.prod(original_shape[axis])
    total_keep = np.prod(new_shape)
    x = np.reshape(x, (total_keep, total_reduce))

    result = np.empty((total_keep,), dtype=x.dtype)
    for idx in range(result.size):
        result[idx] = reduce_op(x[idx, ...])

    if keepdims:
        new_shape = original_shape
        new_shape[axis] = 1

    return np.reshape(result, new_shape)

with the following test snippet

# import the reduce function
import numpy as np


shape = (1, 3, 4, 5)
x = np.arange(np.prod(shape)).reshape(shape)
keepdims = False
axes = (-1, -2, 1)

result = reduce(x, np.sum, axis=axes, keepdims=keepdims)
expected = np.sum(x, axis=axes, keepdims=keepdims)

assert np.allclose(result, expected)

and should be easy to turn into a numba function ... except that I don't know how to use numba properly, so I'm fighting my way through beginner mistakes at the moment (enforcing consistent types, etc.). I will update this post once I get it working.

Update: So apparently np.movaxis is not supported? If so, the above solution won't work in numba :(.

On a tangent, I am quite interested in the idea of testing/choosing numba as the accelerator for scikit-bot instead of going down the cython or raw C route 🚀. ~90% of my performance-critical code relies on the ability to use axes= and keepdims=, so native support for those kwargs would be highly appreciated ❤️

[FirefoxMetzger]

For future reference, here is a numba implementation of a generic reduction op with support for axis= and with keepdims=True. I didn't manage to get keepdims=False to work for the same reasons I couldn't make np.squeeze work #4074 (comment) (inability to set a dynamic return type + inability to set ndim dynamically).

@numba.jit(nopython=True)
def reduce(
    x: np.ndarray, axis: np.ndarray, keepdims: bool
) -> np.ndarray:
    # replace with your favourite reduction op
    reduce_op = np.sum

    if keepdims is False:
        raise NotImplementedError("Numba can't np.squeeze yet.")

    mask = np.zeros(x.ndim, dtype=np.bool8)
    mask[axis] = True

    original_shape = np.array(x.shape)
    squeezed_shape = original_shape[~mask]

    # this could be reversed, but we are calling a reduction op on it anyway
    new_axes = -np.arange(1, len(axis) + 1)
    
    # not that this will copy if reduction happens along a non-contigous axis
    x_work = np.moveaxis(x, axis, new_axes)
    x_work = np.ascontiguousarray(x_work)

    total_reduce = np.prod(original_shape[axis])
    total_keep = np.prod(squeezed_shape)
    tmp_shape = to_fixed_tuple(np.array((total_keep, total_reduce)), 2)
    x_work = np.reshape(x_work, tmp_shape)

    result = np.empty((total_keep,), dtype=x_work.dtype)
    for idx in range(result.size):
        result[idx] = reduce_op(x_work[idx, ...])

    new_shape = original_shape.copy()
    new_shape[axis] = 1
    new_shape_tuple = to_fixed_tuple(new_shape, x.ndim)
    return np.reshape(result, new_shape_tuple)

This makes use of np.moveaxis of which you can find a draft here: #7369

[Hvass-Labs]

I just wanted to mention that I also need the axis-keyword in np.mean and np.std which is apparently not supported in Numba. I just need to pass a single integer, so I can do the calculations on rows or columns in a matrix. If this has already been implemented for Numba's version of np.sum, then why can't it be made to work for the other functions as well? Thanks!

[esc]

If this has already been implemented for Numba's version of np.sum, then why can't it be made to work for the other functions as well?

Guessing here: probably this is technically feasible and would be a meaningful first step to improving axis support. I would also guess that it may be easy to port the code from sum to mean and std. The code for sum with axis is here:

https://github.com/numba/numba/blob/main/numba/np/arraymath.py#L165-L351

Just in case you want to take a closer look 😉

[Hjorthmedh]

Hello, is there any progress on getting axis support for amax/max?

[kc611]

I have a patch that adds axis support to various reduction functions: #8420

The issue is, due to the complexity/approach the function implementations take a performance hit and it ends up slowing the overall testing framework. (This can be noticed in the runtime of the patch CI tests).

The function used to generate the intermediate indices transposes the array according to axis leading to higher array access times:

numba/numba/np/arraymath.py

Lines 745 to 786 in 2ca22c2

	def build_argmax_or_argmin_with_axis_impl(arr, axis, flatten_impl):
	"""
	Given a function that implements the logic for handling a flattened
	array, return the implementation function.
	"""
	check_is_integer(axis, "axis")
	retty = types.intp
	tuple_buffer = tuple(range(arr.ndim))
	def impl(arr, axis=None):
	if axis < 0:
	axis = arr.ndim + axis
	if axis < 0 or axis >= arr.ndim:
	raise ValueError("axis is out of bounds")
	# Short circuit 1-dimensional arrays:
	if arr.ndim == 1:
	return flatten_impl(arr)
	# Make chosen axis the last axis:
	tmp = tuple_buffer
	for i in range(axis, arr.ndim - 1):
	tmp = tuple_setitem(tmp, i, i + 1)
	transpose_index = tuple_setitem(tmp, arr.ndim - 1, axis)
	transposed_arr = arr.transpose(transpose_index)
	# Flatten along that axis; since we've transposed, we can just get
	# batches off the overall flattened array.
	m = transposed_arr.shape[-1]
	raveled = transposed_arr.ravel()
	assert raveled.size == arr.size
	assert transposed_arr.size % m == 0
	out = np.empty(transposed_arr.size // m, retty)
	for i in range(out.size):
	out[i] = flatten_impl(raveled[i * m:(i + 1) * m])
	# Reshape based on axis we didn't flatten over:
	return out.reshape(transposed_arr.shape[:-1])
	return impl

The other approach to this problem would be to build reduction specific iterators. (Similar to the Indexers that Numba uses within it's fancy indexing framework.)

[chn-lee-yumi]

np.linalg.norm has the same problem too.

[MattGam3]

I have a similar issue with cumprod. I tried to modify the base function as

@nb.njit
def np_apply_along_axis(func1d, axis, arr):
    assert arr.ndim == 2
    # assert axis in [0, 1]
    if axis == 0:
        result = np.empty(arr.shape[1], dtype=arr.dtype)
        for i in nb.prange(len(result)):
            t_arr = np.transpose(arr)
            vec = t_arr[i]
            result[i] = func1d(vec)
            # result[i] = func1d(arr[:, i])
    else:
        result = np.empty(arr.shape[0], dtype=arr.dtype)
        for i in nb.prange(len(result)):
            result[i] = func1d(arr[i])
    return result

@nb.njit
def nb_cumprod(array, axis):
    return np_apply_along_axis(np.cumprod, axis, array)

but still get some errors:

> numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
> Failed in nopython mode pipeline (step: nopython frontend)
> No implementation of function Function(<built-in function setitem>) found for signature:
>  
>  >>> setitem(array(float64, 1d, C), int64, array(float64, 1d, C))
>  
> There are 16 candidate implementations:
>    - Of which 16 did not match due to:
>    Overload of function 'setitem': File: <numerous>: Line N/A.
>      With argument(s): '(array(float64, 1d, C), int64, array(float64, 1d, C))':
>     No match.
> 
> def np_apply_along_axis(func1d, axis, arr):
>     <source elided>
>             vec = t_arr[i]
>             result[i] = func1d(vec)
>             ^
> 
> def nb_cumprod(array, axis):
>     return np_apply_along_axis(np.cumprod, axis, array)
>     ^

Do you have a solution / idea why this is happening?

One solution would be

@nb.njit
def nb_cumprod_axis(arr, axis):
    if axis == 0:
        result = np.copy(arr)
        for i in range(1, result.shape[0]):
            result[i] *= result[i - 1]
        return result
    elif axis == 1:
        result = np.copy(arr)
        for i in range(1, result.shape[1]):
            result[:, i] *= result[:, i - 1]
        return result
    else:
        raise ValueError("Invalid axis value. Axis must be 0 or 1.")

but it would not fit in the above framework.

[chulwoohan]

While I was looking for a solution to the same problem and appreciated previous suggestions, I soon realized that most calculations can be decomposed into additions and multiplications. Using this fact, basic stat functions can be made from np.sum(). First, the mean function:

def mean(x, axis=None):
    if axis == None:
        return np.sum(x, axis) / np.prod(x.shape)
    else:
        return np.sum(x, axis) / x.shape[axis]

Variance function can be defined using the formula: var = mean(x**2) - mean(x)**2:

def var(x, axis=None):
    return mean(x**2, axis) - mean(x, axis)**2

The standard deviation:

def std(x, axis=None):
    return np.sqrt(var(x, axis)

Note: A tuple axis is not supported.

[jorenham]

@lucasjinreal Nobody is stopping you 🤷🏻

[pitrou]

@lucasjinreal Please stop this.

[jorenham]

@lucasjinreal I meant no insult; I was just trying to keep it constructive. And I'm afraid that I can't answer your questions, since the presumptions are false.

[chulwoohan]

One hundred years later, still no such a trivial contribution or commits to this issue? mean, max, min, sum ....

None of them foundamental supported in numba.......

Can't you follow my suggestion?

[fingoldo]

Hi, just stumbled upon this problem, too. Can we hope to have this implemented by at least 2030? )

[esc]

Hi, just stumbled upon this problem, too. Can we hope to have this implemented by at least 2030? )

Pull-requests welcome! 😈

[tianlinzx]

Pull-requests welcome! 😈