在 Python/Numba 中添加不同形状的矩阵的最快方法

Question

我想“添加”两个矩阵，一个形状a为 ( K,T) 的矩阵和一个形状b为 的矩阵(K,N)，以产生形状为 的矩阵(K,T,N）

以下工作正常：

import numpy as np 
from numba import njit

@njit
def add_matrices(a, b):
    K, T, N = a.shape[0], a.shape[1], b.shape[1]
    result_matrix = np.empty((K, T, N))
    
    for k in range(K):
        for t in range(T):
            for n in range(N):
                result_matrix[k, t, n] = a[k, t] + b[k, n]
    
    return result_matrix


K = 10
T = 11
N = 12
a = np.ones((K,T))
b = np.ones((K,N))

result = add_matrices(a, b)

有没有一种更快（矢量化？）的方法来做到这一点，不需要 for 循环，我认为这会减慢函数的速度，特别是对于较大的 值K, T,N？

Answer 1

使用广播。

a[:,:,None] + b[:,None,:]

这使得a看起来有大小 [K, T, 1]，并且b有大小[K, 1, N]。Numpy 知道如何将这两者相加。