Repository: numba/numba

Issue:
add np.take (and corresponding index array) support
This should not be too hard to add in the 1D case, and higher dimensional cases may not be too much worse.



Discussion:
(Request from librosa developer)
A bit more context on this one:

I did a quick-skim audit of librosa, and one of the major hurdles to adopting numba is read-only access to subsets of an array when applying pooling operators.  (Think of replacing a column of a matrix by the element-wise average of its neighboring columns.)  This doesn't require full fancy indexing, since we can get a read-only temporary array from `np.take`.

Similarly, this might help unblock some of the issues with implementing sparse matrix ops, since `np.take` can be used to access the underlying data buffer of a compressed sparse matrix.
Do you need the `axis` or `mode` arguments to `np.take`?
I don't personally, no.  I can usually get around the `axis` dependency by wrapping `swapaxes(0, axis)` around the input and output of the function in question.

`mode` hasn't come up for me, largely because my `take` arguments tend to be valid indices (no negatives), so I never need to wrap or clip.  But YMMV.

Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -3347,6 +3347,64 @@ def diag_impl(arr, k=0):
 
     res = context.compile_internal(builder, diag_impl, sig, args)
     return impl_ret_new_ref(context, builder, sig.return_type, res)
+ 
+@lower_builtin(np.take, types.Array, types.Integer)
+@lower_builtin('array.take', types.Array, types.Integer)
+def numpy_take_1(context, builder, sig, args):
+
+    def take_impl(a, indices):
+        if indices > (a.size - 1) or indices < -a.size:
+            raise IndexError("Index out of bounds")
+        return a.ravel()[np.int(indices)]
+
+    res = context.compile_internal(builder, take_impl, sig, args)
+    return impl_ret_new_ref(context, builder, sig.return_type, res)
+
+@lower_builtin('array.take', types.Array, types.Array)
+@lower_builtin(np.take, types.Array, types.Array)
+def numpy_take_2(context, builder, sig, args):
+
+    F_order = sig.args[1].layout == 'F'
+
+    def take_impl(a, indices):
+        ret = np.empty(indices.size, dtype=a.dtype)
+        if F_order:
+            walker = indices.copy() # get C order
+        else:
+            walker = indices
+        it = np.nditer(walker)
+        i = 0
+        flat = a.ravel()
+        for x in it:
+            if x > (a.size - 1) or x < -a.size:
+                raise IndexError("Index out of bounds")
+            ret[i] = flat[x]
+            i = i + 1
+        return ret.reshape(indices.shape)
+
+    res = context.compile_internal(builder, take_impl, sig, args)
+    return impl_ret_new_ref(context, builder, sig.return_type, res)
+
+@lower_builtin('array.take', types.Array, types.List)
+@lower_builtin(np.take, types.Array, types.List)
+def numpy_take_3(context, builder, sig, args):
+
+    def take_impl(a, indices):
+        convert = np.array(indices)
+        ret = np.empty(convert.size, dtype=a.dtype)
+        it = np.nditer(convert)
+        i = 0
+        flat = a.ravel()
+        for x in it:
+            if x > (a.size - 1) or x < -a.size:
+                raise IndexError("Index out of bounds")
+            ret[i] = flat[x]
+            i = i + 1
+        return ret.reshape(convert.shape)
+    
+    res = context.compile_internal(builder, take_impl, sig, args)
+    return impl_ret_new_ref(context, builder, sig.return_type, res)
+
 
 @lower_builtin(np.arange, types.Number)
 def numpy_arange_1(context, builder, sig, args):
diff --git a/numba/typing/arraydecl.py b/numba/typing/arraydecl.py
--- a/numba/typing/arraydecl.py
+++ b/numba/typing/arraydecl.py
@@ -426,6 +426,18 @@ def resolve_flatten(self, ary, args, kws):
         assert not args
         return signature(ary.copy(ndim=1, layout='C'))
 
