Ctrl+K

Einsum path helpers#

ivy.utils.einsum_path_helpers.can_dot(inputs, result, idx_removed)[source]#

Check if we can use BLAS (np.tensordot) call and its beneficial to do so.

Parameters:

  • inputs (list of str) – Specifies the subscripts for summation.

  • result (str) – Resulting summation.

  • idx_removed (set) – Indices that are removed in the summation

Returns:

type (bool) – Returns true if BLAS should and can be used, else False

Notes

If the operations is BLAS level 1 or 2 and is not already aligned we default back to einsum as the memory movement to copy is more costly than the operation itself.

Examples

# Standard GEMM operation >>> can_dot([‘ij’, ‘jk’], ‘ik’, set(‘j’)) True # Can use the standard BLAS, but requires odd data movement >>> can_dot([‘ijj’, ‘jk’], ‘ik’, set(‘j’)) False # DDOT where the memory is not aligned >>> can_dot([‘ijk’, ‘ikj’], ‘’, set(‘ijk’)) False

ivy.utils.einsum_path_helpers.compute_size_by_dict(indices, idx_dict)[source]#

Compute the product of the elements in indices based on the dictionary idx_dict.

Parameters:

  • indices (iterable) – Indices to base the product on.

  • idx_dict (dictionary) – Dictionary of index sizes

Returns:

ret (int) – The resulting product.

Examples

>>> compute_size_by_dict('abbc', {'a': 2, 'b':3, 'c':5})
90
ivy.utils.einsum_path_helpers.find_contraction(positions, input_sets, output_set)[source]#

Find the contraction for a given set of input and output sets.

Parameters:

  • positions (iterable) – Integer positions of terms used in the contraction.

  • input_sets (list) – List of sets that represent the lhs side of the einsum subscript

  • output_set (set) – Set that represents the rhs side of the overall einsum subscript

Returns:

  • new_result (set) – The indices of the resulting contraction

  • remaining (list) – List of sets that have not been contracted, the new set is appended to the end of this list

  • idx_removed (set) – Indices removed from the entire contraction

  • idx_contraction (set) – The indices used in the current contraction

Examples

# A simple dot product test case >>> pos = (0, 1) >>> isets = [set(‘ab’), set(‘bc’)] >>> oset = set(‘ac’) >>> find_contraction(pos, isets, oset) ({‘a’, ‘c’}, [{‘a’, ‘c’}], {‘b’}, {‘a’, ‘b’, ‘c’}) # A more complex case with additional terms in the contraction >>> pos = (0, 2) >>> isets = [set(‘abd’), set(‘ac’), set(‘bdc’)] >>> oset = set(‘ac’) >>> find_contraction(pos, isets, oset) ({‘a’, ‘c’}, [{‘a’, ‘c’}, {‘a’, ‘c’}], {‘b’, ‘d’}, {‘a’, ‘b’, ‘c’, ‘d’})

ivy.utils.einsum_path_helpers.flop_count(idx_contraction, inner, num_terms, size_dictionary)[source]#

Compute the number of FLOPS in the contraction.

Parameters:

  • idx_contraction (iterable) – The indices involved in the contraction

  • inner (bool) – Does this contraction require an inner product?

  • num_terms (int) – The number of terms in a contraction

  • size_dictionary (dict) – The size of each of the indices in idx_contraction

Returns:

flop_count (int) – The total number of FLOPS required for the contraction.

Examples

>>> flop_count('abc', False, 1, {'a': 2, 'b':3, 'c':5})
30

>>> flop_count('abc', True, 2, {'a': 2, 'b':3, 'c':5})
60
ivy.utils.einsum_path_helpers.greedy_path(input_sets, output_set, idx_dict, memory_limit)[source]#

Find the path by contracting the best pair until the input list is exhausted. The best pair is found by minimizing the tuple (-prod(indices_removed), cost). What this amounts to is prioritizing matrix multiplication or inner product operations, then Hadamard like operations, and finally outer operations. Outer products are limited by memory_limit. This algorithm scales cubically with respect to the number of elements in the list input_sets.

Parameters:

  • input_sets (list) – List of sets that represent the lhs side of the einsum subscript

  • output_set (set) – Set that represents the rhs side of the overall einsum subscript

  • idx_dict (dictionary) – Dictionary of index sizes

  • memory_limit (int) – The maximum number of elements in a temporary array

Returns:

path (list) – The greedy contraction order within the memory limit constraint.

Examples

>>> isets = [set('abd'), set('ac'), set('bdc')]
>>> oset = set()
>>> idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
>>> greedy_path(isets, oset, idx_sizes, 5000)
[(0, 2), (0, 1)]
ivy.utils.einsum_path_helpers.optimal_path(input_sets, output_set, idx_dict, memory_limit)[source]#

Compute all possible pair contractions, sieves the results based on memory_limit and returns the lowest cost path. This algorithm scales factorial with respect to the elements in the list input_sets.

Parameters:

  • input_sets (list) – List of sets that represent the lhs side of the einsum subscript

  • output_set (set) – Set that represents the rhs side of the overall einsum subscript

  • idx_dict (dictionary) – Dictionary of index sizes

  • memory_limit (int) – The maximum number of elements in a temporary array

Returns:

path (list) – The optimal contraction order within the memory limit constraint.

Examples

>>> isets = [set('abd'), set('ac'), set('bdc')]
>>> oset = set()
>>> idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
>>> optimal_path(isets, oset, idx_sizes, 5000)
[(0, 2), (0, 1)]
ivy.utils.einsum_path_helpers.parse_einsum_input(operands, subscripts=None)[source]#

Reproduction of einsum c side einsum parsing in python.

Returns:

  • input_strings (str) – Parsed input strings

  • output_string (str) – Parsed output string

  • operands (list of array_like) – The operands to use in the numpy contraction

Examples

The operand list is simplified to reduce printing:

>>> np.random.seed(123)
>>> a = np.random.rand(4, 4)
>>> b = np.random.rand(4, 4, 4)
>>> parse_einsum_input(('...a,...a->...', a, b))
('za,xza', 'xz', [a, b]) # may vary
>>> parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0]))
('za,xza', 'xz', [a, b]) # may vary
ivy.utils.einsum_path_helpers.parse_possible_contraction(positions, input_sets, output_set, idx_dict, memory_limit, path_cost, naive_cost)[source]#

Compute the cost (removed size + flops) and resultant indices for performing the contraction specified by positions.

Parameters:

  • positions (tuple of int) – The locations of the proposed tensors to contract.

  • input_sets (list of sets) – The indices found on each tensors.

  • output_set (set) – The output indices of the expression.

  • idx_dict (dict) – Mapping of each index to its size.

  • memory_limit (int) – The total allowed size for an intermediary tensor.

  • path_cost (int) – The contraction cost so far.

  • naive_cost (int) – The cost of the unoptimized expression.

Returns:

  • cost ((int, int)) – A tuple containing the size of any indices removed, and the flop cost.

  • positions (tuple of int) – The locations of the proposed tensors to contract.

  • new_input_sets (list of sets) – The resulting new list of indices if this proposed contraction is performed.

ivy.utils.einsum_path_helpers.update_other_results(results, best)[source]#

Update the positions and provisional input_sets of results based on performing the contraction result best. Remove any involving the tensors contracted.

Parameters:

  • results (list) – List of contraction results produced by _parse_possible_contraction.

  • best (list) – The best contraction of results i.e. the one that will be performed.

Returns:

mod_results (list) – The list of modified results, updated with outcome of best contraction.

This should have hopefully given you an overview of the einsum_path_helpers submodule, if you have any questions, please feel free to reach out on our discord!

On this page