使用广播实现单向量与向量数组相乘时的Numba类型错误排查

报错信息 #

启用Numba的njit装饰器后，出现如下类型统一错误：

Exception occurred:
Type: TypingError
Message: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Cannot unify array(float64, 2d, C) and array(float64, 1d, C) for 'q1.2', defined at .\rotations.py (82)

File "rotations.py", line 82:
def quaternion_mult(q1, qa):
    <source elided>

    quat_result[:, 0] = (q1[:, 0] * q2[:, 0]) - (q1[:, 1] * q2[:, 1]) - (q1[:, 2] * q2[:, 2]) - (q1[:, 3] * q2[:, 3])
    ^

During: typing of assignment at .\rotations.py (82)

File "rotations.py", line 82:
def quaternion_mult(q1, qa):
    <source elided>

    quat_result[:, 0] = (q1[:, 0] * q2[:, 0]) - (q1[:, 1] * q2[:, 1]) - (q1[:, 2] * q2[:, 2]) - (q1[:, 3] * q2[:, 3])
    ^

During: resolving callee type: type(CPUDispatcher(<function quaternion_mult at 0x00000290EE6FE670>))
During: typing of call at .\rotations.py (102)

During: resolving callee type: type(CPUDispatcher(<function quaternion_mult at 0x00000290EE6FE670>))
During: typing of call at .\rotations.py (102)


File "rotations.py", line 102:
def quaternion_vect_mult(q1, vect_array):
    <source elided>

    temp = quaternion_mult(q1, q2)
    ^

相关函数代码 #

原函数实现如下：

from numba import njit
import numpy as np

@njit(cache=True)
def quaternion_conjugate_vect(q):
    """
    return the conjugate of a quaternion or an array of quaternions
    """
    return q * np.array([1, -1, -1, -1])


@njit(cache=True)
def quaternion_mult(q1, qa):
    """
    multiply an array of quaternions (Nx4) by a single quaternion.

    qa is always a (Nx4) array of quaternions np.ndarray
    q1 is always a single (1x4) quaternion np.ndarray

    """
    N = max(len(qa), len(q1))
    quat_result = np.zeros((N, 4), dtype=np.float64)

    if qa.ndim == 1:
        q2 = qa.copy().reshape((1, -1))
    else:
        q2 = qa

    if q1.ndim == 1:
        q1 = np.reshape(q1, (1, -1))

    quat_result[:, 0] = (q1[:, 0] * q2[:, 0]) - (q1[:, 1] * q2[:, 1]) - (q1[:, 2] * q2[:, 2]) - (q1[:, 3] * q2[:, 3])
    quat_result[:, 1] = (q1[:, 0] * q2[:, 1]) + (q1[:, 1] * q2[:, 0]) + (q1[:, 2] * q2[:, 3]) - (q1[:, 3] * q2[:, 2])
    quat_result[:, 2] = (q1[:, 0] * q2[:, 2]) + (q1[:, 2] * q2[:, 0]) + (q1[:, 3] * q2[:, 1]) - (q1[:, 1] * q2[:, 3])
    quat_result[:, 3] = (q1[:, 0] * q2[:, 3]) + (q1[:, 3] * q2[:, 0]) + (q1[:, 1] * q2[:, 2]) - (q1[:, 2] * q2[:, 1])

    return quat_result


@njit(cache=True)
def quaternion_vect_mult(q1, vect_array):
    """
    Multiplies an array of x,y,z coordinates by a single quaternion q1.
    """
    q2 = np.zeros((len(vect_array), 4), dtype=np.float64)
    q2[:, 1::] = vect_array

    temp = quaternion_mult(q1, q2)
    result = quaternion_mult(temp, quaternion_conjugate_vect(q1))

    return result[:, 1::]

测试场景 #

当传入形状为(4,)的四元组时编译报错，传入(1,4)时正常运行：

quat_single = np.random.random((4,))  # 此输入会触发报错
# quat_single = np.random.random([1,4])  # 此输入可正常运行
coord_array = np.random.random([9,3])

quaternion_vect_mult(quat_single, coord_array)

问题原因 #

Numba的nopython模式在编译阶段进行静态类型推断，无法处理依赖运行时条件（如ndim判断）的变量类型分支。原代码中q1可能在不同分支下是1D或2D数组，Numba无法统一这两种类型，进而抛出类型不匹配错误。而原生Numpy的广播是运行时动态处理，所以不会有此问题。

解决方案 #

修改quaternion_mult函数，强制将输入统一转为2D数组，避免依赖运行时维度判断，让Numba能明确推断变量类型：

@njit(cache=True)
def quaternion_mult(q1, qa):
    """
    multiply an array of quaternions (Nx4) by a single quaternion.

    qa can be (Nx4) array or (4,) single quaternion
    q1 can be (1x4) array or (4,) single quaternion
    """
    # 强制转为2D数组，不管输入是1D还是2D
    q1_2d = q1.reshape(1, -1) if q1.ndim == 1 else q1
    qa_2d = qa.reshape(1, -1) if qa.ndim == 1 else qa

    # 确定输出的行数：如果q1是单四元组则用qa的行数，反之亦然
    N = qa_2d.shape[0] if q1_2d.shape[0] == 1 else q1_2d.shape[0]
    quat_result = np.zeros((N, 4), dtype=np.float64)

    # 利用广播进行元素运算
    quat_result[:, 0] = (q1_2d[:, 0] * qa_2d[:, 0]) - (q1_2d[:, 1] * qa_2d[:, 1]) - (q1_2d[:, 2] * qa_2d[:, 2]) - (q1_2d[:, 3] * qa_2d[:, 3])
    quat_result[:, 1] = (q1_2d[:, 0] * qa_2d[:, 1]) + (q1_2d[:, 1] * qa_2d[:, 0]) + (q1_2d[:, 2] * qa_2d[:, 3]) - (q1_2d[:, 3] * qa_2d[:, 2])
    quat_result[:, 2] = (q1_2d[:, 0] * qa_2d[:, 2]) + (q1_2d[:, 2] * qa_2d[:, 0]) + (q1_2d[:, 3] * qa_2d[:, 1]) - (q1_2d[:, 1] * qa_2d[:, 3])
    quat_result[:, 3] = (q1_2d[:, 0] * qa_2d[:, 3]) + (q1_2d[:, 3] * qa_2d[:, 0]) + (q1_2d[:, 1] * qa_2d[:, 2]) - (q1_2d[:, 2] * qa_2d[:, 1])

    return quat_result

同时，调整quaternion_conjugate_vect函数，确保返回的数组维度与输入一致：

@njit(cache=True)
def quaternion_conjugate_vect(q):
    """
    return the conjugate of a quaternion or an array of quaternions
    """
    conjugate_factor = np.array([1, -1, -1, -1], dtype=np.float64)
    # 保持输入的维度
    return q * conjugate_factor

修改后，无论传入(4,)还是(1,4)的四元组，Numba都能正确编译并运行，同时保留原有的广播计算逻辑。

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