3

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)
)
Him
  • 5,257
  • 3
  • 26
  • 83

1 Answers1

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.

Divakar
  • 218,885
  • 19
  • 262
  • 358
  • I see that the `ij, kj` bit remained from the first multiply. What is the reasoning that went on here? – Him Nov 06 '18 at 18:16
  • Ah! I see it. :) – Him Nov 06 '18 at 18:19
  • 1
    Yes, nested einsums may always be rewritten as a single statement, by associativity of the multilinear product. If it is a good idea performance-wise might depend on the particulars though. If one of the subterms is known to reduce to a small temporary, computing that temporary first might be more efficient, since einsum is not necessarily smart about doing that optimization itself. So if performance is indeed an issue, make sure to benchmark. – Eelco Hoogendoorn Nov 06 '18 at 19:38
  • As a note with `optimize=True` einsum is quite smart about building intermediates and using BLAS when possible. – Daniel Dec 31 '18 at 21:15