If you expect to use a particular contraction repeatedly, it can make things simpler and more efficient not to compute the path each time.
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.
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]
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.
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
# 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
'mn' indices), and will store as many reusable
constant contractions as possible.
.. code:: python
A1, E1 = np.random.rand(*shapes), np.random.rand(*shapes[-1]) out1 = expr(A1, E1) out1.shap #> (8, 9) A2, E2 = np.random.rand(*shapes), 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.
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
We can confirm the advantage of using expressions and constants by timing the
following scenarios, first setting
A = np.random.rand(*shapes) 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.