API Documentation

opt_einsum.contract

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:

Name	Type	Description	Default
subscripts	Union[str, ArrayType]	Specifies the subscripts for summation.	required
*operands	Union[ArrayType, Collection[int]]	These are the arrays for the operation.	()
out	Optional[ArrayType]	A output array in which set the resulting output.	None
use_blas	bool	Do you use BLAS for valid operations, may use extra memory for more intermediates.	True
optimize	OptimizeKind	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.	True
memory_limit	_MemoryLimit	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.	None
backend	BackendType	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.	'auto'

Returns:

Type	Description
ArrayType	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.

Source code in opt_einsum/contract.py

def contract(
    subscripts: Union[str, ArrayType],
    *operands: Union[ArrayType, Collection[int]],
    out: Optional[ArrayType] = None,
    use_blas: bool = True,
    optimize: OptimizeKind = True,
    memory_limit: _MemoryLimit = None,
    backend: BackendType = "auto",
    **kwargs: Any,
) -> ArrayType:
    """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.
    """
    if (optimize is True) or (optimize is None):
        optimize = "auto"

    operands_list = [subscripts] + list(operands)

    # If no optimization, run pure einsum
    if optimize is False:
        return _einsum(*operands_list, out=out, **kwargs)

    # Grab non-einsum kwargs
    gen_expression = kwargs.pop("_gen_expression", False)
    constants_dict = kwargs.pop("_constants_dict", {})

    if gen_expression:
        full_str = operands_list[0]

    # Build the contraction list and operand
    contraction_list: ContractionListType
    operands, contraction_list = contract_path(  # type: ignore
        *operands_list, optimize=optimize, memory_limit=memory_limit, einsum_call=True, use_blas=use_blas
    )

    # check if performing contraction or just building expression
    if gen_expression:
        return ContractExpression(full_str, contraction_list, constants_dict, **kwargs)

    return _core_contract(operands, contraction_list, backend=backend, out=out, **kwargs)

opt_einsum.contract_path

Find a contraction order path, without performing the contraction.

Parameters:

Name	Type	Description	Default
subscripts	Any	Specifies the subscripts for summation.	required
*operands	Any	These are the arrays for the operation.	()
use_blas	bool	Do you use BLAS for valid operations, may use extra memory for more intermediates.	True
optimize	OptimizeKind	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' 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.	True
memory_limit	_MemoryLimit	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.	None
shapes	bool	Whether contract_path should assume arrays (the default) or array shapes have been supplied.	False

Returns:

Name	Type	Description
path	PathType	The optimized einsum contraciton path
PathInfo	PathInfo	A printable object containing various information about the path found.

Notes

The resulting path indicates which terms of the input contraction should be contracted first, the result of this contraction is then appended to the end of the contraction list.

Examples:

We can begin with a chain dot example. In this case, it is optimal to contract the b and c tensors represented by the first element of the path (1, 2). The resulting tensor is added to the end of the contraction and the remaining contraction, (0, 1), is then executed.

a = np.random.rand(2, 2)
b = np.random.rand(2, 5)
c = np.random.rand(5, 2)
path_info = opt_einsum.contract_path('ij,jk,kl->il', a, b, c)
print(path_info[0])
#> [(1, 2), (0, 1)]
print(path_info[1])
#>   Complete contraction:  ij,jk,kl->il
#>          Naive scaling:  4
#>      Optimized scaling:  3
#>       Naive FLOP count:  1.600e+02
#>   Optimized FLOP count:  5.600e+01
#>    Theoretical speedup:  2.857
#>   Largest intermediate:  4.000e+00 elements
#> -------------------------------------------------------------------------
#> scaling                  current                                remaining
#> -------------------------------------------------------------------------
#>    3                   kl,jk->jl                                ij,jl->il
#>    3                   jl,ij->il                                   il->il

A more complex index transformation example.

I = np.random.rand(10, 10, 10, 10)
C = np.random.rand(10, 10)
path_info = oe.contract_path('ea,fb,abcd,gc,hd->efgh', C, C, I, C, C)

print(path_info[0])
#> [(0, 2), (0, 3), (0, 2), (0, 1)]
print(path_info[1])
#>   Complete contraction:  ea,fb,abcd,gc,hd->efgh
#>          Naive scaling:  8
#>      Optimized scaling:  5
#>       Naive FLOP count:  8.000e+08
#>   Optimized FLOP count:  8.000e+05
#>    Theoretical speedup:  1000.000
#>   Largest intermediate:  1.000e+04 elements
#> --------------------------------------------------------------------------
#> scaling                  current                                remaining
#> --------------------------------------------------------------------------
#>    5               abcd,ea->bcde                      fb,gc,hd,bcde->efgh
#>    5               bcde,fb->cdef                         gc,hd,cdef->efgh
#>    5               cdef,gc->defg                            hd,defg->efgh
#>    5               defg,hd->efgh                               efgh->efgh

Source code in opt_einsum/contract.py

def contract_path(
    subscripts: Any,
    *operands: Any,
    use_blas: bool = True,
    optimize: OptimizeKind = True,
    memory_limit: _MemoryLimit = None,
    shapes: bool = False,
    **kwargs: Any,
) -> Tuple[PathType, PathInfo]:
    """Find a contraction order `path`, without performing the contraction.

    Parameters:
          subscripts: Specifies the subscripts for summation.
          *operands: These are the arrays for the operation.
          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'` 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.

          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.

          shapes: Whether ``contract_path`` should assume arrays (the default) or array shapes have been supplied.

    Returns:
          path: The optimized einsum contraciton path
          PathInfo: A printable object containing various information about the path found.

    Notes:
          The resulting path indicates which terms of the input contraction should be
          contracted first, the result of this contraction is then appended to the end of
          the contraction list.

    Examples:
          We can begin with a chain dot example. In this case, it is optimal to
          contract the b and c tensors represented by the first element of the path (1,
          2). The resulting tensor is added to the end of the contraction and the
          remaining contraction, `(0, 1)`, is then executed.

      ```python
      a = np.random.rand(2, 2)
      b = np.random.rand(2, 5)
      c = np.random.rand(5, 2)
      path_info = opt_einsum.contract_path('ij,jk,kl->il', a, b, c)
      print(path_info[0])
      #> [(1, 2), (0, 1)]
      print(path_info[1])
      #>   Complete contraction:  ij,jk,kl->il
      #>          Naive scaling:  4
      #>      Optimized scaling:  3
      #>       Naive FLOP count:  1.600e+02
      #>   Optimized FLOP count:  5.600e+01
      #>    Theoretical speedup:  2.857
      #>   Largest intermediate:  4.000e+00 elements
      #> -------------------------------------------------------------------------
      #> scaling                  current                                remaining
      #> -------------------------------------------------------------------------
      #>    3                   kl,jk->jl                                ij,jl->il
      #>    3                   jl,ij->il                                   il->il
      ```

      A more complex index transformation example.

      ```python
      I = np.random.rand(10, 10, 10, 10)
      C = np.random.rand(10, 10)
      path_info = oe.contract_path('ea,fb,abcd,gc,hd->efgh', C, C, I, C, C)

      print(path_info[0])
      #> [(0, 2), (0, 3), (0, 2), (0, 1)]
      print(path_info[1])
      #>   Complete contraction:  ea,fb,abcd,gc,hd->efgh
      #>          Naive scaling:  8
      #>      Optimized scaling:  5
      #>       Naive FLOP count:  8.000e+08
      #>   Optimized FLOP count:  8.000e+05
      #>    Theoretical speedup:  1000.000
      #>   Largest intermediate:  1.000e+04 elements
      #> --------------------------------------------------------------------------
      #> scaling                  current                                remaining
      #> --------------------------------------------------------------------------
      #>    5               abcd,ea->bcde                      fb,gc,hd,bcde->efgh
      #>    5               bcde,fb->cdef                         gc,hd,cdef->efgh
      #>    5               cdef,gc->defg                            hd,defg->efgh
      #>    5               defg,hd->efgh                               efgh->efgh
      ```
    """
    if (optimize is True) or (optimize is None):
        optimize = "auto"

    # Hidden option, only einsum should call this
    einsum_call_arg = kwargs.pop("einsum_call", False)
    if len(kwargs):
        raise TypeError(f"Did not understand the following kwargs: {kwargs.keys()}")

    # Python side parsing
    operands_ = [subscripts] + list(operands)
    input_subscripts, output_subscript, operands_prepped = parser.parse_einsum_input(operands_, shapes=shapes)

    # Build a few useful list and sets
    input_list = input_subscripts.split(",")
    input_sets = [frozenset(x) for x in input_list]
    if shapes:
        input_shapes = operands_prepped
    else:
        input_shapes = [parser.get_shape(x) for x in operands_prepped]
    output_set = frozenset(output_subscript)
    indices = frozenset(input_subscripts.replace(",", ""))

    # Get length of each unique dimension and ensure all dimensions are correct
    size_dict: Dict[str, int] = {}
    for tnum, term in enumerate(input_list):
        sh = input_shapes[tnum]

        if len(sh) != len(term):
            raise ValueError(
                f"Einstein sum subscript '{input_list[tnum]}' does not contain the "
                f"correct number of indices for operand {tnum}."
            )
        for cnum, char in enumerate(term):
            dim = int(sh[cnum])

            if char in size_dict:
                # For broadcasting cases we always want the largest dim size
                if size_dict[char] == 1:
                    size_dict[char] = dim
                elif dim not in (1, size_dict[char]):
                    raise ValueError(
                        f"Size of label '{char}' for operand {tnum} ({size_dict[char]}) does not match previous "
                        f"terms ({dim})."
                    )
            else:
                size_dict[char] = dim

    # Compute size of each input array plus the output array
    size_list = [helpers.compute_size_by_dict(term, size_dict) for term in input_list + [output_subscript]]
    memory_arg = _choose_memory_arg(memory_limit, size_list)

    num_ops = len(input_list)

    # Compute naive cost
    # This is not quite right, need to look into exactly how einsum does this
    # indices_in_input = input_subscripts.replace(',', '')
    inner_product = (sum(len(x) for x in input_sets) - len(indices)) > 0
    naive_cost = helpers.flop_count(indices, inner_product, num_ops, size_dict)

    # Compute the path
    if optimize is False:
        path_tuple: PathType = [tuple(range(num_ops))]
    elif not isinstance(optimize, (str, paths.PathOptimizer)):
        # Custom path supplied
        path_tuple = optimize  # type: ignore
    elif num_ops <= 2:
        # Nothing to be optimized
        path_tuple = [tuple(range(num_ops))]
    elif isinstance(optimize, paths.PathOptimizer):
        # Custom path optimizer supplied
        path_tuple = optimize(input_sets, output_set, size_dict, memory_arg)
    else:
        path_optimizer = paths.get_path_fn(optimize)
        path_tuple = path_optimizer(input_sets, output_set, size_dict, memory_arg)

    cost_list = []
    scale_list = []
    size_list = []
    contraction_list = []

    # Build contraction tuple (positions, gemm, einsum_str, remaining)
    for cnum, contract_inds in enumerate(path_tuple):
        # Make sure we remove inds from right to left
        contract_inds = tuple(sorted(contract_inds, reverse=True))

        contract_tuple = helpers.find_contraction(contract_inds, input_sets, output_set)
        out_inds, input_sets, idx_removed, idx_contract = contract_tuple

        # Compute cost, scale, and size
        cost = helpers.flop_count(idx_contract, bool(idx_removed), len(contract_inds), size_dict)
        cost_list.append(cost)
        scale_list.append(len(idx_contract))
        size_list.append(helpers.compute_size_by_dict(out_inds, size_dict))

        tmp_inputs = [input_list.pop(x) for x in contract_inds]
        tmp_shapes = [input_shapes.pop(x) for x in contract_inds]

        if use_blas:
            do_blas = blas.can_blas(tmp_inputs, "".join(out_inds), idx_removed, tmp_shapes)
        else:
            do_blas = False

        # Last contraction
        if (cnum - len(path_tuple)) == -1:
            idx_result = output_subscript
        else:
            # use tensordot order to minimize transpositions
            all_input_inds = "".join(tmp_inputs)
            idx_result = "".join(sorted(out_inds, key=all_input_inds.find))

        shp_result = parser.find_output_shape(tmp_inputs, tmp_shapes, idx_result)

        input_list.append(idx_result)
        input_shapes.append(shp_result)

        einsum_str = ",".join(tmp_inputs) + "->" + idx_result

        # for large expressions saving the remaining terms at each step can
        # incur a large memory footprint - and also be messy to print
        if len(input_list) <= 20:
            remaining: Optional[Tuple[str, ...]] = tuple(input_list)
        else:
            remaining = None

        contraction = (contract_inds, idx_removed, einsum_str, remaining, do_blas)
        contraction_list.append(contraction)

    opt_cost = sum(cost_list)

    if einsum_call_arg:
        return operands_prepped, contraction_list  # type: ignore

    path_print = PathInfo(
        contraction_list,
        input_subscripts,
        output_subscript,
        indices,
        path_tuple,
        scale_list,
        naive_cost,
        opt_cost,
        size_list,
        size_dict,
    )

    return path_tuple, path_print

