Умножить Tensor-Slice на Matrix-Row

Я пытаюсь найти способ эффективного вычисления матрично-векторного произведения каждого двумерного среза тензора по глубине (shape: (n, n, m)) с каждой строкой матрицы (shape: (n, m)). То, что я пытаюсь сделать, выглядит так без векторизации:

import numpy as np

np.random.seed(42)

np.random.seed(1)

a = np.arange(16).reshape(4, 4)
b = np.random.randn(4, 4, 4)
c = np.zeros((4, 4))

for i in range(4):
    c[i] = b[..., i] @ a[i]

Который дает:

<<< print(c)
>>> [[  0.53623421  -0.10257152  -1.34855819  -1.72774519]
     [-18.13932187   1.82230599 -11.99348739  15.0787884 ]
     [ 38.5704751   -0.38514407   4.19673794   9.01941574]
     [-68.11165212  -5.52586601  64.69279036  11.3196871 ]]

Самое близкое, с чем я пришел, было:

<<< print(np.einsum("ij,ijk->ki", a, b))
>>> [[  0.53623421,  -2.66288958, -16.91264496, -12.98103047],
     [ -3.95244251,   1.82230599, -20.5351456 ,  34.69343339],
     [  8.07033597,  -0.90215803,   4.19673794,  12.57858867],
     [ -8.18116212,  -3.54815874,  46.60443317,  11.3196871 ]]

Где хотя бы верхний левый и нижний.правый элементы совпадают.