使用 numpy 矩阵将 3D 数组的切片与二维数组的切片相乘

Question

A is a 2 dimensional array with dimensions (s, a) and B is a 3 dimensional array with (s, a, a).

I want B[i, :, :] @ A[i, :] for every i in range(s). The result should be arranged in an array of shape (s, a). In code:

s = 4
a = 3

A = np.random.uniform(size = [s,a])
B = np.random.uniform(size= [s ,a, a])

C = np.zeros_like(A)
for i in range(A.shape[0]):
  C[i,:] = (A[i,:] @ B[i,:,:])

I am looking for C. The catch is that everything should happen in numpy, without slicing A and B.

Answer 1

Numpy does batched multiplication to last two dimensions if the first n-2 dimensions match. You create a extra dimension with None and use matmul

(A[:,None,:]@B).reshape(A.shape[0],-1)