cotengra

Hyper optimized contraction trees for large tensor networks and einsums.

Submodules

• cotengra._version
• cotengra.contract
• cotengra.core
• cotengra.core_multi
• cotengra.experimental
• cotengra.hypergraph
• cotengra.hyperoptimizers
• cotengra.interface
• cotengra.nodeops
• cotengra.oe
• cotengra.parallel
• cotengra.pathfinders
• cotengra.plot
• cotengra.presets
• cotengra.reusable
• cotengra.schematic
• cotengra.scoring
• cotengra.slicer
• cotengra.utils

Attributes

greedy_optimize
optimal_optimize
optimal_outer_optimize
UniformOptimizer	Does no gaussian process tuning by default, just randomly samples - requires
contract_expression	Alias for cotengra.einsum_expression().
contract	Alias for cotengra.einsum().

Classes

ContractionTree	Binary tree representing a tensor network contraction.
ContractionTreeCompressed	A contraction tree for compressed contractions. Currently the only
ContractionTreeMulti	Binary tree representing a tensor network contraction.
HyperGraph	Simple hypergraph builder and writer.
HyperCompressedOptimizer	A compressed contraction path optimizer that samples a series of ordered
HyperMultiOptimizer	A path optimizer that samples a series of contraction trees
HyperOptimizer	A path optimizer that samples a series of contraction trees
ReusableHyperCompressedOptimizer	Like HyperCompressedOptimizer but it will re-instantiate the
ReusableHyperOptimizer	Like HyperOptimizer but it will re-instantiate the optimizer
GreedyOptimizer	Class interface to the greedy optimizer which can be instantiated with
OptimalOptimizer	Class interface to the optimal optimizer which can be instantiated with
RandomGreedyOptimizer	Lightweight random greedy optimizer, that eschews hyper parameter
ReusableRandomGreedyOptimizer	A reusable random greedy optimizer that caches path per contraction.
FlowCutterOptimizer
QuickBBOptimizer
RandomOptimizer	A fully random pathfinder, that randomly selects pairs of tensors to
AutoHQOptimizer	An optimizer that automatically chooses between optimal and
AutoOptimizer	An optimizer that automatically chooses between optimal and
SliceFinder	An object to help find the best indices to slice over in order to reduce

Functions

get_hypergraph(inputs[, output, size_dict, accel])	Single entry-point for creating a, possibly accelerated, HyperGraph.
get_hyper_space()
list_hyper_functions()	Return a list of currently registered hyper contraction finders.
array_contract(arrays, inputs[, output, optimize, ...])	Perform the tensor contraction specified by inputs, output and
array_contract_expression(inputs[, output, size_dict, ...])	Get an callable 'expression' that will contract tensors with indices and
array_contract_path(inputs[, output, size_dict, ...])	Find only the contraction path for the specific contraction, with fast
array_contract_tree(inputs[, output, size_dict, ...])	Get the ContractionTree for the tensor contraction specified by
einsum(*args[, optimize, strip_exponent, ...])	Perform an einsum contraction, using cotengra, using strategy given by
einsum_expression(*args[, optimize, constants, cache])	Get an callable 'expression' that will contract tensors with shapes
einsum_tree(*args[, optimize, canonicalize, ...])	Get the ContractionTree for the einsum equation eq and
ncon(arrays, indices, **kwargs)	Perform a contraction specified by the ncon style indices, using
register_preset(preset, optimizer[, optimizer_tree, ...])	Register a preset optimizer.
edge_path_to_linear(edge_path, inputs)	Convert a path specified by a sequence of edge indices to a path with
edge_path_to_ssa(edge_path, inputs)	Convert a path specified by a sequence of edge indices to a path with
linear_to_ssa(path[, N])	Convert a path with recycled linear ids to a path with static single
ssa_to_linear(ssa_path[, N])	Convert a path with static single assignment ids to a path with recycled
optimize_flowcutter(inputs, output, size_dict[, ...])
optimize_quickbb(inputs, output, size_dict[, ...])
plot_contractions(tree[, order, color_size, ...])
plot_contractions_alt(tree[, x, y, color, size, ...])
plot_scatter(self[, x, y, ylim, cumulative_time, ...])
plot_scatter_alt(self[, x, y, color, color_scheme, ...])	Plot the trials total flops vs max size.
plot_slicings(slice_finder[, color_scheme, ...])
plot_slicings_alt(slice_finder[, color_scheme, ...])
plot_tree(tree[, layout, hypergraph_layout, ...])	Plot a contraction tree using matplotlib.
plot_tree_ring(tree, **kwargs)
plot_tree_span(tree, **kwargs)
plot_tree_tent(tree, **kwargs)
plot_trials(self, *[, x, y, figsize, cumulative_time, ...])
plot_trials_alt(self[, y, width, height])	Plot the trials interactively using altair.
auto_hq_optimize(inputs, output, size_dict, **kwargs)
auto_optimize(inputs, output, size_dict, **kwargs)
hash_contraction(inputs, output, size_dict[, method])	Compute a hash for a particular contraction geometry.
get_symbol(i)	Get the symbol corresponding to int i - runs through the usual 52
get_symbol_map(inputs)	Get a mapping of arbitrary hashable 'indices' to single unicode symbols,
hyper_optimize(inputs, output, size_dict[, ...])

Package Contents

class cotengra.ContractionTree(inputs, output, size_dict, track_childless=False, track_flops=False, track_write=False, track_size=False, objective=None, nodeops='auto')

Binary tree representing a tensor network contraction.

Parameters:

• inputs (sequence of str) – The list of input tensor’s indices.

• output (str) – The output indices.

• size_dict (dict[str, int]) – The size of each index.

• track_childless (bool, optional) – Whether to dynamically keep track of which nodes are childless. Useful if you are ‘divisively’ building the tree.

• track_flops (bool, optional) – Whether to dynamically keep track of the total number of flops. If False You can still compute this once the tree is complete.

• track_write (bool, optional) – Whether to dynamically keep track of the total number of elements written. If False You can still compute this once the tree is complete.

• track_size (bool, optional) – Whether to dynamically keep track of the largest tensor so far. If False You can still compute this once the tree is complete.

• objective (str or Objective, optional) – An default objective function to use for further optimization and scoring, for example reconfiguring or computing the combo cost. If not supplied the default is to create a flops objective when needed.

children

Mapping of each node to two children.

Type:

dict[node, tuple[node]]

info

Information about the tree nodes. The key is the set of inputs (a set of inputs indices) the node contains. Or in other words, the subgraph of the node. The value is a dictionary to cache information about effective ‘leg’ indices, size, flops of formation etc.

Type:

dict[node, dict]

inputs

output

N

appearances

preprocessing

children

nodeops

info

root

track_childless = False

_track_flops = False

_track_write = False

_track_size = False

already_optimized

multiplicity = 1

sliced_inds

sliced_inputs

contraction_cores

_default_objective = None

set_state_from(other)

Set the internal state of this tree to that of other.

copy()

Create a copy of this ContractionTree.

set_default_objective(objective)

Set the objective function for this tree.

get_default_objective()

Get the objective function for this tree.

get_default_combo_factor()

Get the default combo factor for this tree.

get_score(objective=None)

Score this tree using the default objective function.

property nslices

Simple alias for how many independent contractions this tree represents overall.

property nchunks

The number of ‘chunks’ - determined by the number of sliced output indices.

input_to_node(i)

Create a node from a single input index, i.e. the subgraph that only contains the input tensor i.

Parameters:

i (int) – The input index.

Returns:

node

Return type:

node_type

node_to_input(node)

Assuming node has one element, i.e. is a leaf, return the corresponding input index.

Parameters:

node (node_type) – The node to convert.

Returns:

i

Return type:

int

node_to_terms(node)

Turn a node into the corresponding terms a sequence of leaf legs, corresponding to input indices.

gen_leaves()

Generate the nodes representing leaves of the contraction tree, i.e. of size 1 each corresponding to a single input tensor.

get_incomplete_nodes()

Get the set of current nodes that have no children and the set of nodes that have no parents. These are the ‘childless’ and ‘parentless’ nodes respectively, that need to be contracted to complete the tree. The parentless nodes are grouped into the childless nodes that contain them as subgraphs.

Returns:

groups – A mapping of childless nodes to the list of parentless nodes are beneath them.

Return type:

dict[node_type, list[node_type]]

See also

autocomplete

autocomplete(**contract_opts)

Contract all remaining node groups (as computed by tree.get_incomplete_nodes) in the tree to complete it.

Parameters:

contract_opts – Options to pass to tree.contract_nodes.

See also

get_incomplete_nodes, contract_nodes

classmethod from_path(inputs, output, size_dict, *, path=None, ssa_path=None, edge_path=None, optimize='auto', autocomplete='auto', check=False, **kwargs)

Create a (completed) ContractionTree from the usual inputs plus a standard contraction path or ‘ssa_path’ - you need to supply one.

Parameters:

• inputs (Sequence[Sequence[str]]) – The input indices of each tensor, as single unicode characters.

• output (Sequence[str]) – The output indices.

• size_dict (dict[str, int]) – The size of each index.

• path (Sequence[Sequence[int]], optional) – The contraction path, a sequence of pairs of tensor ids to contract. The ids are linear indices into the list of temporary tensors, which are recycled as each contraction pops a pair and appends the result. One of path, ssa_path or edge_path must be supplied.

• ssa_path (Sequence[Sequence[int]], optional) – The contraction path, a sequence of pairs of indices to contract. The indices are single use, as if the result of each contraction is appended to the end of the list of temporary tensors without popping. One of path, ssa_path or edge_path must be supplied.

• edge_path (Sequence[str], optional) – The contraction path, a sequence of indices to contract in order. One of path, ssa_path or edge_path must be supplied.

• optimize (str, optional) – If a contraction within the path contains 3 or more tensors, how to optimize this subcontraction into a binary tree.

• autocomplete ("auto" or bool, optional) – Whether to automatically complete the path, i.e. contract all remaining nodes. If “auto” then a warning is issued if the path is not complete.

• check (bool, optional) – Whether to perform some basic checks while creating the contraction nodes.

Return type:

ContractionTree

classmethod from_info(info, **kwargs)

Create a ContractionTree from an opt_einsum.PathInfo object.

classmethod from_eq(eq, size_dict, **kwargs)

Create a empty ContractionTree directly from an equation and set of shapes.

Parameters:

• eq (str) – The einsum string equation.

• size_dict (dict[str, int]) – The size of each index.

get_eq()

Get the einsum equation corresponding to this tree. Note that this is the total (or original) equation, so includes indices which have been sliced.

Returns:

eq

Return type:

str

get_shapes()

Get the shapes of the input tensors corresponding to this tree.

Returns:

shapes

Return type:

tuple[tuple[int]]

get_inputs_sliced()

Get the input indices corresponding to a single slice of this tree, i.e. with sliced indices removed.

Returns:

inputs

Return type:

tuple[tuple[str]]

get_output_sliced()

Get the output indices corresponding to a single slice of this tree, i.e. with sliced indices removed.

Returns:

output

Return type:

tuple[str]

get_eq_sliced()

Get the einsum equation corresponding to a single slice of this tree, i.e. with sliced indices removed.

Returns:

eq

Return type:

str

get_shapes_sliced()

Get the shapes of the input tensors corresponding to a single slice of this tree, i.e. with sliced indices removed.

Returns:

shapes

Return type:

tuple[tuple[int]]

classmethod from_edge_path(edge_path, inputs, output, size_dict, optimize='auto', autocomplete='auto', check=False, **kwargs)

