Does numba not support np.argsort for 2d arrays?

Numba
1

Numba reports an error when running the np.argsort function of a 2-dimensional array. The 1d array runs normally. Is it not supported?

from numba import njit
from numba.types import unicode_type
import numpy as np

@njit
def foo(arr):
    return np.argsort(arr)

foo(np.random.randn(3, 4).astype('float32'))
2

The np.argsort function also only applies to a single dimension (axis), even with Numpy. The default of Numpy is to apply it to the last axis when unspecified, or on the flattened array with axis=None.

So you could implement that with Numba for example using the guvectorize decorator, which also adds the axis keyword with the same default as Numpy (axis=-1).

For example:

from numba import njit, guvectorize
import numpy as np

@guvectorize("void(float32[:], int64[:])", "(n)->(n)")
def foo(arr, out):
    out[:] = np.argsort(arr)

That should match Numpy, for example sorting over the default last axis:

image