Numba：@jit 中的并行标志在我使用 numpy 二维数组的代码中不起作用

Question

首先，我问了这个问题Numba : Why guvectorize is so slow?。当我发现 numba 时尝试了很多东西，出现了这个更基本的问题

为什么此代码有效（并行标志设置为 False）：

import numpy as np
from numba import jit, prange

@jit(('float64[:,:], float64[:,:], float64[:]'), nopython=True, parallel=False)
def myFoo(result,par,x):
    lenght=len(x)
    for i in prange(lenght):
        result[:,i:i+1] = (par[:,0:1]*np.exp(par[:,1:2]*x[i]))
         
x=np.array(np.arange(0,20,1.0))
par=np.array([[i*10,i*0.1] for i in range(10)])
result = np.empty([par.shape[0],x.shape[0]], dtype=np.float64)
myFoo(result, par, x)

但是当我将 parallel 标志更改为 True 时，它不起作用。错误信息是：

TypingError: Failed in nopython mode pipeline (step: nopython mode backend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<intrinsic assert_equiv>) found for signature:
 
 >>> <unknown function>(Literal[str](Sizes of result, binary_multiply.31 do not match on /tmp/ipykernel_64043/4263189787.py (13)), float64, int64)
 
There are 2 candidate implementations:
    - Of which 2 did not match due to:
    Intrinsic in function 'assert_equiv': File: numba/parfors/array_analysis.py: Line 127.
      With argument(s): '(unicode_type, float64, int64)':
     Rejected as the implementation raised a specific error:
       AssertionError: 
  raised from ... lib/python3.9/site-packages/numba/parfors/array_analysis.py:143

Answer 1

不确定在您的项目中是否有必要在函数范围之外分配结果数组并将其作为参数传递给函数。否则此代码修改运行得很好：

import numpy as np
from numba import jit, prange

@jit(('float64[:,:], float64[:]'), nopython=True, parallel=True)
def myFoo(par,x):
    length=x.shape[0]
    result = np.empty((par.shape[0],x.shape[0]), dtype=np.float64)               
    for i in prange(length):
        result[:,i] = (par[:,0]*np.exp(par[:,1]*x[i]))
    return result

    
x=np.array(np.arange(0,20,1.0))
par=np.array([[i*10,i*0.1] for i in range(10)])
result = myFoo(par, x)
print(result)

从错误消息来看，我假设 numba 在确定要写入的结果数组的引用时遇到问题。
请注意，如果使用列表而不是元组完成分配 result 中的形状规范，则会发出警告。

Numba：@jit 中的并行标志在我使用 numpy 二维数组的代码中不起作用

Numba : parallel flag in @jit does not work in my code with numpy 2D arrays

python

numpy

numpy-ndarray

numba