Create a ContractionTree from an edge elimination ordering.

_add_node(node, check=False, **kwargs)

Add a node to this tree, specified either directly as a existing node type, or as a subgraph (i.e. a sequence of input positions) which is then converted to a node with the corresponding extent and subgraph information.

Note if “ssa” nodes are used, then adding two equivalent subgraphs will result in two new nodes, since the node labels do not themselves encode the subgraph information.

Parameters:

• node (node_type or Sequence[int]) – The node to add, either directly as a node type, or as a subgraph specified by the sequence of input positions it contains.

• check (bool, optional) – Whether to perform some basic checks on the node and tree state before adding the node.

• kwargs (dict, optional) – Additional information to cache about this node, for example its ‘extent’ or ‘subgraph’. If it is being specified as a sequence of input positions, these two will be injected automatically.

Returns:

node – The node that was added, which may be different from the input if the input was specified as a sequence of input positions.

Return type:

node_type

_remove_node(node)

Remove node from this tree and update the flops and maximum size if tracking them respectively, as well as input pre-processing.

compute_leaf_legs(i)

Compute the effective ‘outer’ indices for the ith input tensor. This is not always simply the ith input indices, due to A) potential slicing and B) potential preprocessing.

has_preprocessing()

has_hyper_indices()

Check if there are any ‘hyper’ indices in the contraction, i.e. indices that don’t appear exactly twice, when considering the inputs and output.

get_extent(node)

Get the number of input tensors contained in the subgraph represented by node.

Parameters:

node (node_type) – The node to compute the extent of.

Returns:

extent

Return type:

int

get_subgraph(node) → tuple[int, Ellipsis]

Get the sequence of input tensors contained in subgraph represented by node.

Parameters:

node (node_type) – The node to compute the subgraph of.

Returns:

subgraph – The input tensor indices contained in this subgraph.

Return type:

tuple[int]

get_legs(node)

Get the effective ‘outer’ indices for the collection of tensors in node.

get_involved(node)

Get all the indices involved in the formation of subgraph node.

get_size(node)

Get the tensor size of node.

get_flops(node)

Get the FLOPs for the pairwise contraction that will create node.

get_can_dot(node)

Get whether this contraction can be performed as a dot product (i.e. with tensordot), or else requires einsum, as it has indices that don’t appear exactly twice in either the inputs or the output.

get_inds(node)

Get the indices of this node - an ordered string version of get_legs that starts with tree.inputs and maintains the order they appear in each contraction ‘ABC,abc->ABCabc’, to match tensordot.

get_tensordot_axes(node)

Get the axes arg for a tensordot ocontraction that produces node. The pairs are sorted in order of appearance on the left input.

get_tensordot_perm(node)

Get the permutation required, if any, to bring the tensordot output of this nodes contraction into line with self.get_inds(node).

get_einsum_eq(node)

Get the einsum string describing the contraction that produces node, unlike get_inds the characters are mapped into [a-zA-Z], for compatibility with numpy.einsum for example.

is_leaf(node)

Check if node is a leaf node in this tree.

Parameters:

node (node_type) – The node to check.

Returns:

is_leaf

Return type:

bool

is_root(node)

Check if node is the root node in this tree.

Parameters:

node (node_type) – The node to check.

Returns:

is_root

Return type:

bool

is_descendant(node, ancestor)

Check if node is a descendant of ancestor in this tree.

Parameters:

• node (node_type) – The node to check.

• ancestor (node_type) – The potential ancestor node.

Returns:

is_descendant

Return type:

bool

get_peak_size(node)

Get the peak size for all but only the contractions required to produce node. The value for the root note will be the peak size of the entire contraction.

Parameters:

node (node_type) – The node to compute the peak size of.

Returns:

peak_size

Return type:

int

reorder_for_peak_size()

This reorders the depth first traversal of the tree to minimize the peak size of the contraction.

get_centrality(node)

total_flops(dtype=None, log=None)

Sum the flops contribution from every node in the tree.

Parameters:

dtype ({'float', 'complex', None}, optional) – Scale the answer depending on the assumed data type.

total_write()

Sum the total amount of memory that will be created and operated on.

combo_cost(factor=DEFAULT_COMBO_FACTOR, combine=sum, log=None)

total_cost

max_size(log=None)

The size of the largest intermediate tensor.

max_contraction_size(log=None)

The maximum size of a single contraction in the tree. This includes the size of the two input tensors and the output tensor, and can be a more practical measure of the peak memory required.

Parameters:

log (float, optional) – If provided, return the log of the size to this base.

Returns:

size – The maximum size of a single contraction in the tree, or its log.

Return type:

int or float

peak_size(order=None, log=None)

Get the peak concurrent size of tensors needed - this depends on the traversal order, i.e. the exact contraction path, not just the contraction tree.

contract_stats(force=False)

Simulteneously compute the total flops, write and size of the contraction tree. This is more efficient than calling each of the individual methods separately. Once computed, each quantity is then automatically tracked.

Returns:

stats – The total flops, write and size.

Return type:

dict[str, int]

arithmetic_intensity()

The ratio of total flops to total write - the higher the better for extracting good computational performance.

contraction_scaling()

This is computed simply as the maximum number of indices involved in any single contraction, which will match the scaling assuming that all dimensions are equal.

contraction_cost(log=None)

Get the total number of scalar operations ~ time complexity.

naive_cost(log=None)

Get the naive cost of performing this contraction as a single einsum summation, without any intermediate contractions. This is given the as product of the size of all indices.

Parameters:

log (float, optional) – If provided, return log of the cost to this base.

speedup(log=None)

Speedup compared to naive summation.

Parameters:

log (float, optional) – If provided, return log of the speedup to this base.

contraction_width(log=2)

Get log2 of the size of the largest tensor.

compressed_contract_stats(chi=None, order='surface_order', compress_late=None)

total_flops_compressed(chi=None, order='surface_order', compress_late=None, dtype=None, log=None)

Estimate the total flops for a compressed contraction of this tree with maximum bond size chi. This includes basic estimates of the ops to perform contractions, QRs and SVDs.

contraction_cost_compressed

total_write_compressed(chi=None, order='surface_order', compress_late=None, log=None)

Compute the total size of all intermediate tensors when a compressed contraction is performed with maximum bond size chi, ordered by order. This is relevant maybe for time complexity and e.g. autodiff space complexity (since every intermediate is kept).

combo_cost_compressed(chi=None, order='surface_order', compress_late=None, factor=None, log=None)

total_cost_compressed

max_size_compressed(chi=None, order='surface_order', compress_late=None, log=None)

Compute the maximum sized tensor produced when a compressed contraction is performed with maximum bond size chi, ordered by order. This is close to the ideal space complexity if only tensors that are being directly operated on are kept in memory.

peak_size_compressed(chi=None, order='surface_order', compress_late=None, accel='auto', log=None)

Compute the peak size of combined intermediate tensors when a compressed contraction is performed with maximum bond size chi, ordered by order. This is the practical space complexity if one is not swapping intermediates in and out of memory.

contraction_width_compressed(chi=None, order='surface_order', compress_late=None, log=2)

Compute log2 of the maximum sized tensor produced when a compressed contraction is performed with maximum bond size chi, ordered by order.

_update_tracked(node)

contract_nodes_pair(x, y, legs=None, cost=None, size=None, parent=None, check=False)

Contract node x with node y in the tree to create a new parent node, which is returned.

Parameters:

• x (node_type) – The first node to contract.

• y (node_type) – The second node to contract.

• legs (dict[str, int], optional) – The effective ‘legs’ of the new node if already known. If not given, this is computed from the inputs of x and y.

• cost (int, optional) – The cost of the contraction if already known. If not given, this is computed from the inputs of x and y.

• size (int, optional) – The size of the new node if already known. If not given, this is computed from the inputs of x and y.

• check (bool, optional) – Whether to check the inputs are valid.

Returns:

parent – The new parent node of x and y.

Return type:

node_type

contract_nodes(nodes, optimize='auto', grandparent=None, check=False, extra_opts=None)

Contract an arbitrary number of nodes in the tree to build up a subtree. The root of this subtree (a new intermediate) is returned.

is_complete()

Check every node has two children, unless it is a leaf.

get_default_order()

_traverse_dfs()

Traverse the tree in a depth first, non-recursive, order.

_traverse_ordered(order)

Traverse the tree in the order that minimizes order(node), but still constrained to produce children before parents.

traverse(order=None)

Generate, in order, all the node merges in this tree. Non-recursive! This ensures children are always visited before their parent.

Parameters:

order (None, "dfs", or callable, optional) – How to order the contractions within the tree. If a callable is given (which should take a node as its argument), try to contract nodes that minimize this function first.

Returns:

The bottom up ordered sequence of tree merges, each a tuple of (parent, left_child, right_child).

Return type:

generator[tuple[node]]

See also

descend

descend(mode='dfs')

Generate, from root to leaves, all the node merges in this tree. Non-recursive! This ensures parents are visited before their children.

Parameters:

mode ({'dfs', bfs}, optional) – How expand from a parent.

Returns:

The top down ordered sequence of tree merges, each a tuple of (parent, left_child, right_child).

Return type:

