flopscope 0.2.0__py3-none-any.whl

flopscope/_opt_einsum/_paths.py

@@ -0,0 +1,1653 @@
+"""Contains the path technology behind opt_einsum in addition to several path helpers."""
+
+import bisect
+import functools
+import heapq
+import itertools
+import operator
+import random
+import re
+from collections import Counter, defaultdict
+from collections import Counter as CounterType
+from collections.abc import Callable, Generator, Sequence
+from typing import Any
+
+from ._helpers import compute_size_by_dict, flop_count
+from ._subgraph_symmetry import SubgraphSymmetryOracle
+from ._symmetry import (
+    SymmetryGroup,
+    symmetric_flop_count,
+    unique_elements,
+)
+from ._typing import ArrayIndexType, PathSearchFunctionType, PathType, TensorShapeType
+
+__all__ = [
+    "optimal",
+    "BranchBound",
+    "branch",
+    "greedy",
+    "auto",
+    "auto_hq",
+    "get_path_fn",
+    "DynamicProgramming",
+    "dynamic_programming",
+]
+
+_UNLIMITED_MEM = {-1, None, float("inf")}
+
+
+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: int | None = None,
+        **kwargs: Any,
+    ) -> PathType:
+        raise NotImplementedError
+
+
+def ssa_to_linear(ssa_path: PathType) -> PathType:
+    """Convert a path with static single assignment ids to a path with recycled
+    linear ids.
+
+    Example:
+        ```python
+        ssa_to_linear([(0, 3), (2, 4), (1, 5)])
+        #> [(0, 3), (1, 2), (0, 1)]
+        ```
+    """
+    n = sum(map(len, ssa_path)) - len(ssa_path) + 1
+    ids = list(range(n))
+    path = []
+    ssa = n
+    for scon in ssa_path:
+        con = sorted([bisect.bisect_left(ids, s) for s in scon])
+        for j in reversed(con):
+            ids.pop(j)
+        ids.append(ssa)
+        path.append(con)
+        ssa += 1
+    return [tuple(x) for x in path]
+
+
+def linear_to_ssa(path: PathType) -> PathType:
+    """Convert a path with recycled linear ids to a path with static single
+    assignment ids.
+
+    Exmaple:
+        ```python
+        linear_to_ssa([(0, 3), (1, 2), (0, 1)])
+        #> [(0, 3), (2, 4), (1, 5)]
+        ```
+    """
+    num_inputs = sum(map(len, path)) - len(path) + 1
+    linear_to_ssa = list(range(num_inputs))
+    new_ids = itertools.count(num_inputs)
+    ssa_path = []
+    for ids in path:
+        ssa_path.append(tuple(linear_to_ssa[id_] for id_ in ids))
+        for id_ in sorted(ids, reverse=True):
+            del linear_to_ssa[id_]
+        linear_to_ssa.append(next(new_ids))
+    return ssa_path
+
+
+def calc_k12_flops(
+    inputs: tuple[frozenset[str]],
+    output: frozenset[str],
+    remaining: frozenset[int],
+    i: int,
+    j: int,
+    size_dict: dict[str, int],
+    oracle: SubgraphSymmetryOracle | None = None,
+    ssa_to_subset: dict[int, frozenset[int]] | None = None,
+) -> tuple[frozenset[str], int, SymmetryGroup | None]:
+    """Calculate the resulting indices and flops for a potential pairwise
+    contraction.
+
+    Parameters
+    ----------
+    oracle : SubgraphSymmetryOracle | None
+        Subset-keyed symmetry oracle. When provided together with
+        ``ssa_to_subset``, the symmetry of the output tensor is looked
+        up via ``oracle.sym(ssa_to_subset[i] | ssa_to_subset[j])`` and
+        used to reduce the FLOP count.
+    ssa_to_subset : dict[int, frozenset[int]] | None
+        Mapping from SSA tensor id to the subset of original operand
+        positions that tensor represents.
+
+    Returns
+    -------
+    k12 : frozenset[str]
+        The resulting indices of the potential tensor.
+    cost : int
+        Estimated flop count of the operation.
+    sym12 : SymmetryGroup | None
+        Symmetry of the result tensor (None when no oracle or no symmetry).
+    """
+    k1, k2 = inputs[i], inputs[j]
+    either = k1 | k2
+    keep = frozenset.union(output, *map(inputs.__getitem__, remaining - {i, j}))
+    k12 = either & keep
+    inner = bool(either - k12)
+
+    sym12: SymmetryGroup | None = None
+    if oracle is not None and ssa_to_subset is not None:
+        merged_subset = ssa_to_subset[i] | ssa_to_subset[j]
+        subset_sym = oracle.sym(merged_subset)
+        sym12 = subset_sym.output
+
+        from flopscope._config import get_setting
+
+        idx_removed = either - k12
+        cost = symmetric_flop_count(
+            either,
+            inner,
+            2,
+            size_dict,
+            output_group=subset_sym.output,
+            output_indices=k12,
+            inner_group=subset_sym.inner,
+            inner_indices=idx_removed if idx_removed else None,
+            use_inner_symmetry=bool(get_setting("use_inner_symmetry")),
+        )
+    else:
+        cost = flop_count(either, inner, 2, size_dict)
+
+    return k12, cost, sym12
+
+
+def _compute_oversize_flops(
+    inputs: tuple[frozenset[str]],
+    remaining: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+) -> int:
+    """Compute the flop count for a contraction of all remaining arguments. This
+    is used when a memory limit means that no pairwise contractions can be made.
+    """
+    idx_contraction = frozenset.union(*map(inputs.__getitem__, remaining))  # type: ignore
+    inner = idx_contraction - output
+    num_terms = len(remaining)
+    return flop_count(idx_contraction, bool(inner), num_terms, size_dict)
+
+
+def optimal(
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+    memory_limit: int | None = None,
+    symmetry_oracle: SubgraphSymmetryOracle | None = None,
+) -> PathType:
+    """Computes all possible pair contractions in a depth-first recursive manner."""
+    inputs_set = tuple(map(frozenset, inputs))
+    output_set = frozenset(output)
+    num_operands = len(inputs)
+
+    best_flops = {"flops": float("inf")}
+    best_ssa_path = {"ssa_path": (tuple(range(len(inputs))),)}
+    size_cache: dict[frozenset[str], int] = {}
+
+    # Initial SSA -> subset mapping: each original operand covers itself.
+    initial_ssa_to_subset: dict[int, frozenset[int]] = {
+        k: frozenset({k}) for k in range(num_operands)
+    }
+
+    # Result cache is valid because calc_k12_flops is now pure in
+    # (merged_subset,) — the oracle guarantees that two subsets with the
+    # same frozenset key produce the same symmetry.
+    result_cache: dict[
+        tuple[ArrayIndexType, ArrayIndexType, frozenset[int]],
+        tuple[frozenset[str], int, SymmetryGroup | None],
+    ] = {}
+
+    def _optimal_iterate(path, remaining, inputs, flops, ssa_to_subset):
+        if len(remaining) == 1:
+            best_flops["flops"] = flops
+            best_ssa_path["ssa_path"] = path
+            return
+
+        for i, j in itertools.combinations(remaining, 2):
+            if i > j:
+                i, j = j, i
+
+            merged_subset = ssa_to_subset[i] | ssa_to_subset[j]
+            cache_key = (inputs[i], inputs[j], merged_subset)
+            try:
+                k12, flops12, sym12 = result_cache[cache_key]
+            except KeyError:
+                k12, flops12, sym12 = result_cache[cache_key] = calc_k12_flops(
+                    inputs,
+                    output_set,
+                    remaining,
+                    i,
+                    j,
+                    size_dict,
+                    oracle=symmetry_oracle,
+                    ssa_to_subset=ssa_to_subset,
+                )
+
+            new_flops = flops + flops12
+            if new_flops >= best_flops["flops"]:
+                continue
+
+            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)
+                if size12 > memory_limit:  # type: ignore[operator]
+                    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
+
+            new_ssa_to_subset = dict(ssa_to_subset)
+            new_ssa_to_subset[len(inputs)] = merged_subset
+
+            _optimal_iterate(
+                path=path + ((i, j),),
+                inputs=inputs + (k12,),
+                remaining=remaining - {i, j} | {len(inputs)},
+                flops=new_flops,
+                ssa_to_subset=new_ssa_to_subset,
+            )
+
+    _optimal_iterate(
+        path=(),
+        inputs=inputs_set,
+        remaining=set(range(len(inputs))),
+        flops=0,
+        ssa_to_subset=initial_ssa_to_subset,
+    )
+
+    return ssa_to_linear(best_ssa_path["ssa_path"])
+
+
+# functions for comparing which of two paths is 'better'
+
+
+def better_flops_first(flops: int, size: int, best_flops: int, best_size: int) -> bool:
+    return (flops, size) < (best_flops, best_size)
+
+
+def better_size_first(flops: int, size: int, best_flops: int, best_size: int) -> bool:
+    return (size, flops) < (best_size, best_flops)
+
+
+_BETTER_FNS = {
+    "flops": better_flops_first,
+    "size": better_size_first,
+}
+
+
+def get_better_fn(key: str) -> Callable[[int, int, int, int], bool]:
+    return _BETTER_FNS[key]
+
+
+# functions for assigning a heuristic 'cost' to a potential contraction
+
+
+def cost_memory_removed(
+    size12: int, size1: int, size2: int, k12: int, k1: int, k2: int
+) -> float:
+    """The default heuristic cost, corresponding to the total reduction in
+    memory of performing a contraction.
+    """
+    return size12 - size1 - size2
+
+
+def cost_memory_removed_jitter(
+    size12: int, size1: int, size2: int, k12: int, k1: int, k2: int
+) -> float:
+    """Like memory-removed, but with a slight amount of noise that breaks ties
+    and thus jumbles the contractions a bit.
+    """
+    return random.gauss(1.0, 0.01) * (size12 - size1 - size2)
+
+
+_COST_FNS = {
+    "memory-removed": cost_memory_removed,
+    "memory-removed-jitter": cost_memory_removed_jitter,
+}
+
+
+class BranchBound(PathOptimizer):
+    def __init__(
+        self,
+        nbranch: int | None = 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__(  # type: ignore[override]
+        self,
+        inputs_: list[ArrayIndexType],
+        output_: ArrayIndexType,
+        size_dict: dict[str, int],
+        memory_limit: int | None = None,
+        symmetry_oracle: SubgraphSymmetryOracle | None = 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.
+            symmetry_oracle: Optional subgraph symmetry oracle.
+
+        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[assignment]
+        output: frozenset[str] = frozenset(output_)
+        num_operands = len(inputs_)
+
+        size_cache = {k: compute_size_by_dict(k, size_dict) for k in inputs}
+
+        initial_ssa_to_subset: dict[int, frozenset[int]] = {
+            k: frozenset({k}) for k in range(num_operands)
+        }
+
+        # Result cache is valid with the oracle — key by (k1, k2, merged_subset)
+        result_cache: dict[
+            tuple[frozenset[str], frozenset[str], frozenset[int]],
+            tuple[frozenset[str], int, SymmetryGroup | None],
+        ] = {}
+
+        def _branch_iterate(path, inputs, remaining, flops, size, ssa_to_subset):
+            # 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
+                merged_subset = ssa_to_subset[i] | ssa_to_subset[j]
+                cache_key = (k1, k2, merged_subset)
+                try:
+                    k12, flops12, sym12 = result_cache[cache_key]
+                except KeyError:
+                    k12, flops12, sym12 = result_cache[cache_key] = calc_k12_flops(
+                        inputs,
+                        output,
+                        remaining,
+                        i,
+                        j,
+                        size_dict,
+                        oracle=symmetry_oracle,
+                        ssa_to_subset=ssa_to_subset,
+                    )
+
+                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, sym12
+
+            # 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, sym12 = heapq.heappop(
+                    candidates
+                )
+
+                new_ssa_to_subset = dict(ssa_to_subset)
+                new_ssa_to_subset[len(inputs)] = ssa_to_subset[i] | ssa_to_subset[j]
+
+                _branch_iterate(
+                    path=path + ((i, j),),
+                    inputs=inputs + (k12,),
+                    remaining=(remaining - {i, j}) | {len(inputs)},
+                    flops=new_flops,
+                    size=new_size,
+                    ssa_to_subset=new_ssa_to_subset,
+                )
+                bi += 1
+
+        _branch_iterate(
+            path=(),
+            inputs=inputs,
+            remaining=set(range(len(inputs))),
+            flops=0,
+            size=0,
+            ssa_to_subset=initial_ssa_to_subset,
+        )
+
+        return self.path
+
+
+def branch(
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+    memory_limit: int | None = None,
+    nbranch: int | None = None,
+    cutoff_flops_factor: int = 4,
+    minimize: str = "flops",
+    cost_fn: str = "memory-removed",
+    symmetry_oracle: SubgraphSymmetryOracle | None = None,
+) -> 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,
+        symmetry_oracle=symmetry_oracle,
+    )
+
+
+branch_all = functools.partial(branch, nbranch=None)
+branch_2 = functools.partial(branch, nbranch=2)
+branch_1 = functools.partial(branch, nbranch=1)
+
+GreedyCostType = tuple[int, int, int]
+GreedyContractionType = tuple[
+    GreedyCostType, ArrayIndexType, ArrayIndexType, ArrayIndexType
+]  # Cost, t1,t2->t3
+
+
+def _get_candidate(
+    output: ArrayIndexType,
+    sizes: dict[str, int],
+    remaining: dict[ArrayIndexType, int],
+    footprints: dict[ArrayIndexType, int],
+    dim_ref_counts: dict[int, set[str]],
+    k1: ArrayIndexType,
+    k2: ArrayIndexType,
+    cost_fn: Any,
+) -> GreedyContractionType:
+    either = k1 | k2
+    two = k1 & k2
+    one = either - two
+    k12 = (either & output) | (two & dim_ref_counts[3]) | (one & dim_ref_counts[2])
+    size12 = compute_size_by_dict(k12, sizes)
+    cost = cost_fn(
+        size12,
+        footprints[k1],
+        footprints[k2],
+        k12,
+        k1,
+        k2,
+    )
+    id1 = remaining[k1]
+    id2 = remaining[k2]
+    if id1 > id2:
+        k1, id1, k2, id2 = k2, id2, k1, id1
+    cost = cost, id2, id1  # break ties to ensure determinism
+    return cost, k1, k2, k12
+
+
+def _push_candidate(
+    output: ArrayIndexType,
+    sizes: dict[str, Any],
+    remaining: dict[ArrayIndexType, int],
+    footprints: dict[ArrayIndexType, int],
+    dim_ref_counts: dict[int, set[str]],
+    k1: ArrayIndexType,
+    k2s: list[ArrayIndexType],
+    queue: list[GreedyContractionType],
+    push_all: bool,
+    cost_fn: Any,
+) -> None:
+    candidates = (
+        _get_candidate(
+            output,
+            sizes,
+            remaining,
+            footprints,
+            dim_ref_counts,
+            k1,
+            k2,
+            cost_fn,
+        )
+        for k2 in k2s
+    )
+    if push_all:
+        # want to do this if we e.g. are using a custom 'choose_fn'
+        for candidate in candidates:
+            heapq.heappush(queue, candidate)
+    else:
+        heapq.heappush(queue, min(candidates))
+
+
+def _update_ref_counts(
+    dim_to_keys: dict[str, set[ArrayIndexType]],
+    dim_ref_counts: dict[int, set[str]],
+    dims: ArrayIndexType,
+) -> None:
+    for dim in dims:
+        count = len(dim_to_keys[dim])
+        if count <= 1:
+            dim_ref_counts[2].discard(dim)
+            dim_ref_counts[3].discard(dim)
+        elif count == 2:
+            dim_ref_counts[2].add(dim)
+            dim_ref_counts[3].discard(dim)
+        else:
+            dim_ref_counts[2].add(dim)
+            dim_ref_counts[3].add(dim)
+
+
+def _simple_chooser(queue, remaining):
+    """Default contraction chooser that simply takes the minimum cost option."""
+    cost, k1, k2, k12 = heapq.heappop(queue)
+    if k1 not in remaining or k2 not in remaining:
+        return None  # candidate is obsolete
+    return cost, k1, k2, k12
+
+
+def ssa_greedy_optimize(
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    sizes: dict[str, int],
+    choose_fn: Any = None,
+    cost_fn: Any = "memory-removed",
+    symmetry_oracle: SubgraphSymmetryOracle | None = None,
+    ssa_to_subset: dict[int, frozenset[int]] | None = None,
+) -> PathType:
+    """This is the core function for :func:`greedy` but produces a path with
+    static single assignment ids rather than recycled linear ids.
+    SSA ids are cheaper to work with and easier to reason about.
+    """
+    if len(inputs) == 1:
+        # Perform a single contraction to match output shape.
+        return [(0,)]
+
+    # set the function that assigns a heuristic cost to a possible contraction
+    cost_fn = _COST_FNS.get(cost_fn, cost_fn)
+
+    # set the function that chooses which contraction to take
+    if choose_fn is None:
+        choose_fn = _simple_chooser
+        push_all = False
+    else:
+        # assume chooser wants access to all possible contractions
+        push_all = True
+
+    num_operands = len(inputs)
+    if ssa_to_subset is None:
+        ssa_to_subset = {k: frozenset({k}) for k in range(num_operands)}
+
+    # A dim that is common to all tensors might as well be an output dim, since it
+    # cannot be contracted until the final step. This avoids an expensive all-pairs
+    # comparison to search for possible contractions at each step, leading to speedup
+    # in many practical problems where all tensors share a common batch dimension.
+    fs_inputs = [frozenset(x) for x in inputs]
+    output = frozenset(output) | frozenset.intersection(*fs_inputs)
+
+    # Deduplicate shapes by eagerly computing Hadamard products.
+    remaining: dict[ArrayIndexType, int] = {}  # key -> ssa_id
+    ssa_ids = itertools.count(len(fs_inputs))
+    ssa_path: list[TensorShapeType] = []
+    for ssa_id, key in enumerate(fs_inputs):
+        if key in remaining:
+            ssa_path.append((remaining[key], ssa_id))
+            new_id = next(ssa_ids)
+            # Hadamard dedup: merge subsets
+            old_id = remaining[key]
+            ssa_to_subset[new_id] = ssa_to_subset[old_id] | ssa_to_subset[ssa_id]
+            remaining[key] = new_id
+        else:
+            remaining[key] = ssa_id
+
+    # Keep track of possible contraction dims.
+    dim_to_keys = defaultdict(set)
+    for key in remaining:
+        for dim in key - output:
+            dim_to_keys[dim].add(key)
+
+    # Keep track of the number of tensors using each dim; when the dim is no longer
+    # used it can be contracted. Since we specialize to binary ops, we only care about
+    # ref counts of >=2 or >=3.
+    dim_ref_counts = {
+        count: {dim for dim, keys in dim_to_keys.items() if len(keys) >= count} - output
+        for count in [2, 3]
+    }
+
+    # Compute separable part of the objective function for contractions.
+    footprints = {key: compute_size_by_dict(key, sizes) for key in remaining}
+
+    # Find initial candidate contractions.
+    queue: list[GreedyContractionType] = []
+    for _dim, dim_keys in dim_to_keys.items():
+        dim_keys_list = sorted(dim_keys, key=remaining.__getitem__)
+        for i, k1 in enumerate(dim_keys_list[:-1]):
+            k2s_guess = dim_keys_list[1 + i :]
+            _push_candidate(
+                output,
+                sizes,
+                remaining,
+                footprints,
+                dim_ref_counts,
+                k1,
+                k2s_guess,
+                queue,
+                push_all,
+                cost_fn,
+            )
+
+    # Greedily contract pairs of tensors.
+    while queue:
+        con = choose_fn(queue, remaining)
+        if con is None:
+            continue  # allow choose_fn to flag all candidates obsolete
+        cost, k1, k2, k12 = con
+
+        ssa_id1 = remaining.pop(k1)
+        ssa_id2 = remaining.pop(k2)
+        for dim in k1 - output:
+            dim_to_keys[dim].remove(k1)
+        for dim in k2 - output:
+            dim_to_keys[dim].remove(k2)
+        ssa_path.append((ssa_id1, ssa_id2))
+
+        if k12 in remaining:
+            hadamard_id = next(ssa_ids)
+            ssa_path.append((remaining[k12], hadamard_id))
+            old_id = remaining[k12]
+            merged = (
+                ssa_to_subset[old_id] | ssa_to_subset[ssa_id1] | ssa_to_subset[ssa_id2]
+            )
+            ssa_to_subset[hadamard_id] = merged
+        else:
+            for dim in k12 - output:
+                dim_to_keys[dim].add(k12)
+        new_ssa_id = next(ssa_ids)
+        remaining[k12] = new_ssa_id
+        ssa_to_subset[new_ssa_id] = ssa_to_subset[ssa_id1] | ssa_to_subset[ssa_id2]
+        _update_ref_counts(dim_to_keys, dim_ref_counts, k1 | k2 - output)
+
+        footprints[k12] = compute_size_by_dict(k12, sizes)
+
+        # Find new candidate contractions.
+        k1 = k12
+        k2s = {k2 for dim in k1 for k2 in dim_to_keys[dim]}
+        k2s.discard(k1)
+        if k2s:
+            _push_candidate(
+                output,
+                sizes,
+                remaining,
+                footprints,
+                dim_ref_counts,
+                k1,
+                list(k2s),
+                queue,
+                push_all,
+                cost_fn,
+            )
+
+    # Greedily compute pairwise outer products.
+    final_queue = [
+        (compute_size_by_dict(key & output, sizes), ssa_id, key)
+        for key, ssa_id in remaining.items()
+    ]
+    heapq.heapify(final_queue)
+    _, ssa_id1, k1 = heapq.heappop(final_queue)
+    while final_queue:
+        _, ssa_id2, k2 = heapq.heappop(final_queue)
+        ssa_path.append((min(ssa_id1, ssa_id2), max(ssa_id1, ssa_id2)))
+        k12 = (k1 | k2) & output
+        cost = compute_size_by_dict(k12, sizes)
+        ssa_id12 = next(ssa_ids)
+        ssa_to_subset[ssa_id12] = ssa_to_subset[ssa_id1] | ssa_to_subset[ssa_id2]
+        _, ssa_id1, k1 = heapq.heappushpop(final_queue, (cost, ssa_id12, k12))
+
+    return ssa_path
+
+
+def greedy(
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+    memory_limit: int | None = None,
+    choose_fn: Any = None,
+    cost_fn: str = "memory-removed",
+    symmetry_oracle: SubgraphSymmetryOracle | None = None,
+) -> 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.
+        symmetry_oracle: Optional subgraph symmetry oracle.
+
+    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,
+            symmetry_oracle=symmetry_oracle,
+        )  # type: ignore
+
+    ssa_path = ssa_greedy_optimize(
+        inputs,
+        output,
+        size_dict,
+        cost_fn=cost_fn,
+        choose_fn=choose_fn,
+        symmetry_oracle=symmetry_oracle,
+    )
+    return ssa_to_linear(ssa_path)
+
+
+def _tree_to_sequence(tree: tuple[Any, ...]) -> PathType:
+    """Converts a contraction tree to a contraction path as it has to be
+    returned by path optimizers. A contraction tree can either be an int
+    (=no contraction) or a tuple containing the terms to be contracted. An
+    arbitrary number (>= 1) of terms can be contracted at once. Note that
+    contractions are commutative, e.g. (j, k, l) = (k, l, j). Note that in
+    general, solutions are not unique.
+
+    Parameters:
+        c: Contraction tree
+
+    Returns:
+        path: Contraction path
+
+    Examples:
+        ```python
+        _tree_to_sequence(((1,2),(0,(4,5,3))))
+        #> [(1, 2), (1, 2, 3), (0, 2), (0, 1)]
+        ```
+    """
+    if type(tree) == int:  # noqa: E721
+        return []
+
+    c: list[tuple[Any, ...]] = [
+        tree
+    ]  # list of remaining contractions (lower part of columns shown above)
+    t: list[int] = []  # list of elementary tensors (upper part of columns)
+    s: list[tuple[int, ...]] = []  # resulting contraction sequence
+
+    while len(c) > 0:
+        j = c.pop(-1)
+        s.insert(0, ())
+
+        for i in sorted([i for i in j if type(i) == int]):  # noqa: E721
+            s[0] += (sum(1 for q in t if q < i),)
+            t.insert(s[0][-1], i)
+
+        for i_tup in [i_tup for i_tup in j if type(i_tup) != int]:  # noqa: E721
+            s[0] += (len(t) + len(c),)
+            c.append(i_tup)
+
+    return s
+
+
+def _find_disconnected_subgraphs(
+    inputs: list[frozenset[int]], output: frozenset[int]
+) -> list[frozenset[int]]:
+    """Finds disconnected subgraphs in the given list of inputs. Inputs are
+    connected if they share summation indices. Note: Disconnected subgraphs
+    can be contracted independently before forming outer products.
+
+    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
+
+    Returns:
+        subgraphs: List containing sets of indices for each subgraph
+
+    Examples:
+        ```python
+        _find_disconnected_subgraphs([set("ab"), set("c"), set("ad")], set("bd"))
+        #> [{0, 2}, {1}]
+
+        _find_disconnected_subgraphs([set("ab"), set("c"), set("ad")], set("abd"))
+        #> [{0}, {1}, {2}]
+        ```
+    """
+    subgraphs = []
+    unused_inputs = set(range(len(inputs)))
+
+    i_sum = frozenset.union(*inputs) - output  # all summation indices
+
+    while len(unused_inputs) > 0:
+        g = set()
+        q = [unused_inputs.pop()]
+        while len(q) > 0:
+            j = q.pop()
+            g.add(j)
+            i_tmp = i_sum & inputs[j]
+            n = {k for k in unused_inputs if len(i_tmp & inputs[k]) > 0}
+            q.extend(n)
+            unused_inputs.difference_update(n)
+
+        subgraphs.append(g)
+
+    return [frozenset(x) for x in subgraphs]
+
+
+def _bitmap_select(
+    s: int, seq: list[frozenset[int]]
+) -> Generator[frozenset[int], None, None]:
+    """Select elements of ``seq`` which are marked by the bitmap set ``s``.
+
+    E.g.:
+
+        >>> list(_bitmap_select(0b11010, ['A', 'B', 'C', 'D', 'E']))
+        ['B', 'D', 'E']
+    """
+    return (x for x, b in zip(seq, bin(s)[:1:-1], strict=False) if b == "1")
+
+
+def _dp_calc_legs(g, all_tensors, s, inputs, i1_cut_i2_wo_output, i1_union_i2):
+    """Calculates the effective outer indices of the intermediate tensor
+    corresponding to the subgraph ``s``.
+    """
+    # set of remaining tensors (=g-s)
+    r = g & (all_tensors ^ s)
+    # indices of remaining indices:
+    if r:
+        i_r = frozenset.union(*_bitmap_select(r, inputs))
+    else:
+        i_r = frozenset()
+    # contraction indices:
+    i_contract = i1_cut_i2_wo_output - i_r
+    return i1_union_i2 - i_contract
+
+
+def _dp_compare_flops(
+    cost1: int,
+    cost2: int,
+    i1_union_i2: set[int],
+    size_dict: list[int],
+    cost_cap: int,
+    s1: int,
+    s2: int,
+    xn: dict[int, Any],
+    g: int,
+    all_tensors: int,
+    inputs: list[frozenset[int]],
+    i1_cut_i2_wo_output: set[int],
+    memory_limit: int | None,
+    contract1: int | tuple[int],
+    contract2: int | tuple[int],
+    get_ratio: Callable[[int, frozenset[int]], float] | None = None,
+) -> None:
+    """Performs the inner comparison of whether the two subgraphs (the bitmaps
+    `s1` and `s2`) should be merged and added to the dynamic programming
+    search. Will skip for a number of reasons:
+
+    1. If the number of operations to form `s = s1 | s2` including previous
+       contractions is above the cost-cap.
+    2. If we've already found a better way of making `s`.
+    3. If the intermediate tensor corresponding to `s` is going to break the
+       memory limit.
+
+    When a ``get_ratio`` closure is provided, the step cost is scaled by the
+    exact unique/dense symmetry ratio for the merged subset.
+    """
+    s = s1 | s2
+    i = _dp_calc_legs(g, all_tensors, s, inputs, i1_cut_i2_wo_output, i1_union_i2)
+
+    dense_step = compute_size_by_dict(i1_union_i2, size_dict)
+    if get_ratio is not None:
+        step_cost = int(dense_step * get_ratio(s, i))  # type: ignore[arg-type]
+    else:
+        step_cost = dense_step
+
+    cost = cost1 + cost2 + step_cost
+    if cost <= cost_cap:
+        if s not in xn or cost < xn[s][1]:
+            mem = compute_size_by_dict(i, size_dict)
+            if memory_limit is None or mem <= memory_limit:
+                xn[s] = (i, cost, (contract1, contract2))
+
+
+def _dp_compare_size(
+    cost1: int,
+    cost2: int,
+    i1_union_i2: set[int],
+    size_dict: list[int],
+    cost_cap: int,
+    s1: int,
+    s2: int,
+    xn: dict[int, Any],
+    g: int,
+    all_tensors: int,
+    inputs: list[frozenset[int]],
+    i1_cut_i2_wo_output: set[int],
+    memory_limit: int | None,
+    contract1: int | tuple[int],
+    contract2: int | tuple[int],
+    get_ratio: Callable[[int, frozenset[int]], float] | None = None,
+) -> None:
+    """Like `_dp_compare_flops` but sieves the potential contraction based
+    on the size of the intermediate tensor created, rather than the number of
+    operations, and so calculates that first.
+    """
+    s = s1 | s2
+    i = _dp_calc_legs(g, all_tensors, s, inputs, i1_cut_i2_wo_output, i1_union_i2)
+    mem = compute_size_by_dict(i, size_dict)
+    if get_ratio is not None:
+        mem = int(mem * get_ratio(s, i))  # type: ignore[arg-type]
+    cost = max(cost1, cost2, mem)
+    if cost <= cost_cap:
+        if s not in xn or cost < xn[s][1]:
+            if memory_limit is None or mem <= memory_limit:
+                xn[s] = (i, cost, (contract1, contract2))
+
+
+def _dp_compare_write(
+    cost1: int,
+    cost2: int,
+    i1_union_i2: set[int],
+    size_dict: list[int],
+    cost_cap: int,
+    s1: int,
+    s2: int,
+    xn: dict[int, Any],
+    g: int,
+    all_tensors: int,
+    inputs: list[frozenset[int]],
+    i1_cut_i2_wo_output: set[int],
+    memory_limit: int | None,
+    contract1: int | tuple[int],
+    contract2: int | tuple[int],
+    get_ratio: Callable[[int, frozenset[int]], float] | None = None,
+) -> None:
+    """Like ``_dp_compare_flops`` but sieves the potential contraction based
+    on the total size of memory created, rather than the number of
+    operations, and so calculates that first.
+    """
+    s = s1 | s2
+    i = _dp_calc_legs(g, all_tensors, s, inputs, i1_cut_i2_wo_output, i1_union_i2)
+    mem = compute_size_by_dict(i, size_dict)
+    if get_ratio is not None:
+        mem = int(mem * get_ratio(s, i))  # type: ignore[arg-type]
+    cost = cost1 + cost2 + mem
+    if cost <= cost_cap:
+        if s not in xn or cost < xn[s][1]:
+            if memory_limit is None or mem <= memory_limit:
+                xn[s] = (i, cost, (contract1, contract2))
+
+
+DEFAULT_COMBO_FACTOR = 64
+
+
+def _dp_compare_combo(
+    cost1: int,
+    cost2: int,
+    i1_union_i2: set[int],
+    size_dict: list[int],
+    cost_cap: int,
+    s1: int,
+    s2: int,
+    xn: dict[int, Any],
+    g: int,
+    all_tensors: int,
+    inputs: list[frozenset[int]],
+    i1_cut_i2_wo_output: set[int],
+    memory_limit: int | None,
+    contract1: int | tuple[int],
+    contract2: int | tuple[int],
+    factor: int | float = DEFAULT_COMBO_FACTOR,
+    combine: Callable = sum,
+    get_ratio: Callable[[int, frozenset[int]], float] | None = None,
+) -> None:
+    """Like ``_dp_compare_flops`` but sieves the potential contraction based
+    on some combination of both the flops and size,.
+    """
+    s = s1 | s2
+    i = _dp_calc_legs(g, all_tensors, s, inputs, i1_cut_i2_wo_output, i1_union_i2)
+    mem = compute_size_by_dict(i, size_dict)
+    f = compute_size_by_dict(i1_union_i2, size_dict)
+    if get_ratio is not None:
+        ratio = get_ratio(s, i)  # type: ignore[arg-type]
+        mem = int(mem * ratio)
+        f = int(f * ratio)
+    cost = cost1 + cost2 + combine((f, factor * mem))
+    if cost <= cost_cap:
+        if s not in xn or cost < xn[s][1]:
+            if memory_limit is None or mem <= memory_limit:
+                xn[s] = (i, cost, (contract1, contract2))
+
+
+minimize_finder = re.compile(r"(flops|size|write|combo|limit)-*(\d*)")
+
+
+@functools.lru_cache(128)
+def _parse_minimize(minimize: str | Callable) -> tuple[Callable, int | float]:
+    """This works out what local scoring function to use for the dp algorithm
+    as well as a `naive_scale` to account for the memory_limit checks.
+    """
+    if minimize == "flops":
+        return _dp_compare_flops, 1
+    elif minimize == "size":
+        return _dp_compare_size, 1
+    elif minimize == "write":
+        return _dp_compare_write, 1
+    elif callable(minimize):
+        # default to naive_scale=inf for this and remaining options
+        # as otherwise memory_limit check can cause problems
+        return minimize, float("inf")
+
+    # parse out a customized value for the combination factor
+    match = minimize_finder.fullmatch(minimize)
+    if match is None:
+        raise ValueError(f"Couldn't parse `minimize` value: {minimize}.")
+
+    minimize, custom_factor = match.groups()
+    factor = float(custom_factor) if custom_factor else DEFAULT_COMBO_FACTOR
+    if minimize == "combo":
+        return functools.partial(_dp_compare_combo, factor=factor, combine=sum), float(
+            "inf"
+        )
+    elif minimize == "limit":
+        return functools.partial(_dp_compare_combo, factor=factor, combine=max), float(
+            "inf"
+        )
+    else:
+        raise ValueError(f"Couldn't parse `minimize` value: {minimize}.")
+
+
+def simple_tree_tuple(seq: Sequence[tuple[int, ...]]) -> tuple[Any, ...]:
+    """Make a simple left to right binary tree out of iterable `seq`.
+
+    ```python
+    tuple_nest([1, 2, 3, 4])
+    #> (((1, 2), 3), 4)
+    ```
+
+    """
+    return functools.reduce(lambda x, y: (x, y), seq)  # type: ignore[arg-type]
+
+
+def _dp_parse_out_single_term_ops(
+    inputs: list[frozenset[int]],
+    all_inds: tuple[str, ...],
+    ind_counts: CounterType[str],
+) -> tuple[list[frozenset[int]], list[tuple[int]], list[int | tuple[int]]]:
+    """Take `inputs` and parse for single term index operations, i.e. where
+    an index appears on one tensor and nowhere else.
+
+    If a term is completely reduced to a scalar in this way it can be removed
+    to `inputs_done`. If only some indices can be summed then add a 'single
+    term contraction' that will perform this summation.
+    """
+    i_single = frozenset(i for i, c in enumerate(all_inds) if ind_counts[c] == 1)
+    inputs_parsed: list[frozenset[int]] = []
+    inputs_done: list[tuple[int]] = []
+    inputs_contractions: list[int | tuple[int]] = []
+    for j, i in enumerate(inputs):
+        i_reduced = i - i_single
+        if (not i_reduced) and (len(i) > 0):
+            # input reduced to scalar already - remove
+            inputs_done.append((j,))
+        else:
+            # if the input has any index reductions, add single contraction
+            inputs_parsed.append(i_reduced)
+            inputs_contractions.append((j,) if i_reduced != i else j)
+
+    return inputs_parsed, inputs_done, inputs_contractions
+
+
+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: bool | int = True,
+        search_outer: bool = False,
+    ) -> None:
+        self.minimize = minimize
+        self.search_outer = search_outer
+        self.cost_cap = cost_cap
+
+    def __call__(  # type: ignore[override]
+        self,
+        inputs_: list[ArrayIndexType],
+        output_: ArrayIndexType,
+        size_dict_: dict[str, int],
+        memory_limit_: int | None = None,
+        symmetry_oracle: SubgraphSymmetryOracle | None = 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
+
+        # Build a per-call bitmap->frozenset[int] converter for the oracle.
+        # Maps DP bitmap `s` (bit j set iff tensor j is in the subset) to the
+        # original operand-position frozenset that the oracle was built with.
+        if symmetry_oracle is not None:
+            _bts_cache: dict[int, frozenset[int]] = {}
+
+            def bitmap_to_subset(
+                s: int, _cache: dict[int, frozenset[int]] = _bts_cache
+            ) -> frozenset[int]:
+                if s not in _cache:
+                    result: set[int] = set()
+                    for k in range(len(inputs)):
+                        if s >> k & 1:
+                            orig = inputs_contractions[k]
+                            result.add(orig if isinstance(orig, int) else orig[0])
+                    _cache[s] = frozenset(result)
+                return _cache[s]
+
+            _ratio_cache: dict[int, float] = {}
+
+            def get_ratio(
+                s: int,
+                int_output_legs: frozenset[int],
+                _cache: dict[int, float] = _ratio_cache,
+            ) -> float:
+                """Exact unique/dense ratio for the intermediate at DP bitmap s.
+
+                Lazily computed on first access and cached per subset. Returns
+                1.0 when the oracle reports no symmetry, or when the
+                intermediate has no elements. The int<->str label translation
+                happens once per cache miss and is amortized across all
+                _dp_compare_* helper calls that subsequently reuse this
+                subset's ratio.
+
+                all_inds[ix] is the inverse of symbol2int: the string label
+                for int label ix. Only the labels in this intermediate need
+                translation, keeping the per-miss work bounded.
+                """
+                cached = _cache.get(s, -1.0)
+                if cached >= 0.0:
+                    return cached
+                subset = bitmap_to_subset(s)
+                subset_sym = symmetry_oracle.sym(subset)
+                sym = subset_sym.output
+                if sym is None:
+                    _cache[s] = 1.0
+                    return 1.0
+                str_legs = frozenset(all_inds[ix] for ix in int_output_legs)
+                str_size_dict = {all_inds[ix]: size_dict[ix] for ix in int_output_legs}
+                dense = compute_size_by_dict(str_legs, str_size_dict)
+                if dense <= 0:
+                    _cache[s] = 1.0
+                    return 1.0
+                unique = unique_elements(str_legs, str_size_dict, perm_group=sym)
+                ratio = unique / dense
+                _cache[s] = ratio
+                return ratio
+
+        else:
+            bitmap_to_subset = None  # type: ignore[assignment]
+            get_ratio = None  # type: ignore[assignment]
+
+        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,
+                                            get_ratio=get_ratio,
+                                        )
+
+                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)
+
+
+def dynamic_programming(
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+    memory_limit: int | None = None,
+    symmetry_oracle: SubgraphSymmetryOracle | None = None,
+    **kwargs: Any,
+) -> PathType:
+    optimizer = DynamicProgramming(**kwargs)
+    return optimizer(
+        inputs, output, size_dict, memory_limit, symmetry_oracle=symmetry_oracle
+    )
+
+
+_AUTO_CHOICES = {}
+for i in range(1, 5):
+    _AUTO_CHOICES[i] = optimal
+for i in range(5, 7):
+    _AUTO_CHOICES[i] = branch_all
+for i in range(7, 9):
+    _AUTO_CHOICES[i] = branch_2
+for i in range(9, 15):
+    _AUTO_CHOICES[i] = branch_1
+
+
+def auto(
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+    memory_limit: int | None = None,
+    symmetry_oracle: SubgraphSymmetryOracle | None = None,
+) -> PathType:
+    """Auto-select based on number of inputs. All routed optimizers
+    accept ``symmetry_oracle``; no silent fallback."""
+    fn = _AUTO_CHOICES.get(len(inputs), greedy)
+    return fn(
+        inputs,
+        output,
+        size_dict,
+        memory_limit,
+        symmetry_oracle=symmetry_oracle,
+    )
+
+
+_AUTO_HQ_CHOICES = {}
+for i in range(1, 6):
+    _AUTO_HQ_CHOICES[i] = optimal
+for i in range(6, 17):
+    _AUTO_HQ_CHOICES[i] = dynamic_programming
+
+
+def auto_hq(
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+    memory_limit: int | None = None,
+    symmetry_oracle: SubgraphSymmetryOracle | None = None,
+) -> PathType:
+    """Auto-HQ selection based on number of inputs. All routed
+    optimizers accept ``symmetry_oracle``; no silent fallback."""
+    from ._path_random import random_greedy_128
+
+    fn = _AUTO_HQ_CHOICES.get(len(inputs), random_greedy_128)
+    return fn(
+        inputs,
+        output,
+        size_dict,
+        memory_limit,
+        symmetry_oracle=symmetry_oracle,
+    )
+
+
+_PATH_OPTIONS: dict[str, PathSearchFunctionType] = {
+    "auto": auto,
+    "auto-hq": auto_hq,
+    "optimal": optimal,
+    "branch-all": branch_all,
+    "branch-2": branch_2,
+    "branch-1": branch_1,
+    "greedy": greedy,
+    "eager": greedy,
+    "opportunistic": greedy,
+    "dp": dynamic_programming,
+    "dynamic-programming": dynamic_programming,
+}
+
+
+def register_path_fn(name: str, fn: PathSearchFunctionType) -> None:
+    """Add path finding function ``fn`` as an option with ``name``."""
+    if name in _PATH_OPTIONS:
+        raise KeyError(f"Path optimizer '{name}' already exists.")
+
+    _PATH_OPTIONS[name.lower()] = fn
+
+
+def get_path_fn(path_type: str) -> PathSearchFunctionType:
+    """Get the correct path finding function from str ``path_type``."""
+    path_type = path_type.lower()
+    if path_type not in _PATH_OPTIONS:
+        raise KeyError(
+            f"Path optimizer '{path_type}' not found, valid options are {set(_PATH_OPTIONS.keys())}."
+        )
+
+    return _PATH_OPTIONS[path_type]