Math: Einsum

This module provides a wrapper for the einsum function from opt_einsum package. If opt_einsum is not installed, it falls back to the torch.einsum.

tad_mctc.math.einsum.einsum(subscripts, *operands, out=None, use_blas=True, optimize=True, memory_limit=None, backend='auto', **kwargs)

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

Parameters:

subscripts: Specifies the subscripts for summation. *operands: These are the arrays for the operation. out: A output array in which set the resulting output. use_blas: Do you use BLAS for valid operations, may use extra memory for more intermediates. optimize:- Choose the type of path the contraction will be optimized with

• if a list is given uses this as the path.

• ‘optimal’ An algorithm that explores all possible ways of

contracting the listed tensors. Scales factorially with the number of terms in the contraction. - ‘dp’ A faster (but essentially optimal) algorithm that uses dynamic programming to exhaustively search all contraction paths without outer-products. - ‘greedy’ An cheap algorithm that heuristically chooses the best pairwise contraction at each step. Scales linearly in the number of terms in the contraction. - ‘random-greedy’ Run a randomized version of the greedy algorithm 32 times and pick the best path. - ‘random-greedy-128’ Run a randomized version of the greedy algorithm 128 times and pick the best path. - ‘branch-all’ An algorithm like optimal but that restricts itself to searching ‘likely’ paths. Still scales factorially. - ‘branch-2’ An even more restricted version of ‘branch-all’ that only searches the best two options at each step. Scales exponentially with the number of terms in the contraction. - ‘auto’, None, True Choose the best of the above algorithms whilst aiming to keep the path finding time below 1ms. - ‘auto-hq’ Aim for a high quality contraction, choosing the best of the above algorithms whilst aiming to keep the path finding time below 1sec. - False will not optimize the contraction.

memory_limit:- Give the upper bound of the largest intermediate tensor contract will build.

• None or -1 means there is no limit.

• max_input means the limit is set as largest input tensor.

• A positive integer is taken as an explicit limit on the number of elements.

The default is None. Note that imposing a limit can make contractions exponentially slower to perform.

backend: 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:

The result of the einsum expression.

Notes:

This function should produce a 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 performed optimally. When NumPy is linked to a threaded BLAS, potential speedups are on the order of 20-100 for a six core machine.

tad_mctc.math.einsum.einsum_greedy(subscripts, *operands, out=None, use_blas=True, optimize=True, memory_limit=None, backend='auto', **kwargs)

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

Parameters:

subscripts: Specifies the subscripts for summation. *operands: These are the arrays for the operation. out: A output array in which set the resulting output. use_blas: Do you use BLAS for valid operations, may use extra memory for more intermediates. optimize:- Choose the type of path the contraction will be optimized with

• if a list is given uses this as the path.

• ‘optimal’ An algorithm that explores all possible ways of

contracting the listed tensors. Scales factorially with the number of terms in the contraction. - ‘dp’ A faster (but essentially optimal) algorithm that uses dynamic programming to exhaustively search all contraction paths without outer-products. - ‘greedy’ An cheap algorithm that heuristically chooses the best pairwise contraction at each step. Scales linearly in the number of terms in the contraction. - ‘random-greedy’ Run a randomized version of the greedy algorithm 32 times and pick the best path. - ‘random-greedy-128’ Run a randomized version of the greedy algorithm 128 times and pick the best path. - ‘branch-all’ An algorithm like optimal but that restricts itself to searching ‘likely’ paths. Still scales factorially. - ‘branch-2’ An even more restricted version of ‘branch-all’ that only searches the best two options at each step. Scales exponentially with the number of terms in the contraction. - ‘auto’, None, True Choose the best of the above algorithms whilst aiming to keep the path finding time below 1ms. - ‘auto-hq’ Aim for a high quality contraction, choosing the best of the above algorithms whilst aiming to keep the path finding time below 1sec. - False will not optimize the contraction.

memory_limit:- Give the upper bound of the largest intermediate tensor contract will build.

• None or -1 means there is no limit.

• max_input means the limit is set as largest input tensor.

• A positive integer is taken as an explicit limit on the number of elements.

The default is None. Note that imposing a limit can make contractions exponentially slower to perform.

backend: 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:

The result of the einsum expression.

Notes:

This function should produce a 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 performed optimally. When NumPy is linked to a threaded BLAS, potential speedups are on the order of 20-100 for a six core machine.

tad_mctc.math.einsum.einsum_optimal(subscripts, *operands, out=None, use_blas=True, optimize=True, memory_limit=None, backend='auto', **kwargs)

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

Parameters:

subscripts: Specifies the subscripts for summation. *operands: These are the arrays for the operation. out: A output array in which set the resulting output. use_blas: Do you use BLAS for valid operations, may use extra memory for more intermediates. optimize:- Choose the type of path the contraction will be optimized with

• if a list is given uses this as the path.

• ‘optimal’ An algorithm that explores all possible ways of

contracting the listed tensors. Scales factorially with the number of terms in the contraction. - ‘dp’ A faster (but essentially optimal) algorithm that uses dynamic programming to exhaustively search all contraction paths without outer-products. - ‘greedy’ An cheap algorithm that heuristically chooses the best pairwise contraction at each step. Scales linearly in the number of terms in the contraction. - ‘random-greedy’ Run a randomized version of the greedy algorithm 32 times and pick the best path. - ‘random-greedy-128’ Run a randomized version of the greedy algorithm 128 times and pick the best path. - ‘branch-all’ An algorithm like optimal but that restricts itself to searching ‘likely’ paths. Still scales factorially. - ‘branch-2’ An even more restricted version of ‘branch-all’ that only searches the best two options at each step. Scales exponentially with the number of terms in the contraction. - ‘auto’, None, True Choose the best of the above algorithms whilst aiming to keep the path finding time below 1ms. - ‘auto-hq’ Aim for a high quality contraction, choosing the best of the above algorithms whilst aiming to keep the path finding time below 1sec. - False will not optimize the contraction.

memory_limit:- Give the upper bound of the largest intermediate tensor contract will build.

• None or -1 means there is no limit.

• max_input means the limit is set as largest input tensor.

• A positive integer is taken as an explicit limit on the number of elements.

The default is None. Note that imposing a limit can make contractions exponentially slower to perform.

backend: 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:

The result of the einsum expression.

Notes:

This function should produce a 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 performed optimally. When NumPy is linked to a threaded BLAS, potential speedups are on the order of 20-100 for a six core machine.