generator[tuple[node]

See also

traverse

get_subtree(node, size, search='bfs', seed=None)

Get a subtree spanning down from node which will have size leaves (themselves not necessarily leaves of the actual tree).

Parameters:

• node (node_type) – The node of the tree to start with.

• size (int) – How many subtree leaves to aim for.

• search ({'bfs', 'dfs', 'random'}, optional) –

  How to build the tree:

  • ’bfs’: breadth first expansion

  • ’dfs’: depth first expansion (largest nodes first)

  • ’random’: random expansion

• seed (None, int or random.Random, optional) – Random number generator seed, if search is ‘random’.

Returns:

• sub_leaves (tuple[node_type]) – Nodes which are subtree leaves.

• branches (tuple[node_type]) – Nodes which are between the subtree leaves and root.

remove_ind(ind, project=None, inplace=False)

Remove (i.e. by default slice) index ind from this contraction tree, taking care to update all relevant information about each node.

remove_ind_

restore_ind(ind, inplace=False)

Restore (unslice or un-project) index ind to this contraction tree, taking care to update all relevant information about each node.

Parameters:

• ind (str) – The index to restore.

• inplace (bool, optional) – Whether to perform the restoration inplace or not.

Return type:

ContractionTree

restore_ind_

unslice_rand(seed=None, inplace=False)

Unslice (restore) a random index from this contraction tree.

Parameters:

• seed (None, int or random.Random, optional) – Random number generator seed.

• inplace (bool, optional) – Whether to perform the unslicing inplace or not.

Return type:

ContractionTree

unslice_rand_

unslice_all(inplace=False)

Unslice (restore) all sliced indices from this contraction tree.

Parameters:

inplace (bool, optional) – Whether to perform the unslicing inplace or not.

Return type:

ContractionTree

unslice_all_

calc_subtree_candidates(pwr=2, what='flops')

_subtree_remove_and_optimize(sub_root, sub_leaves, sub_branches, already_optimized, node_cost, minimize, opt, pbar)

_subtree_reconfigure_descend(subtree_size, subtree_search, maxiter, seed, minimize, opt, already_optimized, node_cost, pbar)

_subtree_reconfigure_rand_select(subtree_size, subtree_search, weight_what, weight_pwr, select, maxiter, seed, minimize, opt, already_optimized, node_cost, pbar)

subtree_reconfigure(subtree_size=8, subtree_search='bfs', weight_what='flops', weight_pwr=2, select='max', maxiter='auto', maxiter_auto_cap=1024, seed=None, minimize=None, optimize=None, inplace=False, progbar=False)

Reconfigure subtrees of this tree with locally optimal paths.

Parameters:

• subtree_size (int, optional) – The size of subtree to consider. Cost is exponential in this.

• subtree_search ({'bfs', 'dfs', 'random'}, optional) –

  How to build the subtrees:

  • ’bfs’: breadth-first-search creating balanced subtrees

  • ’dfs’: depth-first-search creating imbalanced subtrees

  • ’random’: random subtree building

• weight_what ({'flops', 'size'}, optional) – When assessing nodes to build and optimize subtrees from whether to score them by the (local) contraction cost, or tensor size.

• weight_pwr (int, optional) – When assessing nodes to build and optimize subtrees from, how to scale their score into a probability: score**(1 / weight_pwr). The larger this is the more explorative the algorithm is when select='random'.

• select ({'descend', 'max', 'min', 'random'}, optional) –

  What order to select node subtrees to optimize:

  • ’descend’: start from the root and then descend into children. In this case the weights and weight_pwr are ignored since this is a deterministic order.

  • ’max’: choose the highest score first

  • ’min’: choose the lowest score first

  • ’random’: choose randomly weighted on score - see weight_pwr.

• maxiter (int, optional) – How many subtree optimizations to perform, the algorithm can terminate before this if all subtrees have been optimized. If ‘auto’, defaults to min(tree.N, maxiter_auto_cap)

• maxiter_auto_cap (int, optional) – The maximum cap to apply to the default value of maxiter when maxiter='auto'.

• seed (int, optional) – A random seed (seeds python system random module).

• minimize ({'flops', 'size'}, optional) – Whether to minimize with respect to contraction flops or size.

• inplace (bool, optional) – Whether to perform the reconfiguration inplace or not.

• progbar (bool, optional) – Whether to show live progress of the reconfiguration.

Return type:

ContractionTree

subtree_reconfigure_

subtree_reconfigure_forest(num_trees=8, num_restarts=10, restart_fraction=0.5, subtree_maxiter=100, subtree_size=10, subtree_search=('random', 'bfs'), subtree_select=('random',), subtree_weight_what=('flops', 'size'), subtree_weight_pwr=(2,), parallel='auto', parallel_maxiter_steps=4, minimize=None, seed=None, progbar=False, inplace=False)

‘Forested’ version of subtree_reconfigure which is more explorative and can be parallelized. It stochastically generates a ‘forest’ reconfigured trees, then only keeps some fraction of these to generate the next forest.

Parameters:

• num_trees (int, optional) – The number of trees to reconfigure at each stage.

• num_restarts (int, optional) – The number of times to halt, prune and then restart the tree reconfigurations.

• restart_fraction (float, optional) – The fraction of trees to keep at each stage and generate the next forest from.

• subtree_maxiter (int, optional) – Number of subtree reconfigurations per step. num_restarts * subtree_maxiter is the max number of total subtree reconfigurations for the final tree produced.

• subtree_size (int, optional) – The size of subtrees to search for and reconfigure.

• subtree_search (tuple[{'random', 'bfs', 'dfs'}], optional) – Tuple of options for the search kwarg of ContractionTree.subtree_reconfigure() to randomly sample.

• subtree_select (tuple[{'random', 'max', 'min'}], optional) – Tuple of options for the select kwarg of ContractionTree.subtree_reconfigure() to randomly sample.

• subtree_weight_what (tuple[{'flops', 'size'}], optional) – Tuple of options for the weight_what kwarg of ContractionTree.subtree_reconfigure() to randomly sample.

• subtree_weight_pwr (tuple[int], optional) – Tuple of options for the weight_pwr kwarg of ContractionTree.subtree_reconfigure() to randomly sample.

• parallel ('auto', False, True, int, or distributed.Client) – Whether to parallelize the search.

• parallel_maxiter_steps (int, optional) – If parallelizing, how many steps to break each reconfiguration into in order to evenly saturate many processes.

• minimize ({'flops', 'size', ..., Objective}, optional) – Whether to minimize the total flops or maximum size of the contraction tree.

• seed (None, int or random.Random, optional) – A random seed to use.

• progbar (bool, optional) – Whether to show live progress.

• inplace (bool, optional) – Whether to perform the subtree reconfiguration inplace.

Return type:

ContractionTree

subtree_reconfigure_forest_

simulated_anneal

simulated_anneal_

parallel_temper

parallel_temper_

slice(target_size=None, target_overhead=None, target_slices=None, temperature=0.01, minimize=None, allow_outer=True, max_repeats=16, reslice=False, seed=None, inplace=False)

Slice this tree (turn some indices into indices which are explicitly summed over rather than being part of contractions). The indices are stored in tree.sliced_inds, and the contraction width updated to take account of the slicing. Calling tree.contract(arrays) moreover which automatically perform the slicing and summation.

Parameters:

• target_size (int, optional) – The target number of entries in the largest tensor of the sliced contraction. The search algorithm will terminate after this is reached.

• target_slices (int, optional) – The target or minimum number of ‘slices’ to consider - individual contractions after slicing indices. The search algorithm will terminate after this is breached. This is on top of the current number of slices.

• target_overhead (float, optional) – The target increase in total number of floating point operations. For example, a value of 2.0 will terminate the search just before the cost of computing all the slices individually breaches twice that of computing the original contraction all at once.

• temperature (float, optional) – How much to randomize the repeated search.

• minimize ({'flops', 'size', ..., Objective}, optional) – Which metric to score the overhead increase against.

• allow_outer (bool, optional) – Whether to allow slicing of outer indices.

• max_repeats (int, optional) – How many times to repeat the search with a slight randomization.

• reslice (bool, optional) – Whether to reslice the tree, i.e. first remove all currently sliced indices and start the search again. Generally any ‘good’ sliced indices will be easily found again.

• seed (None, int or random.Random, optional) – A random seed or generator to use for the search.

• inplace (bool, optional) – Whether the remove the indices from this tree inplace or not.

Return type:

ContractionTree

See also

SliceFinder, ContractionTree.slice_and_reconfigure

slice_

slice_and_reconfigure(target_size, step_size=2, temperature=0.01, minimize=None, allow_outer=True, max_repeats=16, reslice=False, reconf_opts=None, progbar=False, inplace=False)

Interleave slicing (removing indices into an exterior sum) with subtree reconfiguration to minimize the overhead induced by this slicing.

Parameters:

• target_size (int) – Slice the tree until the maximum intermediate size is this or smaller.

• step_size (int, optional) – The minimum size reduction to try and achieve before switching to a round of subtree reconfiguration.

• temperature (float, optional) – The temperature to supply to SliceFinder for searching for indices.

• minimize ({'flops', 'size', ..., Objective}, optional) – The metric to minimize when slicing and reconfiguring subtrees.

• max_repeats (int, optional) – The number of slicing attempts to perform per search.

• progbar (bool, optional) – Whether to show live progress.

• inplace (bool, optional) – Whether to perform the slicing and reconfiguration inplace.

• reconf_opts (None or dict, optional) – Supplied to ContractionTree.subtree_reconfigure() or ContractionTree.subtree_reconfigure_forest(), depending on ‘forested’ key value.

slice_and_reconfigure_

slice_and_reconfigure_forest(target_size, step_size=2, num_trees=8, restart_fraction=0.5, temperature=0.02, max_repeats=32, reslice=False, minimize=None, allow_outer=True, parallel='auto', progbar=False, inplace=False, reconf_opts=None)

‘Forested’ version of ContractionTree.slice_and_reconfigure(). This maintains a ‘forest’ of trees with different slicing and subtree reconfiguration attempts, pruning the worst at each step and generating a new forest from the best.

Parameters:

• target_size (int) – Slice the tree until the maximum intermediate size is this or smaller.

• step_size (int, optional) – The minimum size reduction to try and achieve before switching to a round of subtree reconfiguration.

• num_restarts (int, optional) – The number of times to halt, prune and then restart the tree reconfigurations.

• restart_fraction (float, optional) – The fraction of trees to keep at each stage and generate the next forest from.

• temperature (float, optional) – The temperature at which to randomize the sliced index search.

• max_repeats (int, optional) – The number of slicing attempts to perform per search.

• parallel ('auto', False, True, int, or distributed.Client) – Whether to parallelize the search.

• progbar (bool, optional) – Whether to show live progress.

• inplace (bool, optional) – Whether to perform the slicing and reconfiguration inplace.

• reconf_opts (None or dict, optional) – Supplied to ContractionTree.slice_and_reconfigure().

Return type:

ContractionTree

slice_and_reconfigure_forest_

compressed_reconfigure(minimize=None, order_only=False, max_nodes='auto', max_time=None, local_score=None, exploration_power=0, best_score=None, progbar=False, inplace=False)

Reconfigure this tree according to peak_size_compressed.

Parameters:

• chi (int) – The maximum bond dimension to consider.

• order_only (bool, optional) – Whether to only consider the ordering of the current tree contractions, or all possible contractions, starting with the current.

• 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, new_score, dsize, new_size)
  

  where step is the number of steps so far, new_score is the score of the contraction so far, dsize is the change in memory by the current step, and new_size is the new memory size after contraction.

• 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.

• inplace (bool, optional) – Whether to perform the reconfiguration inplace on this tree.

Return type:

ContractionTree

compressed_reconfigure_

windowed_reconfigure(minimize=None, order_only=False, window_size=20, max_iterations=100, max_window_tries=1000, score_temperature=0.0, queue_temperature=1.0, scorer=None, queue_scorer=None, seed=None, inplace=False, progbar=False, **kwargs)

windowed_reconfigure_

flat_tree(order=None)

Create a nested tuple representation of the contraction tree like:

((0, (1, 2)), ((3, 4), ((5, (6, 7)), (8, 9))))

Such that the contraction will progress like:

((0, (1, 2)), ((3, 4), ((5, (6, 7)), (8, 9))))
((0, 12), (34, ((5, 67), 89)))
(012, (34, (567, 89)))
(012, (34, 56789))
(012, 3456789)
0123456789

Where each integer represents a leaf (i.e. single element node).

get_leaves_ordered()

Return the list of leaves as ordered by the contraction tree.

Return type:

tuple[node_type]

get_path(order=None)

Generate a standard path (with linear recycled ids) from the contraction tree.

Parameters:

order (None, "dfs", or callable, optional) – How to order the contractions within the tree. If a callable is given (which should take a node as its argument), try to contract nodes that minimize this function first.

Returns:

path

Return type:

tuple[tuple[int, int]]

path

get_numpy_path(order=None)

Generate a path compatible with the optimize kwarg of numpy.einsum.

get_ssa_path(order=None)

Generate a single static assignment path from the contraction tree.

Parameters:

order (None, "dfs", or callable, optional) – How to order the contractions within the tree. If a callable is given (which should take a node as its argument), try to contract nodes that minimize this function first.

Returns:

ssa_path

Return type:

tuple[tuple[int, int]]

ssa_path

surface_order(node)

set_surface_order_from_path(ssa_path)

get_path_surface()

path_surface

get_ssa_path_surface()

ssa_path_surface

get_spans()

Get all (which could mean none) potential embeddings of this contraction tree into a spanning tree of the original graph.

Return type:

tuple[dict[frozenset[int], frozenset[int]]]

compute_centralities(combine='mean')

Compute a centrality for every node in this contraction tree.

get_hypergraph(accel=False)

Get a hypergraph representing the uncontracted network (i.e. the leaves).

reset_contraction_indices()