+    @bound_function("array.take")
+    def resolve_take(self, ary, args, kws):
+        assert not kws
+        argty, = args
+        if isinstance(argty, types.Integer):
+            sig = signature(ary.dtype, *args)
+        elif isinstance(argty, types.Array):
+            sig = signature(argty.copy(layout='C', dtype=ary.dtype), *args)
+        elif isinstance(argty, types.List):
+            sig = signature(types.Array(ary.dtype, 1, 'C'), *args)
+        return sig
+
     def generic_resolve(self, ary, attr):
         # Resolution of other attributes, for record arrays
         if isinstance(ary.dtype, types.Record):
diff --git a/numba/typing/npydecl.py b/numba/typing/npydecl.py
--- a/numba/typing/npydecl.py
+++ b/numba/typing/npydecl.py
@@ -1173,6 +1173,22 @@ def typer(ref, k=0):
         return typer
 
 
+@infer_global(np.take)
+class Take(AbstractTemplate):
+
+    def generic(self, args, kws):
+        assert not kws
+        assert len(args) == 2
+        arr, ind = args
+        if isinstance(ind, types.Number):
+            retty = arr.dtype
+        elif isinstance(ind, types.Array):
+            retty = types.Array(ndim=ind.ndim, dtype=arr.dtype, layout='C')
+        elif isinstance(ind, types.List):
+            retty = types.Array(ndim=1, dtype=arr.dtype, layout='C')
+
+        return signature(retty, *args)
+
 # -----------------------------------------------------------------------------
 # Numba helpers
 


Test patch:
diff --git a/numba/tests/test_array_methods.py b/numba/tests/test_array_methods.py
--- a/numba/tests/test_array_methods.py
+++ b/numba/tests/test_array_methods.py
@@ -1,6 +1,6 @@
 from __future__ import division
 
-from itertools import product
+from itertools import product, cycle
 import sys
 
 import numpy as np
@@ -76,6 +76,13 @@ def np_asfortranarray(arr):
 def array_view(arr, newtype):
     return arr.view(newtype)
 
+def array_take(arr, indices):
+    return arr.take(indices)
+
+def array_take_kws(arr, indices, axis):
+    return arr.take(indices, axis=axis)
+
+
 # XXX Can't pass a dtype as a Dispatcher argument for now
 def make_array_view(newtype):
     def array_view(arr):
@@ -646,6 +653,68 @@ def test_cumsum(self):
             cfunc(a, axis=1)
 
 
+    def test_take(self):
+        pyfunc = array_take
+        cfunc = jit(nopython=True)(pyfunc)
+
+        def check(arr, ind):
+            expected = pyfunc(arr, ind)
+            got = cfunc(arr, ind)
+            self.assertPreciseEqual(expected, got)
+            if hasattr(expected, 'order'):
+                self.assertEqual(expected.order == got.order)
+
+        # need to check:
+        # 1. scalar index
+        # 2. 1d array index
+        # 3. nd array index, >2d and F order
+        # 4. reflected list
+        
+        test_indices = []
+        test_indices.append(1)
+        test_indices.append(5)
+        test_indices.append(11)
+        test_indices.append(-2)
+        test_indices.append(np.array([1, 5, 1, 11, 3]))
+        test_indices.append(np.array([[1, 5, 1], [11, 3, 0]], order='F'))
+        test_indices.append(np.array([[[1, 5, 1], [11, 3, 0]]]))
+        test_indices.append(np.array([[[[1, 5]], [[11, 0]],[[1, 2]]]]))
+        test_indices.append([1, 5, 1, 11, 3])
+    
+        layouts = cycle(['C', 'F', 'A'])
+
+        for dt in [np.float64, np.int64, np.complex128]:
+            A = np.arange(12, dtype=dt).reshape((4, 3), order=next(layouts))
+            for ind in test_indices:
+                check(A, ind)
+
+        #check illegal access raises
+        A = np.arange(12, dtype=dt).reshape((4, 3), order=next(layouts))
+        szA = A.size
+        illegal_indices = [szA, -szA - 1, np.array(szA), np.array(-szA - 1),
+                           [szA], [-szA - 1]]
+        for x in illegal_indices:
+            with self.assertRaises(IndexError):
+                cfunc(A, x) # oob raises
+
+        # check float indexing raises
+        with self.assertRaises(TypingError):
+            cfunc(A, [1.7])
+       
+        # check unsupported arg raises
+        with self.assertRaises(TypingError):
+            take_kws = jit(nopython=True)(array_take_kws)
+            take_kws(A, 1, 1)
+
+        # check kwarg unsupported raises
+        with self.assertRaises(TypingError):
+            take_kws = jit(nopython=True)(array_take_kws)
+            take_kws(A, 1, axis=1)
+
+        #exceptions leak refs
+        self.disable_leak_check()
+
+
 class TestArrayComparisons(TestCase):
 
     def test_identity(self):
