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?
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.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With