Compressed contraction path finding using branch and bound.



A contraction tree for compressed contractions. Currently the only


Exhaustively search all possible contraction orders for a given



ssa_to_linear(ssa_path[, N])

Convert a path with static single assignment ids to a path with recycled


get_hypergraph(inputs[, output, size_dict, accel])

Single entry-point for creating a, possibly accelerated, HyperGraph.

do_reconfigure(tree, time, chi)

Module Contents

cotengra.experimental.path_compressed_branchbound.ssa_to_linear(ssa_path, N=None)

Convert a path with static single assignment ids to a path with recycled linear ids. For example:

>>> ssa_to_linear([(0, 3), (2, 4), (1, 5)])
[(0, 3), (1, 2), (0, 1)]
cotengra.experimental.path_compressed_branchbound.get_hypergraph(inputs, output=None, size_dict=None, accel=False)

Single entry-point for creating a, possibly accelerated, HyperGraph.

class cotengra.experimental.path_compressed_branchbound.ContractionTreeCompressed(inputs, output, size_dict, track_childless=False, track_flops=False, track_write=False, track_size=False, objective=None)

Bases: ContractionTree

A contraction tree for compressed contractions. Currently the only difference is that this defaults to the ‘surface’ traversal ordering.

classmethod from_path(inputs, output, size_dict, *, path=None, ssa_path=None, check=False, **kwargs)

Create a (completed) ContractionTreeCompressed from the usual inputs plus a standard contraction path or ‘ssa_path’ - you need to supply one. This also set the default ‘surface’ traversal ordering to be the initial path.


Get the objective function for this tree.

abstract get_contractor(*_, **__)

Get a reusable function which performs the contraction corresponding to this tree, cached.

  • tree (ContractionTree) – The contraction tree.

  • order (str or callable, optional) – Supplied to ContractionTree.traverse(), the order in which to perform the pairwise contractions given by the tree.

  • prefer_einsum (bool, optional) – Prefer to use einsum for pairwise contractions, even if tensordot can perform the contraction.

  • strip_exponent (bool, optional) – If True, the function will strip the exponent from the output array and return it separately.

  • check_zero (bool, optional) – If True, when strip_exponent=True, explicitly check for zero-valued intermediates that would otherwise produce nan, instead terminating early if encountered and returning (0.0, 0.0).

  • implementation (str or tuple[callable, callable], optional) –

    What library to use to actually perform the contractions. Options are:

    • None: let cotengra choose.

    • ”autoray”: dispatch with autoray, using the tensordot and einsum implementation of the backend.

    • ”cotengra”: use the tensordot and einsum implementation of cotengra, which is based on batch matrix multiplication. This is faster for some backends like numpy, and also enables libraries which don’t yet provide tensordot and einsum to be used.

    • ”cuquantum”: use the cuquantum library to perform the whole contraction (not just individual contractions).

    • tuple[callable, callable]: manually supply the tensordot and einsum implementations to use.

  • autojit (bool, optional) – If True, use autoray.autojit() to compile the contraction function.

  • progbar (bool, optional) – Whether to show progress through the contraction by default.


fn – The contraction function, with signature fn(*arrays).

Return type:


class cotengra.experimental.path_compressed_branchbound.CompressedExhaustive(minimize, max_nodes=float('inf'), max_time=None, local_score=None, exploration_power=0.0, best_score=None, progbar=False)

Exhaustively search all possible contraction orders for a given compressed bond dimension, using several forms of pruning and local hueristics to accelerate the search.

  • chi (int) – The max bond dimension of the compressed hypergraph.

  • max_nodes (int, optional) – Set the maximum number of contraction steps to consider.

  • max_time (float, optional) – Set the maximum time to spend on the search.

  • local_score (callable, optional) –

    A function that assigns a score to a potential contraction, with a lower score giving more priority to explore that contraction earlier. It should have signature:

    local_score(step : int, tracker: CompressedStatsTracker) -> float

    where step is the number of steps so far, tracker is a CompressedStatsTracker object that tracks various properties like flops and peak size.

  • exploration_power (float, optional) – If not 0.0, the inverse power to which the step is raised in the default local score function. Higher values favor exploring more promising branches early on - at the cost of increased memory. Ignored if local_score is supplied.

  • best_score (float, optional) – Manually specify an upper bound for best score found so far.

  • progbar (bool, optional) – If True, display a progress bar.

setup(inputs, output, size_dict)

Set-up the optimizer with a specific contraction.

expand_node(i, j, hg, tree_map, ssa_path, tracker, high_priority=False)

Given a current contraction node, expand it by contracting nodes i and j.

  • i (int) – The nodes to contract.

  • j (int) – The nodes to contract.

  • hg (Hypergraph) – The hypergraph to contract.

  • ssa_path (list) – The contraction path so far.

  • tracker (CompressedStatsTracker) – Scoring object that tracks costs of current compressed contraction.

  • high_priority (bool, optional) – If True, the contraction will be assessed before any other normal contractions.


The contraction index, or None if the contraction is guaranteed to be worse than another contraction.

Return type:

int or None

_update_progbar(pbar, c)
run(inputs, output, size_dict)
property ssa_path
property path
explore_path(path, high_priority=True, restrict=False)

Explicitly supply a path to be added to the search space, by default it is added to the priority queue and will be processed first.

  • path (sequence[tuple[int]]) – A contraction path to explore.

  • high_priority (bool, optional) – If True, the path will be assessed before anything else, regardless of cost - the default.

  • restrict (bool, optional) – If True, only allow contractions in this path, so only the order will be optimized.

search(inputs, output, size_dict)

Run and return the best ContractionTreeCompressed.

__call__(inputs, output, size_dict)

Run and return the best path.

cotengra.experimental.path_compressed_branchbound.do_reconfigure(tree, time, chi)
class cotengra.experimental.path_compressed_branchbound.CompressedTreeRefiner(trees, copt, chi, max_refine_time=8, executor=None, pre_dispatch=8, progbar=False, plot=False)
_check_score(key, tree, score=None)
_process_result(tree, key, time, old, new)
refine(num_its=None, bins=30)