opt_einsum is a neat library that can optimize eisum formulations on the fly, but can also be used to simplify formulas to find better structures.
The script eincheck.py
shows an example run on
X = np.random.randn(16, 9, 676)
Y = np.random.randn(16, 32, 676)
np.einsum('bkl,bml,bkn,bmn->mk', X, Y, X, Y)
The "optimized path" gives
np.einsum('bmk,bmk->mk', np.einsum('bml,bkl->bmk', Y, X), np.einsum('bmn,bkn->bmk', Y, X))
but the optimizer does not know that the inputs are (X, Y, X, Y)
,
so it is possible to further simplify things by hand to avoid computing the same thing twice, leading to
np.sum(np.einsum('bml,bkl->bmk', Y, X)**2, axis=0)
giving nicer results
+-----------------+-----------------+-----------------+
| Functions | Time (tot) | Time (per iter) |
+-----------------+-----------------+-----------------+
| [0] naive | 0.047003s | 0.004700s |
| [1] optimized | 0.033002s | 0.003300s |
| [2] optimized2 | 0.017001s | 0.001700s |
+-----------------+-----------------+-----------------+