Reusing Paths

If you expect to use a particular contraction repeatedly, it can make things simpler and more efficient not to compute the path each time. Instead, supplying opt_einsum.contract_expression with the contraction string and the shapes of the tensors generates a opt_einsum.ContractExpression which can then be repeatedly called with any matching set of arrays. For example:

my_expr = oe.contract_expression("abc,cd,dbe->ea", (2, 3, 4), (4, 5), (5, 3, 6))
print(my_expr)
#> <ContractExpression('abc,cd,dbe->ea')>
#>   1.  'dbe,cd->bce' [GEMM]
#>   2.  'bce,abc->ea' [GEMM]

The ContractExpression can be called with 3 arrays that match the original shapes without having to recompute the path:

x, y, z = (np.random.rand(*s) for s in [(2, 3, 4), (4, 5), (5, 3, 6)])
my_expr(x, y, z)
#> array([[ 3.08331541,  4.13708916],
#>        [ 2.92793729,  4.57945185],
#>        [ 3.55679457,  5.56304115],
#>        [ 2.6208398 ,  4.39024187],
#>        [ 3.66736543,  5.41450334],
#>        [ 3.67772272,  5.46727192]])

Note that few checks are performed when calling the expression, and while it will work for a set of arrays with the same ranks as the original shapes but differing sizes, it might no longer be optimal.

Specifying Constants

Often one generates contraction expressions where some of the tensor arguments will remain constant across many calls. opt_einsum.contract_expression allows you to specify the indices of these constant arguments, allowing opt_einsum to build and then reuse as many constant contractions as possible.

Take for example the equation:

eq = "ij,jk,kl,lm,mn->ni"

where we know that only the first and last tensors will vary between calls. We can specify this by marking the middle three as constant - we then need to supply the actual arrays rather than just the shapes to opt_einsum.contract_expression:

#           A       B       C       D       E
shapes = [(9, 5), (5, 5), (5, 5), (5, 5), (5, 8)]

# mark the middle three arrays as constant
constants = [1, 2, 3]

# generate the constant arrays
B, C, D = [np.random.randn(*shapes[i]) for i in constants]

# supplied ops are now mix of shapes and arrays
ops = (9, 5), B, C, D, (5, 8)

expr = oe.contract_expression(eq, *ops, constants=constants)
expr
#> <ContractExpression('ij,[jk,kl,lm],mn->ni', constants=[1, 2, 3])>

The expression now only takes the remaining two arrays as arguments (the tensors with 'ij' and 'mn' indices), and will store as many reusable constant contractions as possible.

.. code:: python

A1, E1 = np.random.rand(*shapes[0]), np.random.rand(*shapes[-1])
out1 = expr(A1, E1)
out1.shap
#> (8, 9)

A2, E2 = np.random.rand(*shapes[0]), np.random.rand(*shapes[-1])
out2 = expr(A2, E2)
out2.shape
#> (8, 9)

np.allclose(out1, out2)
#> False

print(expr)
#> <ContractExpression('ij,[jk,kl,lm],mn->ni', constants=[1, 2, 3])>
#>   1.  'jm,mn->jn' [GEMM]
#>   2.  'jn,ij->ni' [GEMM]

Where we can see that the expression now only has to perform two contractions to compute the output.

Note

The constant part of an expression is lazily generated upon the first call (specific to each backend), though it can also be explicitly built by calling opt_einsum.contract.ContractExpression.evaluate_constants.

We can confirm the advantage of using expressions and constants by timing the following scenarios, first setting A = np.random.rand(*shapes[0]) and E = np.random.rand(*shapes[-1]).

Contract from scratch

%timeit oe.contract(eq, A, B, C, D, E)
#> 239 µs ± 5.06 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Contraction with an expression but no constants

expr_no_consts = oe.contract_expression(eq, *shapes)
%timeit expr_no_consts(A, B, C, D, E)
#> 76.7 µs ± 2.47 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Contraction with an expression and constants marked

%timeit expr(A, E)
#> 40.8 µs ± 1.22 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Although this gives us a rough idea, of course the efficiency savings are hugely dependent on the size of the contraction and number of possible constant contractions.

We also note that even if there are no constant contractions to perform, it can be very advantageous to specify constant tensors for particular backends. For instance, if a GPU backend is used, the constant tensors will be kept on the device rather than being transferred each time.