Why optimize in the einsum can accelerate binary contraction?

Question

In https://numpy.org/doc/stable/reference/generated/numpy.einsum.html

optimize{False, True, ‘greedy’, ‘optimal’}, optional Controls if intermediate optimization should occur. No optimization will occur if False and True will default to the ‘greedy’ algorithm. Also accepts an explicit contraction list from the np.einsum_path function. See np.einsum_path for more details. Defaults to False.

It seems to me the optimize flag is to choose the order in multiple contractions. E.g.,

A B C -> D

for (AB)C or A(BC) or (AC)B which is faster, not for a binary contraction, e.g., AB->C.

For the following code for A[a,b] * B[b,c,d] = C[a,c,d]

import numpy as np
import time
import scipy.stats

# from https://stackoverflow.com/questions/15033511/compute-a-confidence-interval-from-sample-data
def mean_confidence_interval(data, var_name, unit, confidence=0.95):
    a = 1.0 * np.array(data)
    n = len(a)
    m, se = np.mean(a), scipy.stats.sem(a)
    h = se * scipy.stats.t.ppf((1 + confidence) / 2., n-1)
    
    print(var_name, round(m, 5), "\u00B1",  round(h, 5), unit )


def einsum_greedy(A, B, na, nb, nc, nd):
    #res = np.zeros((na,nc,nd))
    res = np.einsum('ab,bcd->acd', A, B, optimize="greedy")
    return res 


def einsum_standard(A, B, na, nb, nc, nd):
   # res = np.zeros((na,nc,nd))
    res = np.einsum('ab,bcd->acd', A, B)
    return res 

def btime_ABC(name_def, name_out, A, B, C, na, nb, nc, nd, n_times):
    global opt_path
    list_time = []
    for i in range(n_times):
        start_time = time.time()
        C = name_def(A, B, na, nb, nc, nd)
        finish_time = time.time()
        list_time.append(finish_time - start_time)

    mean_confidence_interval(list_time, name_out, 's' )



# A[a,b] * B[b,c,d] = C[a,c,d]
na = nb = nc = nd = dim_comm = 90
n_times = 60


print('number of common dimension', dim_comm)
print('number of averaged time', n_times)

A = np.random.random((na,nb))
B = np.random.random((nb,nc,nd))
C1 = np.zeros((na,nc,nd))
C2 = np.zeros((na,nc,nd))

btime_ABC(einsum_standard, 'einsum_standard', A, B, C1, na, nb, nc, nd, n_times)
btime_ABC(einsum_greedy, 'einsum_greedy', A, B, C2, na, nb, nc, nd, n_times)

I got

number of common dimension 90
number of averaged time 60
einsum_standard 0.04799 ± 0.00312 s
einsum_greedy 0.00805 ± 0.00137 s

the optimize flag helps in the binary contraction A[a,b] * B[b,c,d] = C[a,c,d]. Then, why?

Answer 1

My timings:

In [26]: A = np.random.random((90,80))
In [27]: B = np.random.random((80,81,82))
In [28]: timeit np.einsum('ab,bcd->acd',A,B,optimize=False)
39.2 ms ± 1.51 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [29]: timeit np.einsum('ab,bcd->acd',A,B,optimize=True)
9.06 ms ± 70.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Looking at the einsum code, I see, early on a:

# If no optimization, run pure einsum
if optimize is False:
    return c_einsum(*operands, **kwargs)

If not it checks various parameters, and does

operands, contraction_list = einsum_path(*operands, optimize=optimize,
                                         einsum_call=True)
...
# Call tensordot if still possible
if blas:
    ...
new_view = tensordot(*tmp_operands, axes=(tuple(left_pos), tuple(right_pos)))

Since we have only 2 arguments and path is straightforward, I think the True case is just:

In [30]: timeit np.tensordot(A,B,(1,0))
7.62 ms ± 609 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

which is, from past study of tensordot:

In [31]: timeit (A@B.reshape(80,-1)).reshape(90,81,82)
6.44 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

So basically the time difference is between running a compiled "pure einsum", and an alternative that casts this as a matmul problem, which can use the optimized BLAS routines. [29] time appears to be the [31] time plus some overhead.