opt_einsum.contract

opt_einsum.contract(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', use_blas=True, optimize=True, memory_limit=None, backend='numpy')[source]

Evaluates the Einstein summation convention on the operands. A drop in replacment for NumPy’s einsum function that optimizes the order of contraction to reduce overall scaling at the cost of several intermediate arrays.

Parameters:
  • subscripts (str) – Specifies the subscripts for summation.
  • *operands (list of array_like) – These are the arrays for the operation.
  • out (array_like) – A output array in which set the resulting output.
  • dtype (str) – The dtype of the given contraction, see np.einsum.
  • order (str) – The order of the resulting contraction, see np.einsum.
  • casting (str) – The casting procedure for operations of different dtype, see np.einsum.
  • use_blas (bool) – Do you use BLAS for valid operations, may use extra memory for more intermediates.
  • optimize (bool, str, or list, optional (default: greedy)) – Choose the type of path.
    • if a list is given uses this as the path.
    • ‘greedy’ An algorithm that chooses the best pair contraction at each step. Scales cubically with the number of terms in the contraction.
    • ‘optimal’ An algorithm that tries all possible ways of contracting the listed tensors. Scales exponentially with the number of terms in the contraction.
  • memory_limit (int or None (default : None)) – The upper limit of the size of tensor created, by default this will be Give the upper bound of the largest intermediate tensor contract will build. By default (None) will size the memory_limit as the largest input tensor. Users can also specify -1 to allow arbitrarily large tensors to be built.
  • backend (str, optional (default: numpy)) – Which library to use to perform the required tensordot, transpose and einsum calls. Should match the types of arrays supplied, See contract_expression() for generating expressions which convert numpy arrays to and from the backend library automatically.
Returns:

out – The result of the einsum expression.
Return type:

array_like

Notes

This function should produce result identical to that of NumPy’s einsum function. The primary difference is contract will attempt to form intermediates which reduce the overall scaling of the given einsum contraction. By default the worst intermediate formed will be equal to that of the largest input array. For large einsum expressions with many input arrays this can provide arbitrarily large (1000 fold+) speed improvements.

For contractions with just two tensors this function will attempt to use NumPy’s built in BLAS functionality to ensure that the given operation is preformed in an optimal manner. When NumPy is linked to a threaded BLAS, potenital speedsups are on the order of 20-100 for a six core machine.

Examples

See opt_einsum.contract_path() or numpy.einsum()