opt_einsum.contract_expression

Generate a reusable expression for a given contraction with specific shapes, which can, for example, be cached.

Parameters:

subscripts: Specifies the subscripts for summation.
shapes: Shapes of the arrays to optimize the contraction for.
constants: The indices of any constant arguments in `shapes`, in which case the
    actual array should be supplied at that position rather than just a
    shape. If these are specified, then constant parts of the contraction
    between calls will be reused. Additionally, if a GPU-enabled backend is
    used for example, then the constant tensors will be kept on the GPU,
    minimizing transfers.
kwargs: Passed on to `contract_path` or `einsum`. See `contract`.

Returns:

Type	Description
ContractExpression	Callable with signature expr(*arrays, out=None, backend='numpy') where the array's shapes should match shapes.

Notes

The out keyword argument should be supplied to the generated expression rather than this function. The backend keyword argument should also be supplied to the generated expression. If numpy arrays are supplied, if possible they will be converted to and back from the correct backend array type. The generated expression will work with any arrays which have the same rank (number of dimensions) as the original shapes, however, if the actual sizes are different, the expression may no longer be optimal. Constant operations will be computed upon the first call with a particular backend, then subsequently reused.

Examples: Basic usage:

expr = contract_expression("ab,bc->ac", (3, 4), (4, 5))
a, b = np.random.rand(3, 4), np.random.rand(4, 5)
c = expr(a, b)
np.allclose(c, a @ b)
#> True

Supply a as a constant:

expr = contract_expression("ab,bc->ac", a, (4, 5), constants=[0])
expr
#> <ContractExpression('[ab],bc->ac', constants=[0])>

c = expr(b)
np.allclose(c, a @ b)
#> True

Source code in opt_einsum/contract.py

def contract_expression(
    subscripts: Union[str, ArrayType, TensorShapeType],
    *shapes: Union[ArrayType, TensorShapeType, Collection[int]],
    constants: Union[Collection[int], None] = None,
    use_blas: bool = True,
    optimize: OptimizeKind = True,
    memory_limit: _MemoryLimit = None,
    **kwargs: Any,
) -> ContractExpression:
    """Generate a reusable expression for a given contraction with
    specific shapes, which can, for example, be cached.

    Parameters:

        subscripts: Specifies the subscripts for summation.
        shapes: Shapes of the arrays to optimize the contraction for.
        constants: The indices of any constant arguments in `shapes`, in which case the
            actual array should be supplied at that position rather than just a
            shape. If these are specified, then constant parts of the contraction
            between calls will be reused. Additionally, if a GPU-enabled backend is
            used for example, then the constant tensors will be kept on the GPU,
            minimizing transfers.
        kwargs: Passed on to `contract_path` or `einsum`. See `contract`.

    Returns:
        Callable with signature `expr(*arrays, out=None, backend='numpy')` where the array's shapes should match `shapes`.

    Notes:
        The `out` keyword argument should be supplied to the generated expression
        rather than this function.
        The `backend` keyword argument should also be supplied to the generated
        expression. If numpy arrays are supplied, if possible they will be
        converted to and back from the correct backend array type.
        The generated expression will work with any arrays which have
        the same rank (number of dimensions) as the original shapes, however, if
        the actual sizes are different, the expression may no longer be optimal.
        Constant operations will be computed upon the first call with a particular
        backend, then subsequently reused.

    Examples:
    Basic usage:

    ```python
    expr = contract_expression("ab,bc->ac", (3, 4), (4, 5))
    a, b = np.random.rand(3, 4), np.random.rand(4, 5)
    c = expr(a, b)
    np.allclose(c, a @ b)
    #> True
    ```

    Supply `a` as a constant:

    ```python
    expr = contract_expression("ab,bc->ac", a, (4, 5), constants=[0])
    expr
    #> <ContractExpression('[ab],bc->ac', constants=[0])>

    c = expr(b)
    np.allclose(c, a @ b)
    #> True
    ```

    """
    if not optimize:
        raise ValueError("Can only generate expressions for optimized contractions.")

    for arg in ("out", "backend"):
        if kwargs.get(arg, None) is not None:
            raise ValueError(
                f"'{arg}' should only be specified when calling a " "`ContractExpression`, not when building it."
            )

    if not isinstance(subscripts, str):
        subscripts, shapes = parser.convert_interleaved_input((subscripts,) + shapes)

    kwargs["_gen_expression"] = True

    # build dict of constant indices mapped to arrays
    constants = constants or ()
    constants_dict = {i: shapes[i] for i in constants}
    kwargs["_constants_dict"] = constants_dict

    # apart from constant arguments, make dummy arrays
    dummy_arrays = [s if i in constants else shape_only(s) for i, s in enumerate(shapes)]  # type: ignore

    return contract(
        subscripts, *dummy_arrays, use_blas=use_blas, optimize=optimize, memory_limit=memory_limit, **kwargs
    )

opt_einsum.contract.ContractExpression

Helper class for storing an explicit contraction_list which can then be repeatedly called solely with the array arguments.

Source code in opt_einsum/contract.py

class ContractExpression:
    """Helper class for storing an explicit ``contraction_list`` which can
    then be repeatedly called solely with the array arguments.
    """

    def __init__(
        self,
        contraction: str,
        contraction_list: ContractionListType,
        constants_dict: Dict[int, ArrayType],
        **kwargs: Any,
    ):
        self.contraction = format_const_einsum_str(contraction, constants_dict.keys())
        self.contraction_list = contraction_list
        self.kwargs = kwargs

        # need to know _full_num_args to parse constants with, and num_args to call with
        self._full_num_args = contraction.count(",") + 1
        self.num_args = self._full_num_args - len(constants_dict)

        # likewise need to know full contraction list
        self._full_contraction_list = contraction_list

        self._constants_dict = constants_dict
        self._evaluated_constants: Dict[str, Any] = {}
        self._backend_expressions: Dict[str, Any] = {}

    def evaluate_constants(self, backend: Optional[str] = "auto") -> None:
        """Convert any constant operands to the correct backend form, and
        perform as many contractions as possible to create a new list of
        operands, stored in ``self._evaluated_constants[backend]``. This also
        makes sure ``self.contraction_list`` only contains the remaining,
        non-const operations.
        """
        # prepare a list of operands, with `None` for non-consts
        tmp_const_ops = [self._constants_dict.get(i, None) for i in range(self._full_num_args)]
        backend = parse_backend(tmp_const_ops, backend)

        # get the new list of operands with constant operations performed, and remaining contractions
        try:
            new_ops, new_contraction_list = backends.evaluate_constants(backend, tmp_const_ops, self)
        except KeyError:
            new_ops, new_contraction_list = self(*tmp_const_ops, backend=backend, evaluate_constants=True)

        self._evaluated_constants[backend] = new_ops
        self.contraction_list = new_contraction_list

    def _get_evaluated_constants(self, backend: str) -> List[Optional[ArrayType]]:
        """Retrieve or generate the cached list of constant operators (mixed
        in with None representing non-consts) and the remaining contraction
        list.
        """
        try:
            return self._evaluated_constants[backend]
        except KeyError:
            self.evaluate_constants(backend)
            return self._evaluated_constants[backend]

    def _get_backend_expression(self, arrays: Sequence[ArrayType], backend: str) -> Any:
        try:
            return self._backend_expressions[backend]
        except KeyError:
            fn = backends.build_expression(backend, arrays, self)
            self._backend_expressions[backend] = fn
            return fn

    def _contract(
        self,
        arrays: Sequence[ArrayType],
        out: Optional[ArrayType] = None,
        backend: Optional[str] = "auto",
        evaluate_constants: bool = False,
    ) -> ArrayType:
        """The normal, core contraction."""
        contraction_list = self._full_contraction_list if evaluate_constants else self.contraction_list

        return _core_contract(
            list(arrays),
            contraction_list,
            out=out,
            backend=backend,
            evaluate_constants=evaluate_constants,
            **self.kwargs,
        )

    def _contract_with_conversion(
        self,
        arrays: Sequence[ArrayType],
        out: Optional[ArrayType],
        backend: str,
        evaluate_constants: bool = False,
    ) -> ArrayType:
        """Special contraction, i.e., contraction with a different backend
        but converting to and from that backend. Retrieves or generates a
        cached expression using ``arrays`` as templates, then calls it
        with ``arrays``.

        If ``evaluate_constants=True``, perform a partial contraction that
        prepares the constant tensors and operations with the right backend.
        """
        # convert consts to correct type & find reduced contraction list
        if evaluate_constants:
            return backends.evaluate_constants(backend, arrays, self)

        result = self._get_backend_expression(arrays, backend)(*arrays)

        if out is not None:
            out[()] = result
            return out

        return result

    def __call__(
        self,
        *arrays: ArrayType,
        out: Union[None, ArrayType] = None,
        backend: str = "auto",
        evaluate_constants: bool = False,
    ) -> ArrayType:
        """Evaluate this expression with a set of arrays.

        Parameters:
            arrays: The arrays to supply as input to the expression.
            out: If specified, output the result into this array.
            backend: Perform the contraction with this backend library. If numpy arrays
                are supplied then try to convert them to and from the correct
                backend array type.
            evaluate_constants: Pre-evaluates constants with the appropriate backend.

        Returns:
            The contracted result.
        """
        backend = parse_backend(arrays, backend)

        correct_num_args = self._full_num_args if evaluate_constants else self.num_args

        if len(arrays) != correct_num_args:
            raise ValueError(
                f"This `ContractExpression` takes exactly {self.num_args} array arguments "
                f"but received {len(arrays)}."
            )

        if self._constants_dict and not evaluate_constants:
            # fill in the missing non-constant terms with newly supplied arrays
            ops_var, ops_const = iter(arrays), self._get_evaluated_constants(backend)
            ops: Sequence[ArrayType] = [next(ops_var) if op is None else op for op in ops_const]
        else:
            ops = arrays

        try:
            # Check if the backend requires special preparation / calling
            #   but also ignore non-numpy arrays -> assume user wants same type back
            if backends.has_backend(backend) and all(infer_backend(x) == "numpy" for x in arrays):
                return self._contract_with_conversion(ops, out, backend, evaluate_constants=evaluate_constants)

            return self._contract(ops, out=out, backend=backend, evaluate_constants=evaluate_constants)

        except ValueError as err:
            original_msg = str(err.args) if err.args else ""
            msg = (
                "Internal error while evaluating `ContractExpression`. Note that few checks are performed"
                " - the number and rank of the array arguments must match the original expression. "
                f"The internal error was: '{original_msg}'",
            )
            err.args = msg
            raise

    def __repr__(self) -> str:
        if self._constants_dict:
            constants_repr = f", constants={sorted(self._constants_dict)}"
        else:
            constants_repr = ""
        return f"<ContractExpression('{self.contraction}'{constants_repr})>"

    def __str__(self) -> str:
        s = [self.__repr__()]
        for i, c in enumerate(self.contraction_list):
            s.append(f"\n  {i + 1}.  ")
            s.append(f"'{c[2]}'" + (f" [{c[-1]}]" if c[-1] else ""))
            s.append(f"\neinsum_kwargs={self.kwargs}")
        return "".join(s)