Reset all information regarding a) the explicit contraction indices ordering and b) cached contraction expressions. This should probably be called any time structural changes are made to the tree, e.g. reconfiguration.

sort_contraction_indices(priority='flops', make_output_contig=True, make_contracted_contig=True, reset=True)

Set explicit orders for the contraction indices of this self to optimize for one of two things: contiguity in contracted (‘k’) indices, or contiguity of left and right output (‘m’ and ‘n’) indices.

Parameters:

• priority ({'flops', 'size', 'root', 'leaves'}, optional) – Which order to process the intermediate nodes in. Later nodes re-sort previous nodes so are more likely to keep their ordering. E.g. for ‘flops’ the mostly costly contracton will be process last and thus will be guaranteed to have its indices exactly sorted.

• make_output_contig (bool, optional) – When processing a pairwise contraction, sort the parent contraction indices so that the order of indices is the order they appear from left to right in the two child (input) tensors.

• make_contracted_contig (bool, optional) – When processing a pairwise contraction, sort the child (input) tensor indices so that all contracted indices appear contiguously.

• reset (bool, optional) – Reset all indices to the default order before sorting.

print_contractions(sort=None, show_brackets=True)

Print each pairwise contraction, with colorized indices (if colorama is installed), and other information. The color codes are:

• blue: index appears on left and is kept

• green: index appears on right and is kept

• red: contracted index: appears on both sides and is removed

• pink: batch index: appears on both sides and is kept

Any trivial indices that appear only on one term and not in the output are removed and shown by the preprocessing steps.

Parameters:

• sort ({'flops', 'size'}, optional) – Sort the contractions by either the number of floating point operations or the size of the intermediate tensor. By default the contraction are show in the order they are performed.

• show_brackets (bool, optional) – Whether to show the brackets around contiguous sections of the same type of indices.

get_contractor(order=None, prefer_einsum=False, strip_exponent=False, check_zero=False, implementation=None, autojit=False, progbar=False)

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

Parameters:

• 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 eagerly strip the exponent (in log10) from intermediate tensors to control numerical problems from leaving the range of the datatype. This method then returns the scaled ‘mantissa’ output array and the exponent 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.

Returns:

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

Return type:

callable

contract_core(arrays, order=None, prefer_einsum=False, strip_exponent=False, check_zero=False, backend=None, implementation=None, autojit='auto', progbar=False)

Contract arrays with this tree. The order of the axes and output is assumed to be that of tree.inputs and tree.output, but with sliced indices removed. This functon contracts the core tree and thus if indices have been sliced the arrays supplied need to be sliced as well.

Parameters:

• arrays (sequence of array) – The arrays to contract.

• order (str or callable, optional) – Supplied to ContractionTree.traverse().

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

• backend (str, optional) – What library to use for einsum and transpose, will be automatically inferred from the arrays if not given.

• autojit ("auto" or bool, optional) – Whether to use autoray.autojit to jit compile the expression. If “auto”, then let cotengra choose.

• progbar (bool, optional) – Show progress through the contraction.

slice_key(i, strides=None)

Get the combination of sliced index values for overall slice i.

Parameters:

i (int) – The overall slice index.

Returns:

key – The value each sliced index takes for slice i.

Return type:

dict[str, int]

slice_arrays(arrays, i)

Take arrays and slice the relevant inputs according to tree.sliced_inds and the dynary representation of i.

contract_slice(arrays, i, **kwargs)

Get slices i of arrays and then contract them.

gather_slices(slices, backend=None, progbar=False)

Gather all the output contracted slices into a single full result. If none of the sliced indices appear in the output, then this is a simple sum - otherwise the slices need to be partially summed and partially stacked.

gen_output_chunks(arrays, with_key=False, progbar=False, **contract_opts)

Generate each output chunk of the contraction - i.e. take care of summing internally sliced indices only first. This assumes that the sliced_inds are sorted by whether they appear in the output or not (the default order). Useful for performing some kind of reduction over the final tensor object like fn(x).sum() without constructing the entire thing.

Parameters:

• arrays (sequence of array) – The arrays to contract.

• with_key (bool, optional) – Whether to yield the output index configuration key along with the chunk.

• progbar (bool, optional) – Show progress through the contraction chunks.

Yields:

• chunk (array) – A chunk of the contracted result.

• key (dict[str, int]) – The value each sliced output index takes for this chunk.

contract(arrays, order=None, prefer_einsum=False, strip_exponent=False, check_zero=False, backend=None, implementation=None, autojit='auto', progbar=False)

Contract arrays with this tree. This function takes unsliced arrays and handles the slicing, contractions and gathering. The order of the axes and output is assumed to match that of tree.inputs and tree.output.

Parameters:

• arrays (sequence of array) – The arrays to contract.

• order (str or callable, optional) – Supplied to ContractionTree.traverse().

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

• strip_exponent (bool, optional) – If True, eagerly strip the exponent (in log10) from intermediate tensors to control numerical problems from leaving the range of the datatype. This method then returns the scaled ‘mantissa’ output array and the exponent 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).

• backend (str, optional) – What library to use for tensordot, einsum and transpose, it will be automatically inferred from the input arrays if not given.

• autojit (bool, optional) – Whether to use the ‘autojit’ feature of autoray to compile the contraction expression.

• progbar (bool, optional) – Whether to show a progress bar.

Returns:

• output (array) – The contracted output, it will be scaled if strip_exponent==True.

• exponent (float) – The exponent of the output in base 10, returned only if strip_exponent==True.

See also

contract_core, contract_slice, slice_arrays, gather_slices

contract_mpi(arrays, comm=None, root=None, **kwargs)

Contract the slices of this tree and sum them in parallel - assuming we are already running under MPI.

Parameters:

• arrays (sequence of array) – The input (unsliced arrays)

• comm (None or mpi4py communicator) – Defaults to mpi4py.MPI.COMM_WORLD if not given.

• root (None or int, optional) – If root=None, an Allreduce will be performed such that every process has the resulting tensor, else if an integer e.g. root=0, the result will be exclusively gathered to that process using Reduce, with every other process returning None.

• kwargs – Supplied to contract_slice().

benchmark(dtype='float64', max_time=60, min_reps=3, max_reps=100, warmup=True, **contract_opts)

Benchmark the contraction of this tree.

Parameters:

• dtype ({"float32", "float64", "complex64", "complex128"}) – The datatype to use.

• max_time (float, optional) – The maximum time to spend benchmarking in seconds.

• min_reps (int, optional) – The minimum number of repetitions to perform, regardless of time.

• max_reps (int, optional) – The maximum number of repetitions to perform, regardless of time.

• warmup (bool or int, optional) – Whether to perform a warmup run before the benchmark. If an int, the number of warmup runs to perform.

• contract_opts – Supplied to contract_slice().

Returns:

A dictionary of benchmarking results. The keys are:

• ”time_per_slice”float

  The average time to contract a single slice.

• ”est_time_total”float

  The estimated total time to contract all slices.

• ”est_gigaflops”float

  The estimated gigaflops of the contraction.

Return type:

dict

See also

contract_slice

plot_ring

plot_tent

plot_span

plot_flat

plot_circuit

plot_rubberband

plot_contractions

plot_contractions_alt

plot_hypergraph(**kwargs)

describe(info='normal', join=' ')

Return a string describing the contraction tree.

__repr__()

__str__()

class cotengra.ContractionTreeCompressed(inputs, output, size_dict, track_childless=False, track_flops=False, track_write=False, track_size=False, objective=None, nodeops='auto')

Bases: ContractionTree

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

set_state_from(other)

Set the internal state of this tree to that of other.

