Numba jitclass方法数组连续性警告及指定类型报错咨询

问题背景 #

以下代码实现了存储旋转矩阵的Rotator jitclass，其rotate方法完成矩阵与向量的乘法：

import numpy as np
import numpy.typing as npt
import numba as nb
from numba.experimental import jitclass

@jitclass([("_R", nb.float64[:,:])])
class Rotator:
    _R : npt.NDArray[np.float64]
    def __init__(self, R : npt.NDArray[np.float64]):
        self._R = R

    def rotate(self, v : npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return self._R @ v

R = np.eye(3)
print(R.flags)
rotator = Rotator(R)
rotator.rotate(np.array([1., 0., 0.]))

打印的数组flags显示R是C_CONTIGUOUS（行主序连续）：

C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False

但运行时Numba触发性能警告：

<string>:3: NumbaPerformanceWarning: '@' is faster on contiguous arrays, called on (Array(float64, 2, 'A', False, aligned=True), Array(float64, 1, 'C', False, aligned=True))

而将旋转逻辑提取为单独的@njit函数时无此警告：

@njit
def rotate(R : npt.NDArray[np.float64], v : npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    return R @ v

rotate(np.eye(3), np.array([1., 0., 0.]))

尝试将jitclass成员_R声明为行主序数组（@jitclass([("_R", nb.float64[::1,:])])）时，触发TypingError：

<string>:3: NumbaPendingDeprecationWarning: Code using Numba extension API maybe depending on 'old_style' error-capturing, which is deprecated and will be replaced by 'new_style' in a future release. See details at https://numba.readthedocs.io/en/latest/reference/deprecation.html#deprecation-of-old-style-numba-captured-errors
Exception origin:
  File "/Users/xx/miniconda3/envs/xx/lib/python3.10/site-packages/numba/np/arrayobj.py", line 6397, in array_to_array
    assert fromty.mutable != toty.mutable or toty.layout == 'A'

Traceback (most recent call last):
  File "/private/tmp/test.py", line 17, in <module>
    rotator = Rotator(R)
  File "/Users/xx/miniconda3/envs/xx/lib/python3.10/site-packages/numba/experimental/jitclass/base.py", line 124, in __call__
    return cls._ctor(*bind.args[1:], **bind.kwargs)
  File "/Users/xx/miniconda3/envs/xx/lib/python3.10/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/Users/xx/miniconda3/envs/xx/lib/python3.10/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Internal error at <numba.core.typeinfer.CallConstraint object at 0x11edbbbe0>.
Failed in nopython mode pipeline (step: native lowering)

Enable logging at debug level for details.

File "<string>", line 3:
<source missing, REPL/exec in use?>

核心疑问 #

1. 明明输入数组是C连续的，为什么Numba认为jitclass成员中的数组是不连续的（布局标记为'A'）？
2. 如何解决上述性能警告和后续的TypingError？

问题解析与解决方案 #

1. 为什么jitclass成员数组被标记为非连续？ #

当使用nb.float64[:,:]声明jitclass成员时，Numba默认将其类型推断为**任意布局（'A'）**的数组，不会保留输入数组的连续性信息。即使传入的是C连续数组，jitclass内部存储时也会按通用布局类型处理，因此在执行@运算时，Numba会触发连续性警告——它无法确定成员数组的实际布局是否为连续，只能按最保守的通用布局处理。

而单独的@njit函数可以直接根据输入参数的实际类型推断出数组的连续性，因此不会触发警告。

2. 解决警告与TypingError的方案 #

方案一：在jitclass初始化时显式转换数组为连续布局

在__init__方法中，使用np.ascontiguousarray将输入数组转换为C连续布局，同时保持原有的nb.float64[:,:]声明：

import numpy as np
import numpy.typing as npt
import numba as nb
from numba.experimental import jitclass

@jitclass([("_R", nb.float64[:,:])])
class Rotator:
    _R : npt.NDArray[np.float64]
    def __init__(self, R : npt.NDArray[np.float64]):
        # 显式转换为C连续数组
        self._R = np.ascontiguousarray(R)

    def rotate(self, v : npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return self._R @ v

R = np.eye(3)
print(R.flags)
rotator = Rotator(R)
rotator.rotate(np.array([1., 0., 0.]))

这种方式既保证了数组的连续性，又避免了类型声明的错误，同时消除性能警告。

方案二：正确声明连续数组类型（需注意Numba版本兼容性）

如果需要在类型声明阶段就指定连续布局，需使用正确的Numba数组类型语法：对于2D C连续数组，应使用nb.float64[:,::1]而非nb.float64[::1,:]（后者是列维度连续，对应Fortran布局）。修改后的代码如下：

import numpy as np
import numpy.typing as npt
import numba as nb
from numba.experimental import jitclass

# 声明为C连续的2D数组（行主序）
@jitclass([("_R", nb.float64[:,::1])])
class Rotator:
    _R : npt.NDArray[np.float64]
    def __init__(self, R : npt.NDArray[np.float64]):
        self._R = R

    def rotate(self, v : npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return self._R @ v

R = np.eye(3)
print(R.flags)
rotator = Rotator(R)
rotator.rotate(np.array([1., 0., 0.]))

注意：部分旧版本Numba可能对jitclass中连续数组类型的支持存在问题，若仍报错，优先使用方案一。

内容的提问来源于stack exchange，提问作者Carpetfizz