__call__(*arrays, out=None, backend='auto', evaluate_constants=False)

Evaluate this expression with a set of arrays.

Parameters:

Name	Type	Description	Default
arrays	ArrayType	The arrays to supply as input to the expression.	()
out	Union[None, ArrayType]	If specified, output the result into this array.	None
backend	str	Perform the contraction with this backend library. If numpy arrays are supplied then try to convert them to and from the correct backend array type.	'auto'
evaluate_constants	bool	Pre-evaluates constants with the appropriate backend.	False

Returns:

Type	Description
ArrayType	The contracted result.

Source code in opt_einsum/contract.py

def __call__(
    self,
    *arrays: ArrayType,
    out: Union[None, ArrayType] = None,
    backend: str = "auto",
    evaluate_constants: bool = False,
) -> ArrayType:
    """Evaluate this expression with a set of arrays.

    Parameters:
        arrays: The arrays to supply as input to the expression.
        out: If specified, output the result into this array.
        backend: Perform the contraction with this backend library. If numpy arrays
            are supplied then try to convert them to and from the correct
            backend array type.
        evaluate_constants: Pre-evaluates constants with the appropriate backend.

    Returns:
        The contracted result.
    """
    backend = parse_backend(arrays, backend)

    correct_num_args = self._full_num_args if evaluate_constants else self.num_args

    if len(arrays) != correct_num_args:
        raise ValueError(
            f"This `ContractExpression` takes exactly {self.num_args} array arguments "
            f"but received {len(arrays)}."
        )

    if self._constants_dict and not evaluate_constants:
        # fill in the missing non-constant terms with newly supplied arrays
        ops_var, ops_const = iter(arrays), self._get_evaluated_constants(backend)
        ops: Sequence[ArrayType] = [next(ops_var) if op is None else op for op in ops_const]
    else:
        ops = arrays

    try:
        # Check if the backend requires special preparation / calling
        #   but also ignore non-numpy arrays -> assume user wants same type back
        if backends.has_backend(backend) and all(infer_backend(x) == "numpy" for x in arrays):
            return self._contract_with_conversion(ops, out, backend, evaluate_constants=evaluate_constants)

        return self._contract(ops, out=out, backend=backend, evaluate_constants=evaluate_constants)

    except ValueError as err:
        original_msg = str(err.args) if err.args else ""
        msg = (
            "Internal error while evaluating `ContractExpression`. Note that few checks are performed"
            " - the number and rank of the array arguments must match the original expression. "
            f"The internal error was: '{original_msg}'",
        )
        err.args = msg
        raise

evaluate_constants(backend='auto')

Convert any constant operands to the correct backend form, and perform as many contractions as possible to create a new list of operands, stored in self._evaluated_constants[backend]. This also makes sure self.contraction_list only contains the remaining, non-const operations.

Source code in opt_einsum/contract.py

def evaluate_constants(self, backend: Optional[str] = "auto") -> None:
    """Convert any constant operands to the correct backend form, and
    perform as many contractions as possible to create a new list of
    operands, stored in ``self._evaluated_constants[backend]``. This also
    makes sure ``self.contraction_list`` only contains the remaining,
    non-const operations.
    """
    # prepare a list of operands, with `None` for non-consts
    tmp_const_ops = [self._constants_dict.get(i, None) for i in range(self._full_num_args)]
    backend = parse_backend(tmp_const_ops, backend)

    # get the new list of operands with constant operations performed, and remaining contractions
    try:
        new_ops, new_contraction_list = backends.evaluate_constants(backend, tmp_const_ops, self)
    except KeyError:
        new_ops, new_contraction_list = self(*tmp_const_ops, backend=backend, evaluate_constants=True)

    self._evaluated_constants[backend] = new_ops
    self.contraction_list = new_contraction_list

opt_einsum.contract.PathInfo

A printable object to contain information about a contraction path.

Source code in opt_einsum/contract.py

class PathInfo:
    """A printable object to contain information about a contraction path."""

    def __init__(
        self,
        contraction_list: ContractionListType,
        input_subscripts: str,
        output_subscript: str,
        indices: ArrayIndexType,
        path: PathType,
        scale_list: Sequence[int],
        naive_cost: int,
        opt_cost: int,
        size_list: Sequence[int],
        size_dict: Dict[str, int],
    ):
        self.contraction_list = contraction_list
        self.input_subscripts = input_subscripts
        self.output_subscript = output_subscript
        self.path = path
        self.indices = indices
        self.scale_list = scale_list
        self.naive_cost = Decimal(naive_cost)
        self.opt_cost = Decimal(opt_cost)
        self.speedup = self.naive_cost / max(self.opt_cost, Decimal(1))
        self.size_list = size_list
        self.size_dict = size_dict

        self.shapes = [tuple(size_dict[k] for k in ks) for ks in input_subscripts.split(",")]
        self.eq = f"{input_subscripts}->{output_subscript}"
        self.largest_intermediate = Decimal(max(size_list, default=1))

    def __repr__(self) -> str:
        # Return the path along with a nice string representation
        header = ("scaling", "BLAS", "current", "remaining")

        path_print = [
            f"  Complete contraction:  {self.eq}\n",
            f"         Naive scaling:  {len(self.indices)}\n",
            f"     Optimized scaling:  {max(self.scale_list, default=0)}\n",
            f"      Naive FLOP count:  {self.naive_cost:.3e}\n",
            f"  Optimized FLOP count:  {self.opt_cost:.3e}\n",
            f"   Theoretical speedup:  {self.speedup:.3e}\n",
            f"  Largest intermediate:  {self.largest_intermediate:.3e} elements\n",
            "-" * 80 + "\n",
            "{:>6} {:>11} {:>22} {:>37}\n".format(*header),
            "-" * 80,
        ]

        for n, contraction in enumerate(self.contraction_list):
            _, _, einsum_str, remaining, do_blas = contraction

            if remaining is not None:
                remaining_str = ",".join(remaining) + "->" + self.output_subscript
            else:
                remaining_str = "..."
            size_remaining = max(0, 56 - max(22, len(einsum_str)))

            path_run = (
                self.scale_list[n],
                do_blas,
                einsum_str,
                remaining_str,
                size_remaining,
            )
            path_print.append("\n{:>4} {:>14} {:>22}    {:>{}}".format(*path_run))

        return "".join(path_print)

opt_einsum.get_symbol

Get the symbol corresponding to int i - runs through the usual 52 letters before resorting to unicode characters, starting at chr(192) and skipping surrogates.

Examples:

get_symbol(2)
#> 'c'

get_symbol(200)
#> 'Ŕ'

get_symbol(20000)
#> '京'

Source code in opt_einsum/parser.py

def get_symbol(i: int) -> str:
    """Get the symbol corresponding to int ``i`` - runs through the usual 52
    letters before resorting to unicode characters, starting at ``chr(192)`` and skipping surrogates.

    **Examples:**

    ```python
    get_symbol(2)
    #> 'c'

    get_symbol(200)
    #> 'Ŕ'

    get_symbol(20000)
    #> '京'
    ```
    """
    if i < 52:
        return _einsum_symbols_base[i]
    elif i >= 55296:
        # Skip chr(57343) - chr(55296) as surrogates
        return chr(i + 2048)
    else:
        return chr(i + 140)

opt_einsum.shared_intermediates

Context in which contract intermediate results are shared.

Note that intermediate computations will not be garbage collected until 1. this context exits, and 2. the yielded cache is garbage collected (if it was captured).

Parameters:

• cache - (dict) If specified, a user-stored dict in which intermediate results will be stored. This can be used to interleave sharing contexts.

Returns:

• cache - (dict) A dictionary in which sharing results are stored. If ignored, sharing results will be garbage collected when this context is exited. This dict can be passed to another context to resume sharing.

