在 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 知道如何将这两者相加。