numpy composition of einsums?

Question

Suppose that I have a np.einsum that performs some calculation, and then pump that directly into yet another np.einsum to do some other thing. Can I, in general, compose those two einsums into a single einsum?

My specific use case is that I am doing a transpose, a matrix multiplication, and then another matrix multiplication to compute 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

In single einsum call, it would be -

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

The intuition involved would be :

1. Start off from the innermost einsum call : 'ij,kj->ik'.
2. Moving out, the second one is : 'ij,jk->ik'. The first argument in it is the output from step#1. So, let's mould this argument for the second one based on the output from the first one, introducing new strings for new iterators : 'ik,kl->il'. Note that 'kl' is the second arg in this second einsum call, which is a.

Thus, combining, we have : 'ij,kj,kl->il' with the inputs in the same sequence, i.e. b,a for the innermost einsum call and then a incoming as the third input.