Source code in opt_einsum/sharing.py

@contextlib.contextmanager
def shared_intermediates(
    cache: Optional[CacheType] = None,
) -> Generator[CacheType, None, None]:
    """Context in which contract intermediate results are shared.

    Note that intermediate computations will not be garbage collected until
    1. this context exits, and
    2. the yielded cache is garbage collected (if it was captured).

    **Parameters:**

    - **cache** - *(dict)* If specified, a user-stored dict in which intermediate results will be stored. This can be used to interleave sharing contexts.

    **Returns:**

    - **cache** - *(dict)* A dictionary in which sharing results are stored. If ignored,
        sharing results will be garbage collected when this context is
        exited. This dict can be passed to another context to resume
        sharing.
    """
    if cache is None:
        cache = {}
    _add_sharing_cache(cache)
    try:
        yield cache
    finally:
        _remove_sharing_cache()

opt_einsum.paths.optimal

Computes all possible pair contractions in a depth-first recursive manner, sieving results based on memory_limit and the best path found so far.

Parameters:

Name	Type	Description	Default
inputs	List[ArrayIndexType]	List of sets that represent the lhs side of the einsum subscript.	required
output	ArrayIndexType	Set that represents the rhs side of the overall einsum subscript.	required
size_dict	Dict[str, int]	Dictionary of index sizes.	required
memory_limit	Optional[int]	The maximum number of elements in a temporary array.	None

Returns:

Name	Type	Description
path	PathType	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(isets, oset, idx_sizes, 5000)
#> [(0, 2), (0, 1)]

Source code in opt_einsum/paths.py

def optimal(
    inputs: List[ArrayIndexType],
    output: ArrayIndexType,
    size_dict: Dict[str, int],
    memory_limit: Optional[int] = None,
) -> PathType:
    """Computes all possible pair contractions in a depth-first recursive manner,
    sieving results based on `memory_limit` and the best path found so far.

    Parameters:
        inputs: List of sets that represent the lhs side of the einsum subscript.
        output: Set that represents the rhs side of the overall einsum subscript.
        size_dict: Dictionary of index sizes.
        memory_limit: The maximum number of elements in a temporary array.

    Returns:
        path: The optimal contraction order within the memory limit constraint.

    Examples:
    ```python
    isets = [set('abd'), set('ac'), set('bdc')]
    oset = set('')
    idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
    optimal(isets, oset, idx_sizes, 5000)
    #> [(0, 2), (0, 1)]
    ```
    """
    inputs_set = tuple(map(frozenset, inputs))
    output_set = frozenset(output)

    best_flops = {"flops": float("inf")}
    best_ssa_path = {"ssa_path": (tuple(range(len(inputs))),)}
    size_cache: Dict[FrozenSet[str], int] = {}
    result_cache: Dict[Tuple[ArrayIndexType, ArrayIndexType], Tuple[FrozenSet[str], int]] = {}

    def _optimal_iterate(path, remaining, inputs, flops):
        # reached end of path (only ever get here if flops is best found so far)
        if len(remaining) == 1:
            best_flops["flops"] = flops
            best_ssa_path["ssa_path"] = path
            return

        # check all possible remaining paths
        for i, j in itertools.combinations(remaining, 2):
            if i > j:
                i, j = j, i
            key = (inputs[i], inputs[j])
            try:
                k12, flops12 = result_cache[key]
            except KeyError:
                k12, flops12 = result_cache[key] = calc_k12_flops(inputs, output_set, remaining, i, j, size_dict)

            # sieve based on current best flops
            new_flops = flops + flops12
            if new_flops >= best_flops["flops"]:
                continue

            # sieve based on memory limit
            if memory_limit not in _UNLIMITED_MEM:
                try:
                    size12 = size_cache[k12]
                except KeyError:
                    size12 = size_cache[k12] = compute_size_by_dict(k12, size_dict)

                # possibly terminate this path with an all-terms einsum
                if size12 > memory_limit:
                    new_flops = flops + _compute_oversize_flops(inputs, remaining, output_set, size_dict)
                    if new_flops < best_flops["flops"]:
                        best_flops["flops"] = new_flops
                        best_ssa_path["ssa_path"] = path + (tuple(remaining),)
                    continue

            # add contraction and recurse into all remaining
            _optimal_iterate(
                path=path + ((i, j),),
                inputs=inputs + (k12,),
                remaining=remaining - {i, j} | {len(inputs)},
                flops=new_flops,
            )

    _optimal_iterate(path=(), inputs=inputs_set, remaining=set(range(len(inputs))), flops=0)

    return ssa_to_linear(best_ssa_path["ssa_path"])

opt_einsum.paths.greedy

Finds the path by a three stage algorithm:

1. Eagerly compute Hadamard products.
2. Greedily compute contractions to maximize removed_size
3. Greedily compute outer products.

This algorithm scales quadratically with respect to the maximum number of elements sharing a common dim.

Parameters:

Name	Type	Description	Default
inputs	List[ArrayIndexType]	List of sets that represent the lhs side of the einsum subscript	required
output	ArrayIndexType	Set that represents the rhs side of the overall einsum subscript	required
size_dict	Dict[str, int]	Dictionary of index sizes	required
memory_limit	Optional[int]	The maximum number of elements in a temporary array	None
choose_fn	Any	A function that chooses which contraction to perform from the queue	None
cost_fn	str	A function that assigns a potential contraction a cost.	'memory-removed'

Returns:

Name	Type	Description
path	PathType	The contraction order (a list of tuples of ints).

Examples:

isets = [set('abd'), set('ac'), set('bdc')]
oset = set('')
idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
greedy(isets, oset, idx_sizes)
#> [(0, 2), (0, 1)]

Source code in opt_einsum/paths.py

def greedy(
    inputs: List[ArrayIndexType],
    output: ArrayIndexType,
    size_dict: Dict[str, int],
    memory_limit: Optional[int] = None,
    choose_fn: Any = None,
    cost_fn: str = "memory-removed",
) -> PathType:
    """Finds the path by a three stage algorithm:

    1. Eagerly compute Hadamard products.
    2. Greedily compute contractions to maximize `removed_size`
    3. Greedily compute outer products.

    This algorithm scales quadratically with respect to the
    maximum number of elements sharing a common dim.

    Parameters:
        inputs: List of sets that represent the lhs side of the einsum subscript
        output: Set that represents the rhs side of the overall einsum subscript
        size_dict: Dictionary of index sizes
        memory_limit: The maximum number of elements in a temporary array
        choose_fn: A function that chooses which contraction to perform from the queue
        cost_fn: A function that assigns a potential contraction a cost.

    Returns:
        path: The contraction order (a list of tuples of ints).

    Examples:
        ```python
        isets = [set('abd'), set('ac'), set('bdc')]
        oset = set('')
        idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
        greedy(isets, oset, idx_sizes)
        #> [(0, 2), (0, 1)]
        ```
    """
    if memory_limit not in _UNLIMITED_MEM:
        return branch(inputs, output, size_dict, memory_limit, nbranch=1, cost_fn=cost_fn)  # type: ignore

    ssa_path = ssa_greedy_optimize(inputs, output, size_dict, cost_fn=cost_fn, choose_fn=choose_fn)
    return ssa_to_linear(ssa_path)

opt_einsum.paths.branch

Source code in opt_einsum/paths.py

def branch(
    inputs: List[ArrayIndexType],
    output: ArrayIndexType,
    size_dict: Dict[str, int],
    memory_limit: Optional[int] = None,
    nbranch: Optional[int] = None,
    cutoff_flops_factor: int = 4,
    minimize: str = "flops",
    cost_fn: str = "memory-removed",
) -> PathType:
    optimizer = BranchBound(
        nbranch=nbranch, cutoff_flops_factor=cutoff_flops_factor, minimize=minimize, cost_fn=cost_fn
    )
    return optimizer(inputs, output, size_dict, memory_limit)

opt_einsum.paths.PathOptimizer

Base class for different path optimizers to inherit from.

Subclassed optimizers should define a call method with signature:

def __call__(self, inputs: List[ArrayIndexType], output: ArrayIndexType, size_dict: dict[str, int], memory_limit: int | None = None) -> list[tuple[int, ...]]:
    \"\"\"
    Parameters:
        inputs: The indices of each input array.
        outputs: The output indices
        size_dict: The size of each index
        memory_limit: If given, the maximum allowed memory.
    \"\"\"
    # ... compute path here ...
    return path

where path is a list of int-tuples specifying a contraction order.

Source code in opt_einsum/paths.py

class PathOptimizer:
    r"""Base class for different path optimizers to inherit from.

    Subclassed optimizers should define a call method with signature:

    ```python
    def __call__(self, inputs: List[ArrayIndexType], output: ArrayIndexType, size_dict: dict[str, int], memory_limit: int | None = None) -> list[tuple[int, ...]]:
        \"\"\"
        Parameters:
            inputs: The indices of each input array.
            outputs: The output indices
            size_dict: The size of each index
            memory_limit: If given, the maximum allowed memory.
        \"\"\"
        # ... compute path here ...
        return path
    ```

    where `path` is a list of int-tuples specifying a contraction order.
    """

    def _check_args_against_first_call(
        self,
        inputs: List[ArrayIndexType],
        output: ArrayIndexType,
        size_dict: Dict[str, int],
    ) -> None:
        """Utility that stateful optimizers can use to ensure they are not
        called with different contractions across separate runs.
        """
        args = (inputs, output, size_dict)
        if not hasattr(self, "_first_call_args"):
            # simply set the attribute as currently there is no global PathOptimizer init
            self._first_call_args = args
        elif args != self._first_call_args:
            raise ValueError(
                "The arguments specifying the contraction that this path optimizer "
                "instance was called with have changed - try creating a new instance."
            )

    def __call__(
        self,
        inputs: List[ArrayIndexType],
        output: ArrayIndexType,
        size_dict: Dict[str, int],
        memory_limit: Optional[int] = None,
    ) -> PathType:
        raise NotImplementedError

opt_einsum.paths.BranchBound

Bases: PathOptimizer

Source code in opt_einsum/paths.py

