Creating `NumPy` arrays inside a function decorated with `numba`'s `@jit(nopython=True)`?

问题

I would like to create a numpy array inside a function decorated with numba's @jit(nopython=True). For example:

import numpy as np
import numba

@numba.jit(nopython=True)
def funny_func():
    zero_array = np.zeros(10)

    sum_result = 0

    for elem in zero_array:
        sum_result += elem

    return sum_result

print funny_func()

Compiling this script creates the following error:

UntypedAttributeError: Unknown attribute "zeros" of type Module(<module
'numpy' from 'A:\Anaconda\lib\site-packages\numpy\__init__.pyc'>)

So, numba does't support NumPy array creation functions. How then do I create NumPy arrays inside such a decorated "numba function"?

回答1:

As the docs explicitly say:

NumPy array creation is not supported in nopython mode. Numba mitigates this by automatically trying to jit loops in nopython mode. This allows for array creation at the top of a function while still getting almost all the performance of nopython mode.

(I linked to an older version, because I believe this limitation doesn't exist in 0.18, which implies that you're using an older one. Even older versions, I think before 0.12 or so, don't have this documentation because the auto-lifting didn't exist yet, and you had to do it manually, but the same approach below will work.)

If you're using too old a version of Numba to have that feature, or you've done something that's complicated enough to confuse that feature and make it not work, you have to do the same thing manually. For example:

@numba.jit(nopython=True)
def _funny_func(zero_array):    
    sum_result = 0

    for elem in zero_array:
        sum_result += elem

    return sum_result

@numba.jit(nopython=False)
def funny_func():
    zero_array = np.zeros(10)
    return _funny_func(zero_array)