将 numpy.stack() 与 numba njit 一起使用时出现打字错误

最初的问题与使用有关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()

这会引发错误，因为尽管记录为支持“仅 3 参数形式” https://numba.readthedocs.io/en/stable/reference/numpysupported.html of 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>)

再次，根据文档 https://numba.readthedocs.io/en/stable/reference/numpysupported.html, np.stack支持（对此也没有任何评论）。

我缺少什么？

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效率不是很高，因为它创建了临时数组和不连续的视图。由于使用 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 舍入，结果可能略有不同。