classmethod from_path(inputs, output, size_dict, *, path=None, ssa_path=None, autocomplete='auto', 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_default_order()

get_default_objective()

Get the objective function for this tree.

get_default_chi()

get_default_compress_late()

total_flops

Sum the flops contribution from every node in the tree.

Parameters:

dtype ({'float', 'complex', None}, optional) – Scale the answer depending on the assumed data type.

total_write

Sum the total amount of memory that will be created and operated on.

combo_cost

total_cost

max_size

The size of the largest intermediate tensor.

peak_size

Get the peak concurrent size of tensors needed - this depends on the traversal order, i.e. the exact contraction path, not just the contraction tree.

contraction_cost

Get the total number of scalar operations ~ time complexity.

contraction_width

Get log2 of the size of the largest tensor.

total_flops_exact

total_write_exact

combo_cost_exact

total_cost_exact

max_size_exact

peak_size_exact

abstractmethod get_contractor(*_, **__)

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

Parameters:

• 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 eagerly strip the exponent (in log10) from intermediate tensors to control numerical problems from leaving the range of the datatype. This method then returns the scaled ‘mantissa’ output array and the exponent 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.

Returns:

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

Return type:

callable

simulated_anneal(minimize=None, tfinal=0.0001, tstart=0.01, tsteps=50, numiter=50, seed=None, inplace=False, progbar=False, **kwargs)

Perform simulated annealing refinement of this compressed contraction tree.

simulated_anneal_

class cotengra.ContractionTreeMulti(inputs, output, size_dict, sliced_inds, objective, track_cache=False)

Bases: cotengra.core.ContractionTree

Binary tree representing a tensor network contraction.

Parameters:

• inputs (sequence of str) – The list of input tensor’s indices.

• output (str) – The output indices.

• size_dict (dict[str, int]) – The size of each index.

• track_childless (bool, optional) – Whether to dynamically keep track of which nodes are childless. Useful if you are ‘divisively’ building the tree.

• track_flops (bool, optional) – Whether to dynamically keep track of the total number of flops. If False You can still compute this once the tree is complete.

• track_write (bool, optional) – Whether to dynamically keep track of the total number of elements written. If False You can still compute this once the tree is complete.

• track_size (bool, optional) – Whether to dynamically keep track of the largest tensor so far. If False You can still compute this once the tree is complete.

• objective (str or Objective, optional) – An default objective function to use for further optimization and scoring, for example reconfiguring or computing the combo cost. If not supplied the default is to create a flops objective when needed.

children

Mapping of each node to two children.

Type:

dict[node, tuple[node]]

info

Information about the tree nodes. The key is the set of inputs (a set of inputs indices) the node contains. Or in other words, the subgraph of the node. The value is a dictionary to cache information about effective ‘leg’ indices, size, flops of formation etc.

Type:

dict[node, dict]

sliced_inds

_track_cache = False

set_state_from(other)

Set the internal state of this tree to that of other.

_remove_node(node)

Remove node from this tree and update the flops and maximum size if tracking them respectively, as well as input pre-processing.

_update_tracked(node)

get_node_var_inds(node)

Get the set of variable indices that a node depends on.

get_node_is_bright(node)

Get whether a node is ‘bright’, i.e. contains a different set of variable indices to either of its children, if a node is not bright then its children never have to be stored in the cache.

get_node_mult(node)

Get the estimated ‘multiplicity’ of a node, i.e. the number of times it will have to be recomputed for different index configurations.

get_node_cache_mult(node, sliced_ind_ordering)

Get the estimated ‘cache multiplicity’ of a node, i.e. the total number of versions with different index configurations that must be stored simultaneously in the cache.

get_flops(node)

The the estimated total cost of computing a node for all index configurations.

get_cache_contrib(node)

peak_size(log=None)

Get the peak concurrent size of tensors needed - this depends on the traversal order, i.e. the exact contraction path, not just the contraction tree.

reorder_contractions_for_peak_est()

Reorder the contractions to try and reduce the peak memory usage.

reorder_sliced_inds()

exact_multi_stats(configs)

class cotengra.HyperGraph(inputs, output=None, size_dict=None)

Simple hypergraph builder and writer.

Parameters:

• inputs (sequence of list[str] or dict[int, list[str]]) – The nodes. If given as a dict, the keys will be taken as the node enumeration rather than range(len(inputs)).

• output (str, optional) – Output indices.

• size_dict (dict[str, int], optional) – Size of each index.

nodes

Mapping of node to the list of edges incident to it.

Type:

dict[int, list[str]]

edges

Mapping of hyper edges to list of nodes it is incident to.

Type:

dict[str, list[int]]

num_nodes

The number of nodes.

Type:

int

num_edges

The number of hyper-edges.

Type:

int

__slots__ = ('inputs', 'output', 'size_dict', 'nodes', 'edges', 'node_counter')

inputs

output = []

size_dict

edges

node_counter

copy()

Copy this HyperGraph.

classmethod from_edges(edges, output=(), size_dict=())

get_num_nodes()

property num_nodes

get_num_edges()

property num_edges

__len__()

edges_size(es)

Get the combined, i.e. product, size of all edges in es.

bond_size(i, j)

Get the combined, i.e. product, size of edges shared by nodes i and j.

node_size(i)

Get the size of the term represented by node i.

neighborhood_size(nodes)

Get the size of nodes in the immediate neighborhood of nodes.

contract_pair_cost(i, j)

Get the cost of contracting nodes i and j - the product of the dimensions of the indices involved.

neighborhood_compress_cost(chi, nodes)

total_node_size()

Get the total size of all nodes.

output_nodes()

Get the nodes with output indices.

neighbors(i)

Get the neighbors of node i.

neighbor_edges(i)

Get the edges incident to all neighbors of node i, (including its own edges).

has_node(i)

Does this hypergraph have node i?

get_node(i)

Get the edges node i is incident to.

get_edge(e)

Get the nodes edge e is incident to.

has_edge(e)

Does this hypergraph have edge e?

next_node()

Get the next available node identifier.

add_node(inds, node=None)

Add a node with inds, and optional identifier node. The identifier will be generated if not given and returned.

remove_node(i)

Remove node i from this hypergraph.

remove_edge(e)

Remove edge e from this hypergraph.

contract(i, j, node=None)

Combine node i and node j.

compress(chi, edges=None)

‘Compress’ multiedges, combining their size up to a maximum of chi.

compute_contracted_inds(nodes)

Generate the output indices if one were to contract nodes.

candidate_contraction_size(i, j, chi=None)

Get the size of the node created if i and j were contracted, optionally including the effect of first compressing bonds to size chi.

all_shortest_distances(nodes=None)

all_shortest_distances_condensed(nodes=None)

simple_distance(region, p=2)

Compute a simple distance metric from nodes in region to all others. Unlike graph distance, relative connectedness is taken into account.

simple_closeness(p=0.75, mu=0.5)

Compute a rough hypergraph ‘closeness’.

Parameters:

• p (float, optional) – Once any node has had H.num_nodes**p visitors terminate. Set greater than 1.0 for no limit (slower).

• mu (float, optional) – Let the visitor score decay with this power. The higher this is, the more local connectivity is favored.

Returns:

scores – The simple hypergraph closenesses - higher being more central.

Return type:

dict[int, float]

simple_centrality(r=None, smoothness=2, **closeness_opts)

A simple algorithm for large hypergraph centrality. First we find a rough closeness centrality, then relax / smooth this by nodes iteratively radiating their centrality to their neighbors.

Parameters:

• r (None or int, optional) – Number of iterations. Defaults to max(10, int(self.num_nodes**0.5)).

• smoothness (float, optional) – The smoothness. In conjunction with a high value of r this will create a smooth gradient from one of the hypergraph to the other.

• closeness_opts – Supplied to HyperGraph.simple_closeness as the starting point.

Return type:

dict[int, float]

compute_loops(start=None, max_loop_length=None)

Generate all loops up to a certain length in this hypergraph.

Parameters:

• start (sequence of int, optional) – Only generate loops including these nodes, defaults to all.

• max_loop_length (None or int, optional) – The maximum loop length to search for. If None, then this is set automatically by the length of the first loop found.

Yields:

loop (tuple[int]) – A set of nodes that form a loop.

get_laplacian()

Get the graph Laplacian.

get_resistance_distances()

Get the resistance distance between all nodes of the raw graph.

resistance_centrality(rescale=True)

Compute the centrality in terms of the total resistance distance to all other nodes.

to_networkx(as_tree_leaves=None)

Convert to a networkx Graph, with hyperedges represented as nodes.

Parameters:

as_tree_leaves (callable, optional) – If supplied this function is used to convert the nodes to ‘tree leaf’ form, i.e. map node i to single element subgraphs, to match the nodes in a ContractionTree.

compute_weights(weight_edges='const', weight_nodes='const')

plot

__repr__()

cotengra.get_hypergraph(inputs, output=None, size_dict=None, accel=False)

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

class cotengra.HyperCompressedOptimizer(chi=None, methods=('greedy-compressed', 'greedy-span', 'kahypar-agglom'), minimize='peak-compressed', simulated_annealing_opts='auto', reconf_opts='auto', **kwargs)

Bases: HyperOptimizer

A compressed contraction path optimizer that samples a series of ordered contraction trees while optimizing the hyper parameters used to generate them.

Parameters:

• chi (None or int, optional) – The maximum bond dimension to compress to. If None then use the square of the largest existing dimension. If minimize is specified as a full scoring function, this is ignored.

• methods (None or sequence[str] or str, optional) – Which method(s) to use from list_hyper_functions().

• minimize (str, Objective or callable, optional) – How to score each trial, used to train the optimizer and rank the results. If a custom callable, it should take a trial dict as its argument and return a single float.

• max_repeats (int, optional) – The maximum number of trial contraction trees to generate. Default: 128.

• max_time (None or float, optional) – The maximum amount of time to run for. Use None for no limit. You can also set an estimated execution ‘rate’ here like 'rate:1e9' that will terminate the search when the estimated FLOPs of the best contraction found divided by the rate is greater than the time spent searching, allowing quick termination on easy contractions.

• parallel ('auto', False, True, int, or distributed.Client) – Whether to parallelize the search.

• slicing_opts (dict, optional) – If supplied, once a trial contraction path is found, try slicing with the given options, and then update the flops and size of the trial with the sliced versions.

• slicing_reconf_opts (dict, optional) – If supplied, once a trial contraction path is found, try slicing interleaved with subtree reconfiguation with the given options, and then update the flops and size of the trial with the sliced and reconfigured versions.

• reconf_opts (dict, optional) – If supplied, once a trial contraction path is found, try subtree reconfiguation with the given options, and then update the flops and size of the trial with the reconfigured versions.

• optlib ({'cmaes', 'optuna', 'nevergrad', 'skopt', ...}, optional) – Which optimizer to sample and train with.

• space (dict, optional) – The hyper space to search, see get_hyper_space for the default.

• score_compression (float, optional) – Raise scores to this power in order to compress or accentuate the differences. The lower this is, the more the selector will sample from various optimizers rather than quickly specializing.

• max_training_steps (int, optional) – The maximum number of trials to train the optimizer with. Setting this can be helpful when the optimizer itself becomes costly to train (e.g. for Gaussian Processes).

• progbar (bool, optional) – Show live progress of the best contraction found so far.

• optlib_opts – Supplied to the hyper-optimizer library initialization.

compressed = True

multicontraction = False

class cotengra.HyperMultiOptimizer(methods=None, minimize='flops', max_repeats=128, max_time=None, parallel='auto', simulated_annealing_opts='auto', slicing_opts='auto', slicing_reconf_opts='auto', reconf_opts='auto', optlib=None, optlib_opts=None, space=None, score_compression=0.75, on_trial_error='warn', max_training_steps=None, constants=None, progbar=False, **kwargs)

Bases: HyperOptimizer

A path optimizer that samples a series of contraction trees while optimizing the hyper parameters used to generate them. The drivers specified in methods are used to generate the trial contraction trees according to certain hyper-parameters, and the results are scored according to minimize and fed back to optlib to suggest new parameters.

If any of simulated_annealing_opts, slicing_opts, slicing_reconf_opts, or reconf_opts are supplied, then once a trial tree is generated, it will be modified by the corresponding options (and in the order above) and the flops and size of the trial will be updated to the modified version, before the score is reported.

Parameters:

• methods (None or sequence[str] or str, optional) – Which method(s) to use from list_hyper_functions().

• minimize (str, Objective or callable, optional) – How to score each trial, used to train the optimizer and rank the results. If a custom callable, it should take a trial dict as its argument and return a single float. It is also supplied by default to any relevant refinement stages, such as subtree reconfiguration.

• max_repeats (int, optional) – The maximum number of trial contraction trees to generate. Default: 128.

• max_time (None or float, optional) – The maximum amount of time to run for. Use None for no limit. You can also set an estimated execution ‘rate’ here like 'rate:1e9' that will terminate the search when the estimated FLOPs of the best contraction found divided by the rate is greater than the time spent searching, allowing quick termination on easy contractions.

• parallel ('auto', False, True, int, or distributed.Client) – Whether to parallelize the search.

• simulated_annealing_opts (dict, optional) –

  If supplied, once a trial contraction path is found, refine it using simulated annealing with the given options, and then update the flops and size of the trial with the refined version. Notable options and defaults:

  • tsteps=50: number of temperature steps,

  • target_size: simulteneously slice the tree to this size,

  • tfinal=0.05: final temperature,

  • tstart=2: initial temperature.

  See ContractionTree.simulated_anneal() for full details.

• slicing_opts (dict, optional) –

  If supplied, once a trial contraction path is found, try slicing with the given options, and then update the flops and size of the trial with the sliced versions. Notable options:

  • target_size: slice until reaching this size,

  • target_slices: slice into this many slices.

  See ContractionTree.slice() for full details.

• slicing_reconf_opts (dict, optional) –

  If supplied, once a trial contraction path is found, try slicing interleaved with subtree reconfiguation with the given options, and then update the flops and size of the trial with the sliced and reconfigured versions.

  • target_size: slice until reaching this size,

  • reconf_opts: options passed to the subtree reconfiguration stage, see below.

  See ContractionTree.slice_and_reconfigure() for full details.

• reconf_opts (dict, optional) –

  If supplied, once a trial contraction path is found, try subtree reconfiguation with the given options, and then update the flops and size of the trial with the reconfigured versions. Notable options and defaults are:

  • subtree_size=6: size of subtree to optimally reconfigure,

  • maxiter=”auto”: maximum number of subtree reconfigurations, by default scales with size of contraction (up to 1024),

  • select=’max’: which subtrees to prioritize for reconfiguration.

  See ContractionTree.subtree_reconfigure() for full details.

• optlib ({'optuna', 'cmaes', 'nevergrad', 'sses', 'sbplx', ...}, optional) – Which optimizer to sample and train with.

• optlib_opts – Supplied to the hyper-optimizer library initialization.

• space (dict, optional) – The hyper space to search, see get_hyper_space for the default.

• score_compression (float, optional) – Raise scores to this power in order to compress or accentuate the differences. The lower this is, the more the selector will sample from various optimizers rather than quickly specializing.

• on_trial_error ({'warn', 'raise', 'ignore'}, optional) – What to do if a trial fails. If 'warn' (default), a warning will be printed and the trial will be given a score of inf. If 'raise' the error will be raised. If 'ignore' the trial will be given a score of inf silently.

• max_training_steps (int, optional) – The maximum number of trials to train the optimizer with. Setting this can be helpful when the optimizer itself becomes costly to train (e.g. for Gaussian Processes).

• constants (dict[dict], optional) – A dict mapping method name to a dict of constant parameters to pass to the trial function for that method. Any parameters specified here will override those in the search space.

• progbar (bool, optional) – Show live progress of the best contraction found so far.

• kwargs – Extra options to pass to the optimizer library initialization, on top of those in optlib_opts (which take precedence).

compressed = False

multicontraction = True

class cotengra.HyperOptimizer(methods=None, minimize='flops', max_repeats=128, max_time=None, parallel='auto', simulated_annealing_opts='auto', slicing_opts='auto', slicing_reconf_opts='auto', reconf_opts='auto', optlib=None, optlib_opts=None, space=None, score_compression=0.75, on_trial_error='warn', max_training_steps=None, constants=None, progbar=False, **kwargs)

Bases: cotengra.oe.PathOptimizer

A path optimizer that samples a series of contraction trees while optimizing the hyper parameters used to generate them. The drivers specified in methods are used to generate the trial contraction trees according to certain hyper-parameters, and the results are scored according to minimize and fed back to optlib to suggest new parameters.

If any of simulated_annealing_opts, slicing_opts, slicing_reconf_opts, or reconf_opts are supplied, then once a trial tree is generated, it will be modified by the corresponding options (and in the order above) and the flops and size of the trial will be updated to the modified version, before the score is reported.

Parameters:

• methods (None or sequence[str] or str, optional) – Which method(s) to use from list_hyper_functions().

• minimize (str, Objective or callable, optional) – How to score each trial, used to train the optimizer and rank the results. If a custom callable, it should take a trial dict as its argument and return a single float. It is also supplied by default to any relevant refinement stages, such as subtree reconfiguration.

• max_repeats (int, optional) – The maximum number of trial contraction trees to generate. Default: 128.

• max_time (None or float, optional) – The maximum amount of time to run for. Use None for no limit. You can also set an estimated execution ‘rate’ here like 'rate:1e9' that will terminate the search when the estimated FLOPs of the best contraction found divided by the rate is greater than the time spent searching, allowing quick termination on easy contractions.

• parallel ('auto', False, True, int, or distributed.Client) – Whether to parallelize the search.

• simulated_annealing_opts (dict, optional) –

  If supplied, once a trial contraction path is found, refine it using simulated annealing with the given options, and then update the flops and size of the trial with the refined version. Notable options and defaults:

  • tsteps=50: number of temperature steps,

  • target_size: simulteneously slice the tree to this size,

  • tfinal=0.05: final temperature,

  • tstart=2: initial temperature.

  See ContractionTree.simulated_anneal() for full details.

• slicing_opts (dict, optional) –

  If supplied, once a trial contraction path is found, try slicing with the given options, and then update the flops and size of the trial with the sliced versions. Notable options:

  • target_size: slice until reaching this size,

  • target_slices: slice into this many slices.

  See ContractionTree.slice() for full details.

• slicing_reconf_opts (dict, optional) –

  If supplied, once a trial contraction path is found, try slicing interleaved with subtree reconfiguation with the given options, and then update the flops and size of the trial with the sliced and reconfigured versions.

  • target_size: slice until reaching this size,

  • reconf_opts: options passed to the subtree reconfiguration stage, see below.

  See ContractionTree.slice_and_reconfigure() for full details.

• reconf_opts (dict, optional) –

  If supplied, once a trial contraction path is found, try subtree reconfiguation with the given options, and then update the flops and size of the trial with the reconfigured versions. Notable options and defaults are:

  • subtree_size=6: size of subtree to optimally reconfigure,

  • maxiter=”auto”: maximum number of subtree reconfigurations, by default scales with size of contraction (up to 1024),

  • select=’max’: which subtrees to prioritize for reconfiguration.

  See ContractionTree.subtree_reconfigure() for full details.

• optlib ({'optuna', 'cmaes', 'nevergrad', 'sses', 'sbplx', ...}, optional) – Which optimizer to sample and train with.

• optlib_opts – Supplied to the hyper-optimizer library initialization.

• space (dict, optional) – The hyper space to search, see get_hyper_space for the default.

• score_compression (float, optional) – Raise scores to this power in order to compress or accentuate the differences. The lower this is, the more the selector will sample from various optimizers rather than quickly specializing.

• on_trial_error ({'warn', 'raise', 'ignore'}, optional) – What to do if a trial fails. If 'warn' (default), a warning will be printed and the trial will be given a score of inf. If 'raise' the error will be raised. If 'ignore' the trial will be given a score of inf silently.

• max_training_steps (int, optional) – The maximum number of trials to train the optimizer with. Setting this can be helpful when the optimizer itself becomes costly to train (e.g. for Gaussian Processes).

• constants (dict[dict], optional) – A dict mapping method name to a dict of constant parameters to pass to the trial function for that method. Any parameters specified here will override those in the search space.

• progbar (bool, optional) – Show live progress of the best contraction found so far.

• kwargs – Extra options to pass to the optimizer library initialization, on top of those in optlib_opts (which take precedence).

compressed = False

multicontraction = False

max_repeats = 128

_repeats_start = 0

max_time = None

property parallel

method_choices = []

param_choices = []

scores = []

times = []

costs_flops = []

costs_write = []

costs_size = []

simulated_annealing_opts = None

slicing_opts = None

reconf_opts = None

slicing_reconf_opts = None

property minimize

score_compression = 0.75

on_trial_error = 'warn'

best_score

max_training_steps = None

best

trials_since_best = 0

_optimizer

progbar = False

property tree

property path

setup(inputs, output, size_dict)

_maybe_cancel_futures()

_maybe_report_result(setting, trial)

_gen_results(repeats, trial_fn, trial_args)

_get_and_report_next_future()

_gen_results_parallel(repeats, trial_fn, trial_args)

_search(inputs, output, size_dict)

search(inputs, output, size_dict)

Run this optimizer and return the ContractionTree for the best path it finds.

get_tree()

Return the ContractionTree for the best path found.

__call__(inputs, output, size_dict, memory_limit=None)

opt_einsum interface, returns direct path.

get_trials(sort=None)

print_trials(sort=None)

to_df()

Create a single pandas.DataFrame with all trials and scores.

to_dfs_parametrized()

Create a pandas.DataFrame for each method, with all parameters and scores for each trial.

plot_trials

plot_trials_alt

plot_scatter

plot_scatter_alt

plot_parameters_parallel

class cotengra.ReusableHyperCompressedOptimizer(chi=None, methods=('greedy-compressed', 'greedy-span', 'kahypar-agglom'), minimize='peak-compressed', **kwargs)

Bases: ReusableHyperOptimizer

Like HyperCompressedOptimizer but it will re-instantiate the optimizer whenever a new contraction is detected, and also cache the paths found.

Parameters:

• chi (None or int, optional) – The maximum bond dimension to compress to. If None then use the square of the largest existing dimension. If minimize is specified as a full scoring function, this is ignored.

• directory (None, True, or str, optional) – If specified use this directory as a persistent cache. If True auto generate a directory in the current working directory based on the options which are most likely to affect the path (see ReusableHyperOptimizer._get_path_relevant_opts).

• overwrite (bool, optional) – If True, the optimizer will always run, overwriting old results in the cache. This can be used to update paths without deleting the whole cache.

• hash_method ({'a', 'b', ...}, optional) – The method used to hash the contraction tree. The default, 'a', is faster hashwise but doesn’t recognize when indices are permuted.

• cache_only (bool, optional) – If True, the optimizer will only use the cache, and will raise KeyError if a contraction is not found.

• opt_kwargs – Supplied to HyperCompressedOptimizer.

_get_suboptimizer()

_deconstruct_tree(opt, tree)

_reconstruct_tree(inputs, output, size_dict, con)

class cotengra.ReusableHyperOptimizer(*, directory=None, overwrite=False, hash_method='a', cache_only=False, directory_split='auto', **opt_kwargs)

Bases: cotengra.reusable.ReusableOptimizer

Like HyperOptimizer but it will re-instantiate the optimizer whenever a new contraction is detected, and also cache the paths (and sliced indices) found.

Parameters:

• directory (None, True, or str, optional) – If specified use this directory as a persistent cache. If True auto generate a directory in the current working directory based on the options which are most likely to affect the path (see ReusableHyperOptimizer._get_path_relevant_opts).

• overwrite (bool or 'improved', optional) – If True, the optimizer will always run, overwriting old results in the cache. This can be used to update paths without deleting the whole cache. If 'improved' then only overwrite if the new path is better.

• hash_method ({'a', 'b', ...}, optional) – The method used to hash the contraction tree. The default, 'a', is faster hashwise but doesn’t recognize when indices are permuted.

• cache_only (bool, optional) – If True, the optimizer will only use the cache, and will raise KeyError if a contraction is not found.

• directory_split ("auto" or bool, optional) – If specified, the hash will be split into two parts, the first part will be used as a subdirectory, and the second part will be used as the filename. This is useful for avoiding a very large flat diretory. If “auto” it will check the current cache if any and guess from that.

• opt_kwargs – Supplied to HyperOptimizer.

_get_path_relevant_opts()

Get a frozenset of the options that are most likely to affect the path. These are the options that we use when the directory name is not manually specified.

_get_suboptimizer()

_deconstruct_tree(opt, tree)

_reconstruct_tree(inputs, output, size_dict, con)

cotengra.get_hyper_space()

cotengra.list_hyper_functions()

Return a list of currently registered hyper contraction finders.

cotengra.array_contract(arrays, inputs, output=None, optimize='auto', strip_exponent=False, cache_expression=True, backend=None, **kwargs)

Perform the tensor contraction specified by inputs, output and size_dict, using strategy given by optimize. By default the path finding and expression building is cached, so that if the a matching contraction is performed multiple times the overhead is negated.

Parameters:

• arrays (Sequence[array_like]) – The arrays to contract.

• inputs (Sequence[Sequence[Hashable]]) – The inputs terms.

• output (Sequence[Hashable]) – The output term.

• optimize (str, path_like, PathOptimizer, or ContractionTree) –

  The optimization strategy to use. This can be:

  • A string preset, e.g. 'auto', 'greedy', 'optimal'.

  • A PathOptimizer instance.

  • An explicit path, e.g. [(0, 1), (2, 3), ...].

  • An explicit ContractionTree instance.

  If the optimizer provides sliced indices they will be used.

• strip_exponent (bool, optional) – If True, eagerly strip the exponent (in log10) from intermediate tensors to control numerical problems from leaving the range of the datatype. This method then returns the scaled ‘mantissa’ output array and the exponent separately.

• cache_expression (bool, optional) – If True, cache the expression used to contract the arrays. This negates the overhead of pathfinding and building the expression when a contraction is performed multiple times. Only for hashable optimize objects.

• backend (str, optional) – If given, the explicit backend to use for the contraction, by default the backend is dispatched automatically.

• kwargs – Passed to array_contract_expression().

• (array_like (array_like or) – The result of the contraction. If strip_exponent is True, the result is a tuple of the output array mantissae and the exponent (in log10).

• scalar) – The result of the contraction. If strip_exponent is True, the result is a tuple of the output array mantissae and the exponent (in log10).

See also

array_contract_expression, array_contract_tree, einsum

cotengra.array_contract_expression(inputs, output=None, size_dict=None, shapes=None, optimize='auto', constants=None, canonicalize=True, cache=True, **kwargs)

Get an callable ‘expression’ that will contract tensors with indices and shapes described by inputs and size_dict to output. The optimize kwarg can be a path, optimizer or also a contraction tree. In the latter case sliced indices for example will be used if present. The same is true if optimize is an optimizer that can directly produce ContractionTree instances (i.e. has a .search() method).

Parameters:

• inputs (Sequence[Sequence[Hashable]]) – The inputs terms.

• output (Sequence[Hashable]) – The output term.

• size_dict (dict[Hashable, int]) – The size of each index.

• optimize (str, path_like, PathOptimizer, or ContractionTree) –

  The optimization strategy to use. This can be:

  • A string preset, e.g. 'auto', 'greedy', 'optimal'.

  • A PathOptimizer instance.

  • An explicit path, e.g. [(0, 1), (2, 3), ...].

  • An explicit ContractionTree instance.

  If the optimizer provides sliced indices they will be used.

• constants (dict[int, array_like], optional) – A mapping of constant input positions to constant arrays. If given, the final expression will take only the remaining non-constant tensors as inputs. Note this is a different format to the constants kwarg of einsum_expression() since it also provides the constant arrays.

• 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.

• via (tuple[callable, callable], optional) – If given, the first function will be applied to the input arrays and the second to the output array. For example, moving the tensors from CPU to GPU and back.

• sort_contraction_indices (bool, optional) – If True, call tree.sort_contraction_indices() before constructing the contraction function.

• cache (bool, optional) – If True, cache the contraction expression. This negates the overhead of pathfinding and building the expression when a contraction is performed multiple times. Only for hashable optimize objects.

Returns:

expr – A callable, signature expr(*arrays) that will contract arrays with shapes shapes.

Return type:

callable

See also

einsum_expression, array_contract, array_contract_tree

cotengra.array_contract_path(inputs, output=None, size_dict=None, shapes=None, optimize='auto', canonicalize=True, cache=True)

Find only the contraction path for the specific contraction, with fast dispatch of optimize, which can be a preset, path, tree, cotengra optimizer or opt_einsum optimizer. The raw path is a more compact representation of the core tree structure but contains less information on its own, for example sliced indices are not included.

Parameters:

• inputs (Sequence[Sequence[Hashable]]) – The inputs terms.

• output (Sequence[Hashable], optional) – The output term.

• size_dict (dict[Hashable, int], optional) – The size of each index, if given, shapes is ignored.

• shapes (Sequence[tuple[int]], optional) – The shape of each input array. Needed if size_dict not supplied.

• optimize (str, path_like, PathOptimizer, or ContractionTree) –

  The optimization strategy to use. This can be:

  • A string preset, e.g. 'auto', 'greedy', 'optimal'.

  • A PathOptimizer instance.

  • An explicit path, e.g. [(0, 1), (2, 3), ...].

  • An explicit ContractionTree instance.

• canonicalize (bool, optional) – If True, canonicalize the inputs and output so that the indices are relabelled 'a', 'b', 'c', ..., etc. in the order they appear.

• cache (bool, optional) – If True, cache the path for the contraction, so that if the same pathfinding is performed multiple times the overhead is negated. Only for hashable optimize objects.

Returns:

path – The contraction path, whose interpretation is thus: the input tensors are assumed to be stored in a list, i.e. indexed by range(N). Each contraction in the path is a set of indices, the tensors at these locations should be popped from the list and then the result of the contraction appended.

Return type:

tuple[tuple[int]]

cotengra.array_contract_tree(inputs, output=None, size_dict=None, shapes=None, optimize='auto', canonicalize=True, sort_contraction_indices=False)

Get the ContractionTree for the tensor contraction specified by inputs, output and size_dict, with optimization strategy given by optimize. The tree can be used to inspect and also perform the contraction.

Parameters:

• inputs (Sequence[Sequence[Hashable]]) – The inputs terms.

• output (Sequence[Hashable], optional) – The output term.

• size_dict (dict[Hashable, int], optional) – The size of each index, if given, shapes is ignored.

• shapes (Sequence[tuple[int]], optional) – The shape of each input array. Needed if size_dict not supplied.

• optimize (str, path_like, PathOptimizer, or ContractionTree) –

  The optimization strategy to use. This can be:

  • A string preset, e.g. 'auto', 'greedy', 'optimal'.

  • A PathOptimizer instance.

  • An explicit path, e.g. [(0, 1), (2, 3), ...].

  • An explicit ContractionTree instance.

• canonicalize (bool, optional) – If True, canonicalize the inputs and output so that the indices are relabelled 'a', 'b', 'c', ..., etc. in the order they appear.

• sort_contraction_indices (bool, optional) – If True, call tree.sort_contraction_indices().

Return type:

ContractionTree

See also

array_contract, array_contract_expression, einsum_tree

cotengra.einsum(*args, optimize='auto', strip_exponent=False, cache_expression=True, backend=None, **kwargs)

Perform an einsum contraction, using cotengra, using strategy given by optimize. By default the path finding and expression building is cached, so that if a matching contraction is performed multiple times the overhead is negated.

Parameters:

• eq (str) – The equation to use for contraction, for example 'ab,bc->ac'.

• arrays (Sequence[array_like]) – The arrays to contract.

• optimize (str, path_like, PathOptimizer, or ContractionTree) –

  The optimization strategy to use. This can be:

  • A string preset, e.g. 'auto', 'greedy', 'optimal'.

  • A PathOptimizer instance.

  • An explicit path, e.g. [(0, 1), (2, 3), ...].

  • An explicit ContractionTree instance.

  If the optimizer provides sliced indices they will be used.

• strip_exponent (bool, optional) – If True, eagerly strip the exponent (in log10) from intermediate tensors to control numerical problems from leaving the range of the datatype. This method then returns the scaled ‘mantissa’ output array and the exponent separately.

• cache_expression (bool, optional) – If True, cache the expression used to contract the arrays. This negates the overhead of pathfinding and building the expression when a contraction is performed multiple times. Only for hashable optimize objects.

• backend (str, optional) – If given, the explicit backend to use for the contraction, by default the backend is dispatched automatically.

• kwargs – Passed to array_contract_expression().

Returns:

The result of the contraction. If strip_exponent is True, the result is a tuple of the output array mantissae and the exponent (in log10).

Return type:

array_like or (array_like, scalar)

See also

einsum_expression, einsum_tree, array_contract

cotengra.einsum_expression(*args, optimize='auto', constants=None, cache=True, **kwargs)

Get an callable ‘expression’ that will contract tensors with shapes shapes according to equation eq. The optimize kwarg can be a path, optimizer or also a contraction tree. In the latter case sliced indices for example will be used if present. The same is true if optimize is an optimizer that can directly produce ContractionTree instances (i.e. has a .search() method).

Parameters:

• eq (str) – The equation to use for contraction, for example 'ab,bc->ac'. The output will be automatically computed if not supplied, but Ellipses (’…’) are not supported.

• shapes (Sequence[tuple[int]]) – The shapes of the tensors to contract, or the constant tensor itself if marked as constant in constants.

• optimize (str, path_like, PathOptimizer, or ContractionTree) –

  The optimization strategy to use. This can be:

  • A string preset, e.g. 'auto', 'greedy', 'optimal'.

  • A PathOptimizer instance.

  • An explicit path, e.g. [(0, 1), (2, 3), ...].

  • An explicit ContractionTree instance.

  If the optimizer provides sliced indices they will be used.

• constants (Sequence of int, optional) – The indices of tensors to treat as constant, the final expression will take the remaining non-constant tensors as inputs. Note this is a different format to the constants kwarg of array_contract_expression() since the actual constant arrays are inserted into shapes.

• 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.

• via (tuple[callable, callable], optional) – If given, the first function will be applied to the input arrays and the second to the output array. For example, moving the tensors from CPU to GPU and back.

• sort_contraction_indices (bool, optional) – If True, call tree.sort_contraction_indices() before constructing the contraction function.

• cache (bool, optional) – If True, cache the contraction expression. This negates the overhead of pathfinding and building the expression when a contraction is performed multiple times. Only for hashable optimize objects.

Returns:

expr – A callable, signature expr(*arrays) that will contract arrays with shapes matching shapes.

Return type:

callable

See also

einsum, einsum_tree, array_contract_expression

cotengra.einsum_tree(*args, optimize='auto', canonicalize=False, sort_contraction_indices=False)

Get the ContractionTree for the einsum equation eq and optimization strategy optimize. The tree can be used to inspect and also perform the contraction.

Parameters:

• eq (str) – The equation to use for contraction, for example 'ab,bc->ac'.

• shapes (Sequence[tuple[int]]) – The shape of each input array.

• optimize (str, path_like, PathOptimizer, or ContractionTree) –

  The optimization strategy to use. This can be:

  • A string preset, e.g. 'auto', 'greedy', 'optimal'.

  • A PathOptimizer instance.

  • An explicit path, e.g. [(0, 1), (2, 3), ...].

  • An explicit ContractionTree instance.

• canonicalize (bool, optional) – If True, canonicalize the inputs and output so that the indices are relabelled 'a', 'b', 'c', ..., etc. in the order they appear.

• sort_contraction_indices (bool, optional) – If True, call tree.sort_contraction_indices().

Return type:

ContractionTree

See also

einsum, einsum_expression, array_contract_tree

cotengra.ncon(arrays, indices, **kwargs)

Perform a contraction specified by the ncon style indices, using cotengra. This is very similar to array_contract, but the indices should be intergers, with negative integers specifying outputs. The output order is determined by the order of the negative integers, like [-1, -2, -3, ...].

See paper https://arxiv.org/abs/1402.0939 and python implementation https://github.com/mhauru/ncon.

The contraction order is found as with other cotengra routines, and does not default to contracting in the order given by the indices.

Parameters:

• arrays (Sequence[array_like]) – The arrays to contract.

• indices (Sequence[Sequence[int]]) – The indices to contract, with negative integers specifying outputs.

• kwargs (dict) – Supplied to array_contract.

Returns:

The result of the contraction. If strip_exponent is True, the result is a tuple of the output array mantissae and the exponent (in log10).

Return type:

array_like or (array_like, scalar)

cotengra.register_preset(preset, optimizer, optimizer_tree=None, register_opt_einsum='auto', compressed=False)

Register a preset optimizer.

Parameters:

• preset (str or Sequence[str]) – The name of the preset, or a sequence of alias names.

• optimizer (callable) – The optimizer function that returns a path.

• optimizer_tree (callable, optional) – The optimizer function that returns a tree.

• register_opt_einsum (bool or str, optional) – If True or 'auto', register the preset with opt_einsum.

• compressed (bool, optional) – If True, the preset presents a compressed contraction optimizer.

class cotengra.GreedyOptimizer(costmod=1.0, temperature=0.0, max_neighbors=DEFAULT_MAX_NEIGHBORS, simplify=True, accel='auto')

Bases: cotengra.oe.PathOptimizer

Class interface to the greedy optimizer which can be instantiated with default options.

__slots__ = ('costmod', 'temperature', 'max_neighbors', 'simplify', '_optimize_fn')

costmod = 1.0

temperature = 0.0

max_neighbors = 16

simplify = True

_optimize_fn

maybe_update_defaults(**kwargs)

ssa_path(inputs, output, size_dict, **kwargs)

search(inputs, output, size_dict, **kwargs)

__call__(inputs, output, size_dict, **kwargs)

class cotengra.OptimalOptimizer(minimize='flops', cost_cap=2, search_outer=False, simplify=True, accel='auto')

Bases: cotengra.oe.PathOptimizer

Class interface to the optimal optimizer which can be instantiated with default options.

__slots__ = ('minimize', 'cost_cap', 'search_outer', 'simplify', '_optimize_fn')

minimize = 'flops'

cost_cap = 2

search_outer = False

simplify = True

_optimize_fn

maybe_update_defaults(**kwargs)

ssa_path(inputs, output, size_dict, **kwargs)

search(inputs, output, size_dict, **kwargs)

__call__(inputs, output, size_dict, **kwargs)

class cotengra.RandomGreedyOptimizer(max_repeats=32, costmod=(0.1, 4.0), temperature=(0.001, 1.0), max_neighbors=DEFAULT_MAX_NEIGHBORS, seed=None, simplify=True, accel='auto', parallel='auto')

Bases: cotengra.oe.PathOptimizer

Lightweight random greedy optimizer, that eschews hyper parameter tuning and contraction tree construction. This is a stateful optimizer that should not be re-used on different contractions.

Parameters:

• max_repeats (int, optional) – The number of random greedy trials to perform.

• costmod ((float, float), optional) –

  When assessing local greedy scores how much to weight the size of the tensors removed compared to the size of the tensor added:

  score = size_ab / costmod - (size_a + size_b) * costmod
  

  It is sampled uniformly from the given range.

• temperature ((float, float), optional) –

  When asessing local greedy scores, how much to randomly perturb the score. This is implemented as:

  score -> sign(score) * log(|score|) - temperature * gumbel()
  

  which implements boltzmann sampling. It is sampled log-uniformly from the given range.

• max_neighbors (int, optional) – When looking for pairs of nodes to contract, skip any index that connects to more than this many nodes. This is useful to avoid combinatorial explosions when dealing with essentially batch indices.

• seed (int, optional) – The seed for the random number generator. Note that deterministic behavior is only guaranteed within the python or rust backend (the accel parameter) and parallel settings.

• simplify (bool, optional) –

  Whether to perform simplifications before optimizing. These are:

  • ignore any indices that appear in all terms

  • combine any repeated indices within a single term

  • reduce any non-output indices that only appear on a single term

  • combine any scalar terms

  • combine any tensors with matching indices (hadamard products)

  Such simpifications may be required in the general case for the proper functioning of the core optimization, but may be skipped if the input indices are already in a simplified form.

• accel (bool or str, optional) – Whether to use the accelerated cotengrust backend. If “auto” the backend is used if available.

• parallel (bool or str, optional) – Whether to use parallel processing. If “auto” the default is to use processes unless the accelerated backend is used, in which case threads are used.

best_ssa_path

The best contraction path found so far.

Type:

list[list[int]]

best_flops

The flops (/ contraction cost / number of multiplications) of the best contraction path found so far.

Type:

float

minimize = 'flops'

max_repeats = 32

max_neighbors = 16

simplify = True

rng

best_ssa_path = None

best_flops

tree = None

_optimize_fn

_pool = None

maybe_update_defaults(**kwargs)

ssa_path(inputs, output, size_dict, **kwargs)

search(inputs, output, size_dict, **kwargs)

__call__(inputs, output, size_dict, **kwargs)

class cotengra.ReusableRandomGreedyOptimizer(*, directory=None, overwrite=False, hash_method='a', cache_only=False, directory_split='auto', **opt_kwargs)

Bases: cotengra.reusable.ReusableOptimizer

A reusable random greedy optimizer that caches path per contraction.

_get_path_relevant_opts()

Get a frozenset of the options that are most likely to affect the path. These are the options that we use when the directory name is not manually specified.

_get_suboptimizer()

_deconstruct_tree(opt, tree)

_reconstruct_tree(inputs, output, size_dict, con)

cotengra.edge_path_to_linear(edge_path, inputs)

Convert a path specified by a sequence of edge indices to a path with recycled linear ids.

Parameters:

• edge_path (sequence[str | int]) – The path specified by a sequence of edge indices.

• inputs (tuple[tuple[str | int]]) – The indices of each input tensor.

Returns:

path – The contraction path in recycled linear id format.

Return type:

tuple[tuple[int]]

cotengra.edge_path_to_ssa(edge_path, inputs)

Convert a path specified by a sequence of edge indices to a path with tuples of single static assignment (SSA) indices.

Parameters:

• edge_path (sequence[str | int]) – The path specified by a sequence of edge indices.

• inputs (tuple[tuple[str | int]]) – The indices of each input tensor.

Returns:

path – The contraction path in static single assignment (SSA) form.

Return type:

tuple[tuple[int]]

cotengra.linear_to_ssa(path, N=None)

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

>>> linear_to_ssa([(0, 3), (1, 2), (0, 1)])
[(0, 3), (2, 4), (1, 5)]

cotengra.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)]

