Повторите умножение числа вектор-матрица-вектор

Question

У меня есть набор векторов (n), другой набор векторов и набор 3x3 2D-массивов (T).

n = np.array([
    [[1, 2, 3]],
    [[2, 2, 3]],
    [[3, 2, 3]],
    [[4, 2, 3]],
    [[5, 2, 3]],
    [[6, 2, 3]]
])
s = np.array([
    [[1, 1, 5]],
    [[2, 2, 5]],
    [[3, 3, 5]],
    [[4, 4, 5]],
    [[5, 5, 5]],
    [[6, 6, 5]]
])

T = np.array([
    [[1, 2, 3],
     [1, 2, 3],
     [2, 2, 3]],

    [[2, 2, 3],
     [3, 2, 3],
     [4, 2, 3]],

    [[3, 2, 3],
     [5, 2, 3],
     [6, 2, 3]],

    [[4, 2, 3],
     [7, 2, 3],
     [8, 2, 3]]
])

Прямо сейчас мой текущий код перебирает n, s, а затем T:

result = np.array(n.shape[0], s.shape[0], T.shape[0])

for i in range(n.shape[0]):
    for j in range(s.shape[0]):
        for k in range(T.shape[0]):
            result[i][j][k] = np.sum(n[i] * T[k] * s[j].T)

Я пытался использовать np.apply_along_axis, но для его работы требуется одномерный массив. В идеале я пытаюсь найти решение, которое не требует циклов for.

Я пытался заставить np.tensordot() работать (и сделать это за две операции), но пока безуспешно.

У кого-нибудь есть идеи относительно более «тупого» способа сделать это?

Answer 1

np.einsum('imn,jnm,kmn->ijk', n, s, T)