Search code examples
pythonparallel-processingjitnumba

How can I set `parallel=True` in numba while I use np.power()?


I'm going to boost my code using numba. However, power function in the parallel mode does not work well i.e. for the function below:

import numpy as np
import numba

@numba.njit(parallel=True, fastmath = True)
def decay_rate(mV, mp):
    drate=(np.power(mp,-3))
    return drate

It says:

The keyword argument 'parallel=True' was specified but no transformation for parallel execution was possible.

The function is more complex than what is demonstrated above (numba overhead worth it!). Also, I've tried the ** for the power before but the result is the same.

How can I fix it?


Solution

  • The message simply states that you specified parallel but there was nothing found to be parallelized. That's because np.power (the same applies to the array-power operator **) is a function defined somewhere else - and you cannot parallelize a function you're calling. However you can parallelize functions you're writing or call functions in parallel.

    If you want to compute the power in parallel, you could do so using numba.prange (see also Explicit Parallel Loops (numba documentation)):

    import numba as nb
    import numpy as np
    
    @nb.njit(parallel=True, fastmath=True)
    def decay_rate(mp):
        drate = np.empty(mp.size)
        for idx in nb.prange(mp.size):
            drate[idx] = np.power(mp[idx], -3)
        return drate
    

    The prange explicitly tells numba to parallelize the loop.