class cotengra.FlowCutterOptimizer(max_time=10, seed=None, executable='flow_cutter_pace17')

Bases: cotengra.oe.PathOptimizer

max_time = 10

rng

executable = 'flow_cutter_pace17'

run_flowcutter(file, max_time=None)

compute_edge_path(lg)

build_tree(inputs, output, size_dict, memory_limit=None)

__call__(inputs, output, size_dict, memory_limit=None)

cotengra.optimize_flowcutter(inputs, output, size_dict, memory_limit=None, max_time=10, seed=None)

class cotengra.QuickBBOptimizer(max_time=10, executable='quickbb_64', seed=None)

Bases: cotengra.oe.PathOptimizer

max_time = 10

executable = 'quickbb_64'

run_quickbb(fname, outfile, statfile, max_time=None)

build_tree(inputs, output, size_dict)

__call__(inputs, output, size_dict, memory_limit=None)

cotengra.optimize_quickbb(inputs, output, size_dict, memory_limit=None, max_time=60, seed=None)

class cotengra.RandomOptimizer(seed=None)

Bases: cotengra.oe.PathOptimizer

A fully random pathfinder, that randomly selects pairs of tensors to contract (even if they are not connected). This is useful for testing purposes, and as a baseline for comparison with other pathfinders.

Parameters:

seed (None, int or np.random.Generator, optional) – Random seed. If None, a random seed is selected. Default is None.

rng

__call__(inputs, outputs, size_dict)

search(inputs, outputs, size_dict)

cotengra.plot_contractions(tree, order=None, color_size=(0.6, 0.4, 0.7), color_cost=(0.3, 0.7, 0.5), figsize=(8, 3))

cotengra.plot_contractions_alt(tree, x='peak-size', y='flops', color='stage', size='scaling', width=400, height=400, point_opacity=0.8, color_scheme='lightmulti', x_scale='log', y_scale='log', color_scale='log', size_scale='linear')

cotengra.plot_scatter(self, x='size', y='flops', ylim=None, cumulative_time=False, plot_best=False, figsize=(5, 5))

cotengra.plot_scatter_alt(self, x='size', y='flops', color='run:Q', color_scheme='purplebluegreen', shape='method:N', width=400, height=400)

Plot the trials total flops vs max size.

cotengra.plot_slicings(slice_finder, color_scheme='RdYlBu_r', relative_flops=False, figsize=(6, 3), point_opacity=0.8)

cotengra.plot_slicings_alt(slice_finder, color_scheme='redyellowblue', relative_flops=False)

cotengra.plot_tree(tree, layout='ring', hypergraph_layout='auto', hypergraph_layout_opts=None, k=0.01, iterations=500, span=None, order=None, order_y_pow=1.0, edge_scale=1.0, node_scale=1.0, highlight=(), edge_color='size', edge_colormap='GnBu', edge_max_width=None, node_colormap='YlOrRd', node_color='flops', node_max_size=None, figsize=(5, 5), raw_edge_color=None, raw_edge_alpha=None, tree_root_height=True, tree_alpha=0.8, colorbars=True, plot_raw_graph=True, plot_leaf_labels=False, ax=None)