diff --git a/numba/tests/test_fancy_indexing.py b/numba/tests/test_fancy_indexing.py
--- a/numba/tests/test_fancy_indexing.py
+++ b/numba/tests/test_fancy_indexing.py
@@ -6,6 +6,7 @@
 
 import numba.unittest_support as unittest
 from numba import types, jit, typeof
+from numba.errors import TypingError
 from .support import MemoryLeakMixin, TestCase, tag
 
 
@@ -15,6 +16,11 @@ def getitem_usecase(a, b):
 def setitem_usecase(a, idx, b):
     a[idx] = b
 
+def np_take(A, indices):
+    return np.take(A, indices)
+
+def np_take_kws(A, indices, axis):
+    return np.take(A, indices, axis=axis)
 
 class TestFancyIndexing(MemoryLeakMixin, TestCase):
 
@@ -163,5 +169,60 @@ def test_setitem_array(self):
         self.check_setitem_indices(arr, indices)
 
 
+    def test_np_take(self):
+        # shorter version of array.take test in test_array_methods
+        pyfunc = np_take
+        cfunc = jit(nopython=True)(pyfunc)
+
+        def check(arr, ind):
+            expected = pyfunc(arr, ind)
+            got = cfunc(arr, ind)
+            self.assertPreciseEqual(expected, got)
+            if hasattr(expected, 'order'):
+                self.assertEqual(expected.order == got.order)
+
+        # need to check:
+        # 1. scalar index
+        # 2. 1d array index
+        # 3. nd array index
+        # 4. reflected list
+        
+        test_indices = []
+        test_indices.append(1)
+        test_indices.append(np.array([1, 5, 1, 11, 3]))
+        test_indices.append(np.array([[[1], [5]], [[1], [11]]]))
+        test_indices.append([1, 5, 1, 11, 3])
+    
+        for dt in [np.int64, np.complex128]:
+            A = np.arange(12, dtype=dt).reshape((4, 3))
+            for ind in test_indices:
+                check(A, ind)
+
+        #check illegal access raises
+        szA = A.size
+        illegal_indices = [szA, -szA - 1, np.array(szA), np.array(-szA - 1),
+                           [szA], [-szA - 1]]
+        for x in illegal_indices:
+            with self.assertRaises(IndexError):
+                cfunc(A, x) # oob raises
+
+        # check float indexing raises
+        with self.assertRaises(TypingError):
+            cfunc(A, [1.7])
+       
+        # check unsupported arg raises
+        with self.assertRaises(TypingError):
+            take_kws = jit(nopython=True)(np_take_kws)
+            take_kws(A, 1, 1)
+
+        # check kwarg unsupported raises
+        with self.assertRaises(TypingError):
+            take_kws = jit(nopython=True)(np_take_kws)
+            take_kws(A, 1, axis=1)
+
+        #exceptions leak refs
+        self.disable_leak_check()
+
+
 if __name__ == '__main__':
     unittest.main()