python - numba.errors.TypingError : Failed in nopython mode pipeline (step: nopython frontend)

Question

我想使用 numba 加速我的 python 代码，如下所示:

import numpy as np
from numba import njit, jit
import time

@njit
def dist_gpu(a, b):
    d = np.linalg.norm(a[:, None, :] - b[None, :, :], axis=2)
    d = np.transpose(d)
    sorted_d = np.sort(d)
    sorted_ind = np.argsort(d)
    return sorted_d, sorted_ind

def get_a_b(r=10**4,c=10**1):
    a = np.random.uniform(-1, 1, (r, c)).astype('f')
    b = np.random.uniform(-1, 1, (r, c)).astype('f')
    return a,b

if __name__ == "__main__":
    a, b = get_a_b()
    st_t = time.time()
    dist_gpu(a,b)
    print('it took {} s'.format(time.time()-st_t))

我在函数中使用装饰器 @njit 或 @jit(nopython=True) 时出现以下错误:

Traceback (most recent call last):
  File "stackoverflow_Q.py", line 22, in <module>
    dist_gpu(a,b)
  File "/usr/local/lib/python2.7/dist-packages/numba/dispatcher.py", line 401, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/usr/local/lib/python2.7/dist-packages/numba/dispatcher.py", line 344, in error_rewrite
    reraise(type(e), e, None)
  File "<string>", line 2, in reraise
numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (array(float32, 2d, C), Tuple(slice<a:b>, none, slice<a:b>))
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
In definition 6:
    All templates rejected with literals.
In definition 7:
    All templates rejected without literals.
In definition 8:
    TypeError: unsupported array index type none in Tuple(slice<a:b>, none, slice<a:b>)
In definition 9:
    TypeError: unsupported array index type none in Tuple(slice<a:b>, none, slice<a:b>)
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: typing of intrinsic-call at stackoverflow_Q.py (8)
[2] During: typing of static-get-item at stackoverflow_Q.py (8)

File "stackoverflow_Q.py", line 8:
def dist_gpu(a, b):
    d = np.linalg.norm(a[:, None, :] - b[None, :, :], axis=2)
    ^

This is not usually a problem with Numba itself but instead often caused by
the use of unsupported features or an issue in resolving types.

To see Python/NumPy features supported by the latest release of Numba visit:
http://numba.pydata.org/numba-doc/latest/reference/pysupported.html
and
http://numba.pydata.org/numba-doc/latest/reference/numpysupported.html

For more information about typing errors and how to debug them visit:
http://numba.pydata.org/numba-doc/latest/user/troubleshoot.html#my-code-doesn-t-compile

If you think your code should work with Numba, please report the error message
and traceback, along with a minimal reproducer at:
https://github.com/numba/numba/issues/new

任何人都可以帮助我使用 numba 做错事，即使没有它也不会发生此错误吗？

Answer 1

也许试试

@njit
def dist_gpu(a, b):
   d = np.linalg.norm(np.expand_dims(a, 1) - np.expand_dims(b, 0), axis=2)

AFAIK，Numba 不(还)喜欢 None 和 np.newaxis... 一次只扩展 1 个轴 np.expand_dims 有效。