Plot a contraction tree using matplotlib.

cotengra.plot_tree_ring(tree, **kwargs)

cotengra.plot_tree_span(tree, **kwargs)

cotengra.plot_tree_tent(tree, **kwargs)

cotengra.plot_trials(self, *, x='trial', y='score', figsize=(8, 3), cumulative_time=True, plot_best=True, **kwargs)

cotengra.plot_trials_alt(self, y=None, width=800, height=300)

Plot the trials interactively using altair.

class cotengra.AutoHQOptimizer(**kwargs)

Bases: AutoOptimizer

An optimizer that automatically chooses between optimal and hyper-optimization, designed for everyday use on harder contractions or those that will be repeated many times, and thus warrant a more extensive search.

class cotengra.AutoOptimizer(optimal_cutoff=250, minimize='combo', cache=True, **hyperoptimizer_kwargs)

Bases: cotengra.oe.PathOptimizer

An optimizer that automatically chooses between optimal and hyper-optimization, designed for everyday use.

minimize = 'combo'

optimal_cutoff = 250

_optimize_optimal_fn

kwargs

_hyperoptimizers_by_thread

_get_optimizer_hyper_threadsafe()

search(inputs, output, size_dict, **kwargs)

__call__(inputs, output, size_dict, **kwargs)

cotengra.auto_hq_optimize(inputs, output, size_dict, **kwargs)

cotengra.auto_optimize(inputs, output, size_dict, **kwargs)

cotengra.greedy_optimize

cotengra.optimal_optimize

cotengra.optimal_outer_optimize

cotengra.hash_contraction(inputs, output, size_dict, method='a')

Compute a hash for a particular contraction geometry.

class cotengra.SliceFinder(tree_or_info, target_size=None, target_overhead=None, target_slices=None, temperature=0.01, minimize='flops', allow_outer=True, seed=None)

An object to help find the best indices to slice over in order to reduce the memory footprint of a contraction as much as possible whilst introducing as little extra overhead. It searches for and stores ContractionCosts.

Parameters:

• tree_or_info (ContractionTree or opt_einsum.PathInfo) – Object describing the target full contraction to slice.

• target_size (int, optional) – The target number of entries in the largest tensor of the sliced contraction. The search algorithm will terminate after this is reached.

• target_slices (int, optional) – The target or minimum number of ‘slices’ to consider - individual contractions after slicing indices. The search algorithm will terminate after this is breached. This is on top of the current number of slices.

• target_overhead (float, optional) – The target increase in total number of floating point operations. For example, a value of 2.0 will terminate the search just before the cost of computing all the slices individually breaches twice that of computing the original contraction all at once.

• temperature (float, optional) – When sampling combinations of indices, how far to randomly stray from what looks like the best (local) choice.

info

costs

temperature = 0.01

rng

target_size = None

target_overhead = None

target_slices = None

minimize

_maybe_default(attr, value)

best(k=None, target_size=None, target_overhead=None, target_slices=None)

Return the best contraction slicing, subject to target filters.

trial(target_size=None, target_overhead=None, target_slices=None, temperature=None)

A single slicing attempt, greedily select indices from the popular pool, subject to the score function, terminating when any of the target criteria are met.

search(max_repeats=16, temperature=None, target_size=None, target_overhead=None, target_slices=None)

Repeat trial several times and return the best found so far.

plot_slicings

plot_slicings_alt

cotengra.get_symbol(i)

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) #> ‘京’

cotengra.get_symbol_map(inputs)

Get a mapping of arbitrary hashable ‘indices’ to single unicode symbols, matching the canonicalization of the expression.

cotengra.UniformOptimizer

Does no gaussian process tuning by default, just randomly samples - requires no optimization library.

cotengra.contract_expression

Alias for cotengra.einsum_expression().

cotengra.contract

Alias for cotengra.einsum().

cotengra.hyper_optimize(inputs, output, size_dict, memory_limit=None, get='path', **opts)