NumPy состав Einsums?

Question

Предположим, у меня есть np.einsum, который выполняет какое-то вычисление, а затем закачивает это непосредственно в еще один np.einsum, чтобы сделать что-то другое. Можно ли вообще объединить эти два einsum в один einsum?

Мой конкретный вариант использования - это выполнение транспонирования, умножения матриц, а затем еще одного умножения матриц для вычисления b a^T a:

import numpy as np
from numpy import array

a = array([[1, 2],
       [3, 4]])
b = array([[1, 2],
       [3, 4],
       [5, 6]])

matrix_multiply_by_transpose = 'ij,kj->ik'
matrix_multiply = 'ij,jk->ik'
test_answer = np.einsum(matrix_multiply,
    np.einsum(matrix_multiply_by_transpose, 
        b, a
    ),
    a
)

assert np.array_equal(test_answer, 
    np.einsum(an_answer_to_this_question, b, a, a))
#or, the ultimate most awesomest answer ever, if such a thing even exists
assert np.array_equal(test_answer,
    np.einsum(the_bestest_answer(matrix_multiply_by_transpose, matrix_multiply),
        b, a, a)
)

Answer 1

В одном вызове einsum это будет -

np.einsum('ij,kj,kl->il',b,a,a)

Интуиция будет:

1. Начните с самого внутреннего einsum вызова: 'ij,kj->ik'.
2. Выйдя, второй: 'ij,jk->ik'. Первым аргументом в нем является вывод из step#1. Итак, давайте сформируем этот аргумент для второго на основе выходных данных первого, введя новые строки для новых итераторов: 'ik,kl->il'. Обратите внимание, что 'kl' является вторым аргументом в этом втором вызове einsum, то есть a.

Таким образом, объединяя, мы имеем: 'ij,kj,kl->il' с входами в той же последовательности, то есть b,a для самого внутреннего вызова einsum и затем a входящий в качестве третьего входа.