class BranchBound(PathOptimizer):
    def __init__(
        self,
        nbranch: Optional[int] = None,
        cutoff_flops_factor: int = 4,
        minimize: str = "flops",
        cost_fn: str = "memory-removed",
    ):
        """Explores possible pair contractions in a depth-first recursive manner like
        the `optimal` approach, but with extra heuristic early pruning of branches
        as well sieving by `memory_limit` and the best path found so far.


        Parameters:
            nbranch: How many branches to explore at each contraction step. If None, explore
                all possible branches. If an integer, branch into this many paths at
                each step. Defaults to None.
            cutoff_flops_factor: If at any point, a path is doing this much worse than the best path
                found so far was, terminate it. The larger this is made, the more paths
                will be fully explored and the slower the algorithm. Defaults to 4.
            minimize: Whether to optimize the path with regard primarily to the total
                estimated flop-count, or the size of the largest intermediate. The
                option not chosen will still be used as a secondary criterion.
            cost_fn: A function that returns a heuristic 'cost' of a potential contraction
                with which to sort candidates. Should have signature
                `cost_fn(size12, size1, size2, k12, k1, k2)`.
        """
        if (nbranch is not None) and nbranch < 1:
            raise ValueError(f"The number of branches must be at least one, `nbranch={nbranch}`.")

        self.nbranch = nbranch
        self.cutoff_flops_factor = cutoff_flops_factor
        self.minimize = minimize
        self.cost_fn: Any = _COST_FNS.get(cost_fn, cost_fn)

        self.better = get_better_fn(minimize)
        self.best: Dict[str, Any] = {"flops": float("inf"), "size": float("inf")}
        self.best_progress: Dict[int, float] = defaultdict(lambda: float("inf"))

    @property
    def path(self) -> PathType:
        return ssa_to_linear(self.best["ssa_path"])

    def __call__(
        self,
        inputs_: List[ArrayIndexType],
        output_: ArrayIndexType,
        size_dict: Dict[str, int],
        memory_limit: Optional[int] = None,
    ) -> PathType:
        """Parameters:
            inputs_: List of sets that represent the lhs side of the einsum subscript
            output_: Set that represents the rhs side of the overall einsum subscript
            size_dict: Dictionary of index sizes
            memory_limit: The maximum number of elements in a temporary array.

        Returns:
            path: The contraction order within the memory limit constraint.

        Examples:
        ```python
        isets = [set('abd'), set('ac'), set('bdc')]
        oset = set('')
        idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
        optimal(isets, oset, idx_sizes, 5000)
        #> [(0, 2), (0, 1)]
        """
        self._check_args_against_first_call(inputs_, output_, size_dict)

        inputs: Tuple[FrozenSet[str]] = tuple(map(frozenset, inputs_))  # type: ignore
        output: FrozenSet[str] = frozenset(output_)

        size_cache = {k: compute_size_by_dict(k, size_dict) for k in inputs}
        result_cache: Dict[Tuple[FrozenSet[str], FrozenSet[str]], Tuple[FrozenSet[str], int]] = {}

        def _branch_iterate(path, inputs, remaining, flops, size):
            # reached end of path (only ever get here if flops is best found so far)
            if len(remaining) == 1:
                self.best["size"] = size
                self.best["flops"] = flops
                self.best["ssa_path"] = path
                return

            def _assess_candidate(k1: FrozenSet[str], k2: FrozenSet[str], i: int, j: int) -> Any:
                # find resulting indices and flops
                try:
                    k12, flops12 = result_cache[k1, k2]
                except KeyError:
                    k12, flops12 = result_cache[k1, k2] = calc_k12_flops(inputs, output, remaining, i, j, size_dict)

                try:
                    size12 = size_cache[k12]
                except KeyError:
                    size12 = size_cache[k12] = compute_size_by_dict(k12, size_dict)

                new_flops = flops + flops12
                new_size = max(size, size12)

                # sieve based on current best i.e. check flops and size still better
                if not self.better(new_flops, new_size, self.best["flops"], self.best["size"]):
                    return None

                # compare to how the best method was doing as this point
                if new_flops < self.best_progress[len(inputs)]:
                    self.best_progress[len(inputs)] = new_flops
                # sieve based on current progress relative to best
                elif new_flops > self.cutoff_flops_factor * self.best_progress[len(inputs)]:
                    return None

                # sieve based on memory limit
                if (memory_limit not in _UNLIMITED_MEM) and (size12 > memory_limit):  # type: ignore
                    # terminate path here, but check all-terms contract first
                    new_flops = flops + _compute_oversize_flops(inputs, remaining, output_, size_dict)
                    if new_flops < self.best["flops"]:
                        self.best["flops"] = new_flops
                        self.best["ssa_path"] = path + (tuple(remaining),)
                    return None

                # set cost heuristic in order to locally sort possible contractions
                size1, size2 = size_cache[inputs[i]], size_cache[inputs[j]]
                cost = self.cost_fn(size12, size1, size2, k12, k1, k2)

                return cost, flops12, new_flops, new_size, (i, j), k12

            # check all possible remaining paths
            candidates = []
            for i, j in itertools.combinations(remaining, 2):
                if i > j:
                    i, j = j, i
                k1, k2 = inputs[i], inputs[j]

                # initially ignore outer products
                if k1.isdisjoint(k2):
                    continue

                candidate = _assess_candidate(k1, k2, i, j)
                if candidate:
                    heapq.heappush(candidates, candidate)

            # assess outer products if nothing left
            if not candidates:
                for i, j in itertools.combinations(remaining, 2):
                    if i > j:
                        i, j = j, i
                    k1, k2 = inputs[i], inputs[j]
                    candidate = _assess_candidate(k1, k2, i, j)
                    if candidate:
                        heapq.heappush(candidates, candidate)

            # recurse into all or some of the best candidate contractions
            bi = 0
            while (self.nbranch is None or bi < self.nbranch) and candidates:
                _, _, new_flops, new_size, (i, j), k12 = heapq.heappop(candidates)
                _branch_iterate(
                    path=path + ((i, j),),
                    inputs=inputs + (k12,),
                    remaining=(remaining - {i, j}) | {len(inputs)},
                    flops=new_flops,
                    size=new_size,
                )
                bi += 1

        _branch_iterate(path=(), inputs=inputs, remaining=set(range(len(inputs))), flops=0, size=0)

        return self.path

__call__(inputs_, output_, size_dict, memory_limit=None)

Parameters:

Name	Type	Description	Default
inputs_	List[ArrayIndexType]	List of sets that represent the lhs side of the einsum subscript	required
output_	ArrayIndexType	Set that represents the rhs side of the overall einsum subscript	required
size_dict	Dict[str, int]	Dictionary of index sizes	required
memory_limit	Optional[int]	The maximum number of elements in a temporary array.	None

Returns:

Name	Type	Description
path	PathType	The contraction order within the memory limit constraint.

