python - einsums 的 numpy 组成？

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 作为第三个输入传入。