Python中适配Numba的(K,3)矩阵向量化计算方案咨询

原始单向量场景 #

当m、n为长度3的向量时，可通过以下代码生成3x3矩阵：

e_matrix = np.array([m[i]*m[j]-n[i]*n[j] for i in range(3) for j in range(3)])

批量场景需求 #

现在m、n为shape(K,3)的矩阵，需要生成shape(3,3,K)的e_matrix，已知朴素循环实现：

e_matrix = np.zeros((3,3,K))

for k in range(K):
    e_matrix[:,:,k] = np.array([m[k,i]*m[k,j]-n[k,i]*n[k,j] for i in range(3) for j in range(3)]) 

兼容Numba的向量化方案 #

可以利用NumPy的广播机制实现完全向量化，且该方案兼容Numba JIT编译（无tensordot等受限操作）：

import numpy as np

# 生成(K,3,3)的外积矩阵
m_outer = m[:, :, None] * m[:, None, :]
n_outer = n[:, :, None] * n[:, None, :]
# 转置得到目标shape(3,3,K)
e_matrix = (m_outer - n_outer).transpose(1, 2, 0)

Numba JIT编译验证

将上述逻辑封装为JIT函数，可实现更快的运行速度：

from numba import jit

@jit(nopython=True)
def compute_e_matrix(m, n):
    m_outer = m[:, :, None] * m[:, None, :]
    n_outer = n[:, :, None] * n[:, None, :]
    return (m_outer - n_outer).transpose(1, 2, 0)

# 调用示例
K = 100
m = np.random.rand(K, 3)
n = np.random.rand(K, 3)
e_matrix = compute_e_matrix(m, n)

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