Add axis support to np.take

[guilhermeleobas]

Non-scientific script to measure time:

import time
from numba import njit
import numpy as np

@njit
def foo(a, indices, axis):
    return np.take(a, indices, axis)


def measure_time(fn, *args):
    start = time.time()
    fn(*args)
    end = time.time()
    print(f'{fn.__name__}: {end - start}s')



# N = 100
N = 10 ** 9
shape = (10,) * int(np.log10(N))
a = np.arange(N).reshape(shape)
indices = (1, 0, 2)
axis=0
foo(a, indices, axis)
measure_time(foo.py_func, a, indices, axis)
measure_time(foo, a, indices, axis)

[guilhermeleobas]

This could be rewritten as:

def _get_take_getitem(ndim):
    @register_jitable
    def _take_getitem(a, idx, axis):
        if axis == 0:
            return a[idx, ...]
        elif ndim >= 2 and axis == 1:
            return a[:, idx, ...]
        elif ndim >= 3 and axis == 2:
            return a[:, :, idx, ...]
        elif ndim >= 4 and axis == 3:
            return a[:, :, :, idx, ...]
        elif ndim >= 5 and axis == 4:
            return a[:, :, :, :, idx, ...]
        elif ndim >= 6 and axis == 2:
            return a[:,:,:,:,:, idx, ...]
        elif ndim >= 7 and axis == 2:
            return a[:,:,:,:,:,:, idx, ...]
        elif ndim >= 8 and axis == 2:
            return a[:,:,:,:,:,:,:, idx, ...]
        elif ndim >= 9 and axis == 2:
            return a[:, :, :, :, :, :, :, idx, ...]
        elif ndim >= 11 and axis == 10:
            ...
        elif ndim >= 12 and axis == 10:
        elif ndim >= 13 and axis == 10:
        elif ndim >= 14 and axis == 10:
        elif ndim >= 15 and axis == 10:
        elif ndim >= 16 and axis == 10:
        elif ndim >= 17 and axis == 10:
        elif ndim >= 18 and axis == 10:
        elif ndim >= 19 and axis == 10:
        elif ndim >= 20 and axis == 20:
        elif ndim >= 21 and axis == 20:
        elif ndim >= 22 and axis == 20:
        elif ndim >= 23 and axis == 20:
        elif ndim >= 24 and axis == 20:
        elif ndim >= 25 and axis == 20:
        elif ndim >= 26 and axis == 20:
        elif ndim >= 27 and axis == 20:
        elif ndim >= 28 and axis == 20:
        elif ndim >= 29 and axis == 20:
        elif ndim >= 30 and axis == 30:
        elif ndim >= 31 and axis == 31:
        elif ndim >= 32 and axis == 32:
    return _take_getitem

Dead branch prune would take care to remove branches that are dead.

[kc611]

The string itself should not contain a register_jitable, the output function of this string should be jitted instead.

[kc611]

The string itself should not contain a register_jitable, the output function of this string should be jitted instead.

[kc611]

The squeeze implementation needs to exclude the axis based on the axis provided not the value 1:

For instance:

import time
from numba import njit
import numpy as np

@njit
def foo(a, indices, axis):
    return np.take(a, indices, axis)


def measure_time(fn, *args):
    start = time.time()
    res = fn(*args)
    end = time.time()
    print(f'{fn.__name__}: {end - start}s')
    print(res.shape)
    print("ndim:", res.ndim)



shape = (10, 1, 11, 1, 12, 1, 13)
a = np.ones(shape)
indices = 0
axis=1
measure_time(foo.py_func, a, indices, axis)
# foo: 0.00011563301086425781s
# (10, 11, 1, 12, 1, 13)
# ndim: 6
measure_time(foo, a, indices, axis) # fails

shape = (10, 1, 11, 1, 12, 1, 13)
a = np.ones(shape)
indices = 0
axis = 2
measure_time(foo.py_func, a, indices, axis)
# foo: 6.4373016357421875e-06s
# (10, 1, 1, 12, 1, 13)
# ndim: 6
measure_time(foo, a, indices, axis) # fails

[guilhermeleobas]

Good catch. I'll fix this

[kc611]

This is potentially very wasteful. If we're anyways generating the getitem and setitem specifically for this function we may want to combine them into a single function that copies the elements directly in the out resultant. Is that something you've tried out ?

[guilhermeleobas]

I've split getitem from setitem because the latter is used in #9313

[guilhermeleobas]

CI failure is not related to this PR.

[gmarkall]

I've given the CI a prod.