Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

numba: sorting an array in place

Tags:

jit

numpy

numba

Numba has the amazing ability to speed up loops with JIT compilation. The key twist, however, is that when using numpy, one isn't allowed to create any new arrays. Fortunately, most numpy functions include an optional out parameter for writing output to -- except numpy.sort. The most obvious alternative is numpy.ndarray.sort, which IS in place,

@njit("void(f8[:])")
def sort_inplace(arr):
  arr.sort()

but this fails to compile,

...
...
...
/Users/duckworthd/anaconda/lib/python2.7/site-packages/numba/typeinfer.pyc in propagate(self)
    293                 print("propagate".center(80, '-'))
    294             oldtoken = newtoken
--> 295             self.constrains.propagate(self.context, self.typevars)
    296             newtoken = self.get_state_token()
    297             if config.DEBUG:

/Users/duckworthd/anaconda/lib/python2.7/site-packages/numba/typeinfer.pyc in propagate(self, context, typevars)
    112                 raise
    113             except Exception as e:
--> 114                 raise TypingError("Internal error:\n%s" % e, constrain.loc)
    115
    116

TypingError: Internal error:
Attribute 'sort' of array(float64, 1d, A) is not typed

Short of re-implementing a sorting algorithm, is there a way to sort a numpy array in a JIT-compiled numba loop?

like image 933
duckworthd Avatar asked May 29 '14 05:05

duckworthd


1 Answers

Numba should be able to compile this in "nopython" mode, but unfortunately we haven't added support for ndarray.sort() yet. It compiles in "python" mode though which is slower since it has to go through the python object layer, but since it looks like ndarray.sort() is implemented in C it may not make much of a difference. I went ahead and added a bug report in numba's github issue tracker.

Another thing to note is that numba will compile loops in nopython mode if possible while the rest of the function is compiled in python mode. This allows you to use unsupported functions like ndarray.sort(), and functions like numpy.arange() that create and return a new array, while still having fast compiled loops (as long as you're not calling ndarray.sort() or numpy.arange() inside your loops, but you probably don't want to do that anyway if performance is a big concern).

So to sum up: you may be able to get away with using the jit decorator instead of the njit decorator until numba properly supports calling ndarray.sort() in nopython mode, as long as you're not sorting arrays inside a loop.

like image 193
jayvius Avatar answered Oct 15 '22 03:10

jayvius