arrays - テンソルのnumpyドット積（3d x 2d）

Question

現在私は使用しています

Na = (3, 2, 4)
Nb = Na[1:]
A = np.arange(np.prod(Na)).reshape(Na)
b = np.arange(np.prod(Nb)).reshape(Nb)

計算したい：

r = np.empty((A.shape[0], A.shape[2])
for i in range(A.shape[2]):
    r[:, i] = np.dot(A[:, :, i], b[:, i])

つまり、A は「4 つの行列のスタック」((3,2) の形状)、つまり 3 次元配列、b は「4 つのベクトルのスタック」((3,) の形状)、つまり 2 次元配列です。望ましい結果は、「4 つの行列ベクトル積のスタック」、つまりベクトルのスタック、つまり再び 2 次元配列 (形状 (3, 4)) です。

np.einsum と np.tensordot を中程度の深さで調べましたが、これらを使用して作成したソリューションは、ループ ソリューションと少なくとも同じくらい長く、読みにくいものでした。

ただし、その単純な問題にはワンライナーが必要だと思います。