python - 将 numpy.stack() 与 numba njit 一起使用时出现 TypingError

Question

最初的问题与使用 np.linspace 和数组作为开始和停止参数有关，尽管现在我遇到了我想出的解决方法的问题。

采取以下措施:

from numba import njit
import numpy as np

@njit
def f1():
  start = np.array([0.1, 1.0], np.float32)
  stop = np.array([1.0, 10.0], np.float32)
  return np.linspace(start, stop, 10)

f1()

这会引发错误，因为 though documented as supporting "only the 3-argument form"在 linspace 中，它们的实际含义是“具有用于开始和停止的标量值的 3 参数形式”。

所以我想出了以下解决方法:

import numpy as np
from numba import njit

@njit
def f2():
  start = np.array([0.1, 1.0], np.float32)
  stop = np.array([1.0, 10.0], np.float32)
  pts_0 = np.linspace(start[0], stop[0], 10).astype(np.float32) # works
  pts_1 = np.linspace(start[1], stop[1], 10).astype(np.float32) # works
  return np.stack([pts_0, pts_1]).T                             # error

引发此错误:

---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)
c:\Users\X\Desktop\X\data_analysis.ipynb Cell 46' in <cell line: 18>()
     15   pts_1 = np.linspace(start[1], stop[1], 10).astype(np.float32)
     16   return np.stack([pts_0, pts_1]).T
---> 18 r = f2()

File c:\Users\X\miniconda3\envs\X\lib\site-packages\numba\core\dispatcher.py:468, in _DispatcherBase._compile_for_args(self, *args, **kws)
    464         msg = (f"{str(e).rstrip()} \n\nThis error may have been caused "
    465                f"by the following argument(s):\n{args_str}\n")
    466         e.patch_message(msg)
--> 468     error_rewrite(e, 'typing')
    469 except errors.UnsupportedError as e:
    470     # Something unsupported is present in the user code, add help info
    471     error_rewrite(e, 'unsupported_error')

File c:\Users\X\miniconda3\envs\X\lib\site-packages\numba\core\dispatcher.py:409, in _DispatcherBase._compile_for_args.<locals>.error_rewrite(e, issue_type)
    407     raise e
    408 else:
--> 409     raise e.with_traceback(None)

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function stack at 0x00000186F280CAF0>) found for signature:
 
 >>> stack(list(array(float32, 1d, C))<iv=None>)

同样，根据 documentation , np.stack 受支持(对此也没有旁注)。

我错过了什么？

Answer 1

np.stack 受支持，但目前它需要一个元组而不是一个列表。这是一个固定的代码:

@njit
def f2():
  start = np.array([0.1, 1.0], np.float32)
  stop = np.array([1.0, 10.0], np.float32)
  pts_0 = np.linspace(start[0], stop[0], 10).astype(np.float32) # works
  pts_1 = np.linspace(start[1], stop[1], 10).astype(np.float32) # works
  return np.stack((pts_0, pts_1)).T                             # works

顺便说一句，请注意 np.stack((pts_0, pts_1)).T 效率不高，因为它创建了临时数组 和一个非连续 View .由于使用 Numba 的目的是加快代码速度，因此请考虑在这里使用应该更快的基本循环。 astype(np.float32) 也是如此:循环可以就地转换值。内存和分配很昂贵，这通常是使 Numpy 变慢的原因(也缺乏特定用途的功能)。这样的事情将来会变慢(有关更多信息，请考虑阅读更多关于“内存墙”的内容)，因此需要避免它们。

这是一个带有基本循环的明显更快的版本:

@njit
def f2():
    start1, start2 = np.float32(0.1), np.float32(1.0)
    stop1, stop2 = np.float32(1.0), np.float32(10.0)
    steps = 10
    delta = np.float32(1 / (steps - 1))
    res = np.empty((steps, 2), dtype=np.float32)
    for i in range(steps):
        res[i, 0] = start1 + (stop1 - start1) * (delta * i)
        res[i, 1] = start2 + (stop2 - start2) * (delta * i)
    return res

请注意，由于 32 位 FP 舍入，结果可能会略有不同。