Examples: ```python isets = [set('abd'), set('ac'), set('bdc')] oset = set('') idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4} optimal(isets, oset, idx_sizes, 5000)

> [(0, 2), (0, 1)]

Source code in opt_einsum/paths.py

def __call__(
    self,
    inputs_: List[ArrayIndexType],
    output_: ArrayIndexType,
    size_dict: Dict[str, int],
    memory_limit: Optional[int] = None,
) -> PathType:
    """Parameters:
        inputs_: List of sets that represent the lhs side of the einsum subscript
        output_: Set that represents the rhs side of the overall einsum subscript
        size_dict: Dictionary of index sizes
        memory_limit: The maximum number of elements in a temporary array.

    Returns:
        path: The contraction order within the memory limit constraint.

    Examples:
    ```python
    isets = [set('abd'), set('ac'), set('bdc')]
    oset = set('')
    idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4}
    optimal(isets, oset, idx_sizes, 5000)
    #> [(0, 2), (0, 1)]
    """
    self._check_args_against_first_call(inputs_, output_, size_dict)

    inputs: Tuple[FrozenSet[str]] = tuple(map(frozenset, inputs_))  # type: ignore
    output: FrozenSet[str] = frozenset(output_)

    size_cache = {k: compute_size_by_dict(k, size_dict) for k in inputs}
    result_cache: Dict[Tuple[FrozenSet[str], FrozenSet[str]], Tuple[FrozenSet[str], int]] = {}

    def _branch_iterate(path, inputs, remaining, flops, size):
        # reached end of path (only ever get here if flops is best found so far)
        if len(remaining) == 1:
            self.best["size"] = size
            self.best["flops"] = flops
            self.best["ssa_path"] = path
            return

        def _assess_candidate(k1: FrozenSet[str], k2: FrozenSet[str], i: int, j: int) -> Any:
            # find resulting indices and flops
            try:
                k12, flops12 = result_cache[k1, k2]
            except KeyError:
                k12, flops12 = result_cache[k1, k2] = calc_k12_flops(inputs, output, remaining, i, j, size_dict)

            try:
                size12 = size_cache[k12]
            except KeyError:
                size12 = size_cache[k12] = compute_size_by_dict(k12, size_dict)

            new_flops = flops + flops12
            new_size = max(size, size12)

            # sieve based on current best i.e. check flops and size still better
            if not self.better(new_flops, new_size, self.best["flops"], self.best["size"]):
                return None

            # compare to how the best method was doing as this point
            if new_flops < self.best_progress[len(inputs)]:
                self.best_progress[len(inputs)] = new_flops
            # sieve based on current progress relative to best
            elif new_flops > self.cutoff_flops_factor * self.best_progress[len(inputs)]:
                return None

            # sieve based on memory limit
            if (memory_limit not in _UNLIMITED_MEM) and (size12 > memory_limit):  # type: ignore
                # terminate path here, but check all-terms contract first
                new_flops = flops + _compute_oversize_flops(inputs, remaining, output_, size_dict)
                if new_flops < self.best["flops"]:
                    self.best["flops"] = new_flops
                    self.best["ssa_path"] = path + (tuple(remaining),)
                return None

            # set cost heuristic in order to locally sort possible contractions
            size1, size2 = size_cache[inputs[i]], size_cache[inputs[j]]
            cost = self.cost_fn(size12, size1, size2, k12, k1, k2)

            return cost, flops12, new_flops, new_size, (i, j), k12

        # check all possible remaining paths
        candidates = []
        for i, j in itertools.combinations(remaining, 2):
            if i > j:
                i, j = j, i
            k1, k2 = inputs[i], inputs[j]

            # initially ignore outer products
            if k1.isdisjoint(k2):
                continue

            candidate = _assess_candidate(k1, k2, i, j)
            if candidate:
                heapq.heappush(candidates, candidate)

        # assess outer products if nothing left
        if not candidates:
            for i, j in itertools.combinations(remaining, 2):
                if i > j:
                    i, j = j, i
                k1, k2 = inputs[i], inputs[j]
                candidate = _assess_candidate(k1, k2, i, j)
                if candidate:
                    heapq.heappush(candidates, candidate)

        # recurse into all or some of the best candidate contractions
        bi = 0
        while (self.nbranch is None or bi < self.nbranch) and candidates:
            _, _, new_flops, new_size, (i, j), k12 = heapq.heappop(candidates)
            _branch_iterate(
                path=path + ((i, j),),
                inputs=inputs + (k12,),
                remaining=(remaining - {i, j}) | {len(inputs)},
                flops=new_flops,
                size=new_size,
            )
            bi += 1

    _branch_iterate(path=(), inputs=inputs, remaining=set(range(len(inputs))), flops=0, size=0)

    return self.path

__init__(nbranch=None, cutoff_flops_factor=4, minimize='flops', cost_fn='memory-removed')

Explores possible pair contractions in a depth-first recursive manner like the optimal approach, but with extra heuristic early pruning of branches as well sieving by memory_limit and the best path found so far.

Parameters:

Name	Type	Description	Default
nbranch	Optional[int]	How many branches to explore at each contraction step. If None, explore all possible branches. If an integer, branch into this many paths at each step. Defaults to None.	None
cutoff_flops_factor	int	If at any point, a path is doing this much worse than the best path found so far was, terminate it. The larger this is made, the more paths will be fully explored and the slower the algorithm. Defaults to 4.	4
minimize	str	Whether to optimize the path with regard primarily to the total estimated flop-count, or the size of the largest intermediate. The option not chosen will still be used as a secondary criterion.	'flops'
cost_fn	str	A function that returns a heuristic 'cost' of a potential contraction with which to sort candidates. Should have signature cost_fn(size12, size1, size2, k12, k1, k2).	'memory-removed'

Source code in opt_einsum/paths.py

def __init__(
    self,
    nbranch: Optional[int] = None,
    cutoff_flops_factor: int = 4,
    minimize: str = "flops",
    cost_fn: str = "memory-removed",
):
    """Explores possible pair contractions in a depth-first recursive manner like
    the `optimal` approach, but with extra heuristic early pruning of branches
    as well sieving by `memory_limit` and the best path found so far.


    Parameters:
        nbranch: How many branches to explore at each contraction step. If None, explore
            all possible branches. If an integer, branch into this many paths at
            each step. Defaults to None.
        cutoff_flops_factor: If at any point, a path is doing this much worse than the best path
            found so far was, terminate it. The larger this is made, the more paths
            will be fully explored and the slower the algorithm. Defaults to 4.
        minimize: Whether to optimize the path with regard primarily to the total
            estimated flop-count, or the size of the largest intermediate. The
            option not chosen will still be used as a secondary criterion.
        cost_fn: A function that returns a heuristic 'cost' of a potential contraction
            with which to sort candidates. Should have signature
            `cost_fn(size12, size1, size2, k12, k1, k2)`.
    """
    if (nbranch is not None) and nbranch < 1:
        raise ValueError(f"The number of branches must be at least one, `nbranch={nbranch}`.")

    self.nbranch = nbranch
    self.cutoff_flops_factor = cutoff_flops_factor
    self.minimize = minimize
    self.cost_fn: Any = _COST_FNS.get(cost_fn, cost_fn)

    self.better = get_better_fn(minimize)
    self.best: Dict[str, Any] = {"flops": float("inf"), "size": float("inf")}
    self.best_progress: Dict[int, float] = defaultdict(lambda: float("inf"))

opt_einsum.path_random.RandomOptimizer

Bases: PathOptimizer

Base class for running any random path finder that benefits from repeated calling, possibly in a parallel fashion. Custom random optimizers should subclass this, and the setup method should be implemented with the following signature:

def setup(self, inputs, output, size_dict):
    # custom preparation here ...
    return trial_fn, trial_args

Where trial_fn itself should have the signature::

def trial_fn(r, *trial_args):
    # custom computation of path here
    return ssa_path, cost, size

Where r is the run number and could for example be used to seed a random number generator. See RandomGreedy for an example.

Parameters:

Name	Type	Description	Default
max_repeats	int	The maximum number of repeat trials to have.	32
max_time	Optional[float]	The maximum amount of time to run the algorithm for.	None
minimize	str	Whether to favour paths that minimize the total estimated flop-count or the size of the largest intermediate created.	'flops'
parallel	Union[bool, Decimal, int]	Whether to parallelize the random trials, by default False. If True, use a concurrent.futures.ProcessPoolExecutor with the same number of processes as cores. If an integer is specified, use that many processes instead. Finally, you can supply a custom executor-pool which should have an API matching that of the python 3 standard library module concurrent.futures. Namely, a submit method that returns Future objects, themselves with result and cancel methods.	False
pre_dispatch	int	If running in parallel, how many jobs to pre-dispatch so as to avoid submitting all jobs at once. Should also be more than twice the number of workers to avoid under-subscription. Default: 128.	128

Attributes:

Name	Type	Description
path	PathType	The best path found so far.
costs	List[int]	The list of each trial's costs found so far.
sizes	List[int]	The list of each trial's largest intermediate size so far.

Source code in opt_einsum/path_random.py

class RandomOptimizer(paths.PathOptimizer):
    """Base class for running any random path finder that benefits
    from repeated calling, possibly in a parallel fashion. Custom random
    optimizers should subclass this, and the `setup` method should be
    implemented with the following signature:

    ```python
    def setup(self, inputs, output, size_dict):
        # custom preparation here ...
        return trial_fn, trial_args
    ```

    Where `trial_fn` itself should have the signature::

    ```python
    def trial_fn(r, *trial_args):
        # custom computation of path here
        return ssa_path, cost, size
    ```

    Where `r` is the run number and could for example be used to seed a
    random number generator. See `RandomGreedy` for an example.


    Parameters:
        max_repeats: The maximum number of repeat trials to have.
        max_time: The maximum amount of time to run the algorithm for.
        minimize:  Whether to favour paths that minimize the total estimated flop-count or
            the size of the largest intermediate created.
        parallel: Whether to parallelize the random trials, by default `False`. If
            `True`, use a `concurrent.futures.ProcessPoolExecutor` with the same
            number of processes as cores. If an integer is specified, use that many
            processes instead. Finally, you can supply a custom executor-pool which
            should have an API matching that of the python 3 standard library
            module `concurrent.futures`. Namely, a `submit` method that returns
            `Future` objects, themselves with `result` and `cancel` methods.
        pre_dispatch: If running in parallel, how many jobs to pre-dispatch so as to avoid
            submitting all jobs at once. Should also be more than twice the number
            of workers to avoid under-subscription. Default: 128.

    Attributes:
        path: The best path found so far.
        costs: The list of each trial's costs found so far.
        sizes: The list of each trial's largest intermediate size so far.
    """

    def __init__(
        self,
        max_repeats: int = 32,
        max_time: Optional[float] = None,
        minimize: str = "flops",
        parallel: Union[bool, Decimal, int] = False,
        pre_dispatch: int = 128,
    ):
        if minimize not in ("flops", "size"):
            raise ValueError("`minimize` should be one of {'flops', 'size'}.")

        self.max_repeats = max_repeats
        self.max_time = max_time
        self.minimize = minimize
        self.better = paths.get_better_fn(minimize)
        self._parallel: Union[bool, Decimal, int] = False
        self.parallel = parallel
        self.pre_dispatch = pre_dispatch

        self.costs: List[int] = []
        self.sizes: List[int] = []
        self.best: Dict[str, Any] = {"flops": float("inf"), "size": float("inf")}

        self._repeats_start = 0
        self._executor: Any
        self._futures: Any

    @property
    def path(self) -> PathType:
        """The best path found so far."""
        return paths.ssa_to_linear(self.best["ssa_path"])

    @property
    def parallel(self) -> Union[bool, Decimal, int]:
        return self._parallel

    @parallel.setter
    def parallel(self, parallel: Union[bool, Decimal, int]) -> None:
        # shutdown any previous executor if we are managing it
        if getattr(self, "_managing_executor", False):
            self._executor.shutdown()

        self._parallel = parallel
        self._managing_executor = False

        if parallel is False:
            self._executor = None
            return

        if parallel is True:
            from concurrent.futures import ProcessPoolExecutor

            self._executor = ProcessPoolExecutor()
            self._managing_executor = True
            return

        if isinstance(parallel, (int, Decimal)):
            from concurrent.futures import ProcessPoolExecutor

            self._executor = ProcessPoolExecutor(int(parallel))
            self._managing_executor = True
            return

        # assume a pool-executor has been supplied
        self._executor = parallel

    def _gen_results_parallel(self, repeats: Iterable[int], trial_fn: Any, args: Any) -> Generator[Any, None, None]:
        """Lazily generate results from an executor without submitting all jobs at once."""
        self._futures = deque()

        # the idea here is to submit at least ``pre_dispatch`` jobs *before* we
        # yield any results, then do both in tandem, before draining the queue
        for r in repeats:
            if len(self._futures) < self.pre_dispatch:
                self._futures.append(self._executor.submit(trial_fn, r, *args))
                continue
            yield self._futures.popleft().result()

        while self._futures:
            yield self._futures.popleft().result()

    def _cancel_futures(self) -> None:
        if self._executor is not None:
            for f in self._futures:
                f.cancel()

    def setup(
        self,
        inputs: List[ArrayIndexType],
        output: ArrayIndexType,
        size_dict: Dict[str, int],
    ) -> Tuple[Any, Any]:
        raise NotImplementedError

    def __call__(
        self,
        inputs: List[ArrayIndexType],
        output: ArrayIndexType,
        size_dict: Dict[str, int],
        memory_limit: Optional[int] = None,
    ) -> PathType:
        self._check_args_against_first_call(inputs, output, size_dict)

        # start a timer?
        if self.max_time is not None:
            t0 = time.time()

        trial_fn, trial_args = self.setup(inputs, output, size_dict)

        r_start = self._repeats_start + len(self.costs)
        r_stop = r_start + self.max_repeats
        repeats = range(r_start, r_stop)

        # create the trials lazily
        if self._executor is not None:
            trials = self._gen_results_parallel(repeats, trial_fn, trial_args)
        else:
            trials = (trial_fn(r, *trial_args) for r in repeats)

        # assess the trials
        for ssa_path, cost, size in trials:
            # keep track of all costs and sizes
            self.costs.append(cost)
            self.sizes.append(size)

            # check if we have found a new best
            found_new_best = self.better(cost, size, self.best["flops"], self.best["size"])

            if found_new_best:
                self.best["flops"] = cost
                self.best["size"] = size
                self.best["ssa_path"] = ssa_path

            # check if we have run out of time
            if (self.max_time is not None) and (time.time() > t0 + self.max_time):
                break

        self._cancel_futures()
        return self.path

    def __del__(self):
        # if we created the parallel pool-executor, shut it down
        if getattr(self, "_managing_executor", False):
            self._executor.shutdown()

path: PathType property

The best path found so far.

opt_einsum.path_random.RandomGreedy

Bases: RandomOptimizer

Source code in opt_einsum/path_random.py

class RandomGreedy(RandomOptimizer):
    def __init__(
        self,
        cost_fn: str = "memory-removed-jitter",
        temperature: float = 1.0,
        rel_temperature: bool = True,
        nbranch: int = 8,
        **kwargs: Any,
    ):
        """Parameters:
        cost_fn: A function that returns a heuristic 'cost' of a potential contraction
                with which to sort candidates. Should have signature
                `cost_fn(size12, size1, size2, k12, k1, k2)`.
        temperature: When choosing a possible contraction, its relative probability will be
                proportional to `exp(-cost / temperature)`. Thus the larger
                `temperature` is, the further random paths will stray from the normal
                'greedy' path. Conversely, if set to zero, only paths with exactly the
                same cost as the best at each step will be explored.
        rel_temperature: Whether to normalize the ``temperature`` at each step to the scale of
                the best cost. This is generally beneficial as the magnitude of costs
                can vary significantly throughout a contraction. If False, the
                algorithm will end up branching when the absolute cost is low, but
                stick to the 'greedy' path when the cost is high - this can also be
                beneficial.
        nbranch: How many potential paths to calculate probability for and choose from at each step.
        kwargs: Supplied to RandomOptimizer.
        """
        self.cost_fn = cost_fn
        self.temperature = temperature
        self.rel_temperature = rel_temperature
        self.nbranch = nbranch
        super().__init__(**kwargs)

    @property
    def choose_fn(self) -> Any:
        """The function that chooses which contraction to take - make this a
        property so that ``temperature`` and ``nbranch`` etc. can be updated
        between runs.
        """
        if self.nbranch == 1:
            return None

        return functools.partial(
            thermal_chooser,
            temperature=self.temperature,
            nbranch=self.nbranch,
            rel_temperature=self.rel_temperature,
        )

    def setup(
        self,
        inputs: List[ArrayIndexType],
        output: ArrayIndexType,
        size_dict: Dict[str, int],
    ) -> Tuple[Any, Any]:
        fn = _trial_greedy_ssa_path_and_cost
        args = (inputs, output, size_dict, self.choose_fn, self.cost_fn)
        return fn, args

choose_fn: Any property

The function that chooses which contraction to take - make this a property so that temperature and nbranch etc. can be updated between runs.

__init__(cost_fn='memory-removed-jitter', temperature=1.0, rel_temperature=True, nbranch=8, **kwargs)

cost_fn: A function that returns a heuristic 'cost' of a potential contraction with which to sort candidates. Should have signature cost_fn(size12, size1, size2, k12, k1, k2). temperature: When choosing a possible contraction, its relative probability will be proportional to exp(-cost / temperature). Thus the larger temperature is, the further random paths will stray from the normal 'greedy' path. Conversely, if set to zero, only paths with exactly the same cost as the best at each step will be explored. rel_temperature: Whether to normalize the temperature at each step to the scale of the best cost. This is generally beneficial as the magnitude of costs can vary significantly throughout a contraction. If False, the algorithm will end up branching when the absolute cost is low, but stick to the 'greedy' path when the cost is high - this can also be beneficial. nbranch: How many potential paths to calculate probability for and choose from at each step. kwargs: Supplied to RandomOptimizer.

Source code in opt_einsum/path_random.py

def __init__(
    self,
    cost_fn: str = "memory-removed-jitter",
    temperature: float = 1.0,
    rel_temperature: bool = True,
    nbranch: int = 8,
    **kwargs: Any,
):
    """Parameters:
    cost_fn: A function that returns a heuristic 'cost' of a potential contraction
            with which to sort candidates. Should have signature
            `cost_fn(size12, size1, size2, k12, k1, k2)`.
    temperature: When choosing a possible contraction, its relative probability will be
            proportional to `exp(-cost / temperature)`. Thus the larger
            `temperature` is, the further random paths will stray from the normal
            'greedy' path. Conversely, if set to zero, only paths with exactly the
            same cost as the best at each step will be explored.
    rel_temperature: Whether to normalize the ``temperature`` at each step to the scale of
            the best cost. This is generally beneficial as the magnitude of costs
            can vary significantly throughout a contraction. If False, the
            algorithm will end up branching when the absolute cost is low, but
            stick to the 'greedy' path when the cost is high - this can also be
            beneficial.
    nbranch: How many potential paths to calculate probability for and choose from at each step.
    kwargs: Supplied to RandomOptimizer.
    """
    self.cost_fn = cost_fn
    self.temperature = temperature
    self.rel_temperature = rel_temperature
    self.nbranch = nbranch
    super().__init__(**kwargs)

opt_einsum.paths.DynamicProgramming

Bases: PathOptimizer

Finds the optimal path of pairwise contractions without intermediate outer products based a dynamic programming approach presented in Phys. Rev. E 90, 033315 (2014) (the corresponding preprint is publicly available at https://arxiv.org/abs/1304.6112). This method is especially well-suited in the area of tensor network states, where it usually outperforms all the other optimization strategies.

This algorithm shows exponential scaling with the number of inputs in the worst case scenario (see example below). If the graph to be contracted consists of disconnected subgraphs, the algorithm scales linearly in the number of disconnected subgraphs and only exponentially with the number of inputs per subgraph.

Parameters:

Name	Type	Description	Default
minimize	str	What to minimize: - 'flops' - minimize the number of flops - 'size' - minimize the size of the largest intermediate - 'write' - minimize the size of all intermediate tensors - 'combo' - minimize flops + alpha * write summed over intermediates, a default ratio of alpha=64 is used, or it can be customized with f'combo-{alpha}' - 'limit' - minimize max(flops, alpha * write) summed over intermediates, a default ratio of alpha=64 is used, or it can be customized with f'limit-{alpha}' - callable - a custom local cost function	'flops'
cost_cap	Union[bool, int]	How to implement cost-capping: - True - iteratively increase the cost-cap - False - implement no cost-cap at all - int - use explicit cost cap	True
search_outer	bool	In rare circumstances the optimal contraction may involve an outer product, this option allows searching such contractions but may well slow down the path finding considerably on all but very small graphs.	False

Source code in opt_einsum/paths.py

class DynamicProgramming(PathOptimizer):
    """Finds the optimal path of pairwise contractions without intermediate outer
    products based a dynamic programming approach presented in
    Phys. Rev. E 90, 033315 (2014) (the corresponding preprint is publicly
    available at https://arxiv.org/abs/1304.6112). This method is especially
    well-suited in the area of tensor network states, where it usually
    outperforms all the other optimization strategies.

    This algorithm shows exponential scaling with the number of inputs
    in the worst case scenario (see example below). If the graph to be
    contracted consists of disconnected subgraphs, the algorithm scales
    linearly in the number of disconnected subgraphs and only exponentially
    with the number of inputs per subgraph.

    Parameters:
        minimize: What to minimize:
            - 'flops' - minimize the number of flops
            - 'size' - minimize the size of the largest intermediate
            - 'write' - minimize the size of all intermediate tensors
            - 'combo' - minimize `flops + alpha * write` summed over intermediates, a default ratio of alpha=64
            is used, or it can be customized with `f'combo-{alpha}'`
            - 'limit' - minimize `max(flops, alpha * write)` summed over intermediates, a default ratio of alpha=64
            is used, or it can be customized with `f'limit-{alpha}'`
            - callable - a custom local cost function

        cost_cap: How to implement cost-capping:
            - True - iteratively increase the cost-cap
            - False - implement no cost-cap at all
            - int - use explicit cost cap

        search_outer: In rare circumstances the optimal contraction may involve an outer
            product, this option allows searching such contractions but may well
            slow down the path finding considerably on all but very small graphs.
    """

    def __init__(self, minimize: str = "flops", cost_cap: Union[bool, int] = True, search_outer: bool = False) -> None:
        self.minimize = minimize
        self.search_outer = search_outer
        self.cost_cap = cost_cap

    def __call__(
        self,
        inputs_: List[ArrayIndexType],
        output_: ArrayIndexType,
        size_dict_: Dict[str, int],
        memory_limit_: Optional[int] = None,
    ) -> PathType:
        """Parameters:
            inputs_: List of sets that represent the lhs side of the einsum subscript
            output_: Set that represents the rhs side of the overall einsum subscript
            size_dict_: Dictionary of index sizes
            memory_limit_: The maximum number of elements in a temporary array.

        Returns:
            path: The contraction order (a list of tuples of ints).

        Examples:
            ```python
            n_in = 3  # exponential scaling
            n_out = 2 # linear scaling
            s = dict()
            i_all = []
            for _ in range(n_out):
                i = [set() for _ in range(n_in)]
                for j in range(n_in):
                    for k in range(j+1, n_in):
                        c = oe.get_symbol(len(s))
                        i[j].add(c)
                        i[k].add(c)
                        s[c] = 2
                i_all.extend(i)
            o = DynamicProgramming()
            o(i_all, set(), s)
            #> [(1, 2), (0, 4), (1, 2), (0, 2), (0, 1)]
            ```
        """
        _check_contraction, naive_scale = _parse_minimize(self.minimize)
        _check_outer = (lambda x: True) if self.search_outer else (lambda x: x)

        ind_counts = Counter(itertools.chain(*inputs_, output_))
        all_inds = tuple(ind_counts)

        # convert all indices to integers (makes set operations ~10 % faster)
        symbol2int = {c: j for j, c in enumerate(all_inds)}
        inputs = [frozenset(symbol2int[c] for c in i) for i in inputs_]
        output = frozenset(symbol2int[c] for c in output_)
        size_dict_canonical = {symbol2int[c]: v for c, v in size_dict_.items() if c in symbol2int}
        size_dict = [size_dict_canonical[j] for j in range(len(size_dict_canonical))]
        naive_cost = naive_scale * len(inputs) * functools.reduce(operator.mul, size_dict, 1)

        inputs, inputs_done, inputs_contractions = _dp_parse_out_single_term_ops(inputs, all_inds, ind_counts)

        if not inputs:
            # nothing left to do after single axis reductions!
            return _tree_to_sequence(simple_tree_tuple(inputs_done))

        # a list of all necessary contraction expressions for each of the
        # disconnected subgraphs and their size
        subgraph_contractions = inputs_done
        subgraph_contractions_size = [1] * len(inputs_done)

        if self.search_outer:
            # optimize everything together if we are considering outer products
            subgraphs = [frozenset(range(len(inputs)))]
        else:
            subgraphs = _find_disconnected_subgraphs(inputs, output)

        # the bitmap set of all tensors is computed as it is needed to
        # compute set differences: s1 - s2 transforms into
        # s1 & (all_tensors ^ s2)
        all_tensors = (1 << len(inputs)) - 1

        for g in subgraphs:
            # dynamic programming approach to compute x[n] for subgraph g;
            # x[n][set of n tensors] = (indices, cost, contraction)
            # the set of n tensors is represented by a bitmap: if bit j is 1,
            # tensor j is in the set, e.g. 0b100101 = {0,2,5}; set unions
            # (intersections) can then be computed by bitwise or (and);
            x: List[Any] = [None] * 2 + [{} for j in range(len(g) - 1)]
            x[1] = {1 << j: (inputs[j], 0, inputs_contractions[j]) for j in g}

            # convert set of tensors g to a bitmap set:
            bitmap_g = functools.reduce(lambda x, y: x | y, (1 << j for j in g))

            # try to find contraction with cost <= cost_cap and increase
            # cost_cap successively if no such contraction is found;
            # this is a major performance improvement; start with product of
            # output index dimensions as initial cost_cap
            subgraph_inds = frozenset.union(*_bitmap_select(bitmap_g, inputs))
            if self.cost_cap is True:
                cost_cap = compute_size_by_dict(subgraph_inds & output, size_dict)
            elif self.cost_cap is False:
                cost_cap = float("inf")  # type: ignore
            else:
                cost_cap = self.cost_cap
            # set the factor to increase the cost by each iteration (ensure > 1)
            if len(subgraph_inds) == 0:
                cost_increment = 2
            else:
                cost_increment = max(min(map(size_dict.__getitem__, subgraph_inds)), 2)

            while len(x[-1]) == 0:
                for n in range(2, len(x[1]) + 1):
                    xn = x[n]

                    # try to combine solutions from x[m] and x[n-m]
                    for m in range(1, n // 2 + 1):
                        for s1, (i1, cost1, contract1) in x[m].items():
                            for s2, (i2, cost2, contract2) in x[n - m].items():
                                # can only merge if s1 and s2 are disjoint
                                # and avoid e.g. s1={0}, s2={1} and s1={1}, s2={0}
                                if (not s1 & s2) and (m != n - m or s1 < s2):
                                    i1_cut_i2_wo_output = (i1 & i2) - output

                                    # maybe ignore outer products:
                                    if _check_outer(i1_cut_i2_wo_output):
                                        i1_union_i2 = i1 | i2
                                        _check_contraction(
                                            cost1,
                                            cost2,
                                            i1_union_i2,
                                            size_dict,
                                            cost_cap,
                                            s1,
                                            s2,
                                            xn,
                                            bitmap_g,
                                            all_tensors,
                                            inputs,
                                            i1_cut_i2_wo_output,
                                            memory_limit_,
                                            contract1,
                                            contract2,
                                        )

                if (cost_cap > naive_cost) and (len(x[-1]) == 0):
                    raise RuntimeError("No contraction found for given `memory_limit`.")

                # increase cost cap for next iteration:
                cost_cap = cost_increment * cost_cap

            i, cost, contraction = list(x[-1].values())[0]
            subgraph_contractions.append(contraction)
            subgraph_contractions_size.append(compute_size_by_dict(i, size_dict))

        # sort the subgraph contractions by the size of the subgraphs in
        # ascending order (will give the cheapest contractions); note that
        # outer products should be performed pairwise (to use BLAS functions)
        subgraph_contractions = [
            subgraph_contractions[j]
            for j in sorted(
                range(len(subgraph_contractions_size)),
                key=subgraph_contractions_size.__getitem__,
            )
        ]

        # build the final contraction tree
        tree = simple_tree_tuple(subgraph_contractions)
        return _tree_to_sequence(tree)

__call__(inputs_, output_, size_dict_, memory_limit_=None)

Parameters:

Name	Type	Description	Default
inputs_	List[ArrayIndexType]	List of sets that represent the lhs side of the einsum subscript	required
output_	ArrayIndexType	Set that represents the rhs side of the overall einsum subscript	required
size_dict_	Dict[str, int]	Dictionary of index sizes	required
memory_limit_	Optional[int]	The maximum number of elements in a temporary array.	None

Returns:

Name	Type	Description
path	PathType	The contraction order (a list of tuples of ints).

Examples:

n_in = 3  # exponential scaling
n_out = 2 # linear scaling
s = dict()
i_all = []
for _ in range(n_out):
    i = [set() for _ in range(n_in)]
    for j in range(n_in):
        for k in range(j+1, n_in):
            c = oe.get_symbol(len(s))
            i[j].add(c)
            i[k].add(c)
            s[c] = 2
    i_all.extend(i)
o = DynamicProgramming()
o(i_all, set(), s)
#> [(1, 2), (0, 4), (1, 2), (0, 2), (0, 1)]

Source code in opt_einsum/paths.py

def __call__(
    self,
    inputs_: List[ArrayIndexType],
    output_: ArrayIndexType,
    size_dict_: Dict[str, int],
    memory_limit_: Optional[int] = None,
) -> PathType:
    """Parameters:
        inputs_: List of sets that represent the lhs side of the einsum subscript
        output_: Set that represents the rhs side of the overall einsum subscript
        size_dict_: Dictionary of index sizes
        memory_limit_: The maximum number of elements in a temporary array.

    Returns:
        path: The contraction order (a list of tuples of ints).

    Examples:
        ```python
        n_in = 3  # exponential scaling
        n_out = 2 # linear scaling
        s = dict()
        i_all = []
        for _ in range(n_out):
            i = [set() for _ in range(n_in)]
            for j in range(n_in):
                for k in range(j+1, n_in):
                    c = oe.get_symbol(len(s))
                    i[j].add(c)
                    i[k].add(c)
                    s[c] = 2
            i_all.extend(i)
        o = DynamicProgramming()
        o(i_all, set(), s)
        #> [(1, 2), (0, 4), (1, 2), (0, 2), (0, 1)]
        ```
    """
    _check_contraction, naive_scale = _parse_minimize(self.minimize)
    _check_outer = (lambda x: True) if self.search_outer else (lambda x: x)

    ind_counts = Counter(itertools.chain(*inputs_, output_))
    all_inds = tuple(ind_counts)

    # convert all indices to integers (makes set operations ~10 % faster)
    symbol2int = {c: j for j, c in enumerate(all_inds)}
    inputs = [frozenset(symbol2int[c] for c in i) for i in inputs_]
    output = frozenset(symbol2int[c] for c in output_)
    size_dict_canonical = {symbol2int[c]: v for c, v in size_dict_.items() if c in symbol2int}
    size_dict = [size_dict_canonical[j] for j in range(len(size_dict_canonical))]
    naive_cost = naive_scale * len(inputs) * functools.reduce(operator.mul, size_dict, 1)

    inputs, inputs_done, inputs_contractions = _dp_parse_out_single_term_ops(inputs, all_inds, ind_counts)

    if not inputs:
        # nothing left to do after single axis reductions!
        return _tree_to_sequence(simple_tree_tuple(inputs_done))

    # a list of all necessary contraction expressions for each of the
    # disconnected subgraphs and their size
    subgraph_contractions = inputs_done
    subgraph_contractions_size = [1] * len(inputs_done)

    if self.search_outer:
        # optimize everything together if we are considering outer products
        subgraphs = [frozenset(range(len(inputs)))]
    else:
        subgraphs = _find_disconnected_subgraphs(inputs, output)

    # the bitmap set of all tensors is computed as it is needed to
    # compute set differences: s1 - s2 transforms into
    # s1 & (all_tensors ^ s2)
    all_tensors = (1 << len(inputs)) - 1

    for g in subgraphs:
        # dynamic programming approach to compute x[n] for subgraph g;
        # x[n][set of n tensors] = (indices, cost, contraction)
        # the set of n tensors is represented by a bitmap: if bit j is 1,
        # tensor j is in the set, e.g. 0b100101 = {0,2,5}; set unions
        # (intersections) can then be computed by bitwise or (and);
        x: List[Any] = [None] * 2 + [{} for j in range(len(g) - 1)]
        x[1] = {1 << j: (inputs[j], 0, inputs_contractions[j]) for j in g}

        # convert set of tensors g to a bitmap set:
        bitmap_g = functools.reduce(lambda x, y: x | y, (1 << j for j in g))

        # try to find contraction with cost <= cost_cap and increase
        # cost_cap successively if no such contraction is found;
        # this is a major performance improvement; start with product of
        # output index dimensions as initial cost_cap
        subgraph_inds = frozenset.union(*_bitmap_select(bitmap_g, inputs))
        if self.cost_cap is True:
            cost_cap = compute_size_by_dict(subgraph_inds & output, size_dict)
        elif self.cost_cap is False:
            cost_cap = float("inf")  # type: ignore
        else:
            cost_cap = self.cost_cap
        # set the factor to increase the cost by each iteration (ensure > 1)
        if len(subgraph_inds) == 0:
            cost_increment = 2
        else:
            cost_increment = max(min(map(size_dict.__getitem__, subgraph_inds)), 2)

        while len(x[-1]) == 0:
            for n in range(2, len(x[1]) + 1):
                xn = x[n]

                # try to combine solutions from x[m] and x[n-m]
                for m in range(1, n // 2 + 1):
                    for s1, (i1, cost1, contract1) in x[m].items():
                        for s2, (i2, cost2, contract2) in x[n - m].items():
                            # can only merge if s1 and s2 are disjoint
                            # and avoid e.g. s1={0}, s2={1} and s1={1}, s2={0}
                            if (not s1 & s2) and (m != n - m or s1 < s2):
                                i1_cut_i2_wo_output = (i1 & i2) - output

                                # maybe ignore outer products:
                                if _check_outer(i1_cut_i2_wo_output):
                                    i1_union_i2 = i1 | i2
                                    _check_contraction(
                                        cost1,
                                        cost2,
                                        i1_union_i2,
                                        size_dict,
                                        cost_cap,
                                        s1,
                                        s2,
                                        xn,
                                        bitmap_g,
                                        all_tensors,
                                        inputs,
                                        i1_cut_i2_wo_output,
                                        memory_limit_,
                                        contract1,
                                        contract2,
                                    )

            if (cost_cap > naive_cost) and (len(x[-1]) == 0):
                raise RuntimeError("No contraction found for given `memory_limit`.")

            # increase cost cap for next iteration:
            cost_cap = cost_increment * cost_cap

        i, cost, contraction = list(x[-1].values())[0]
        subgraph_contractions.append(contraction)
        subgraph_contractions_size.append(compute_size_by_dict(i, size_dict))

    # sort the subgraph contractions by the size of the subgraphs in
    # ascending order (will give the cheapest contractions); note that
    # outer products should be performed pairwise (to use BLAS functions)
    subgraph_contractions = [
        subgraph_contractions[j]
        for j in sorted(
            range(len(subgraph_contractions_size)),
            key=subgraph_contractions_size.__getitem__,
        )
    ]

    # build the final contraction tree
    tree = simple_tree_tuple(subgraph_contractions)
    return _tree_to_sequence(tree)