Fastest way to add matrices of different shapes in Python/Numba

Question

I want to "add" two matrices, a matrix a with shape (K,T) and a matrix b of shape (K,N), to result in a matrix of shape (K,T,N)

The following works ok:

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)

Is there a faster (vectorized?) way to do it that doesn't require the for loops, which I think is slowing down the function, especially for larger values of K, T,N?

Answer 1

Use broadcasting.

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

This makes a appear to have size [K, T, 1], and b to have size [K, 1, N]. Numpy knows how to add these two together.