cotengra.core_multi

Module Contents

Classes

ContractionTreeMulti

Binary tree representing a tensor network contraction.

class cotengra.core_multi.ContractionTreeMulti(inputs, output, size_dict, sliced_inds, objective, track_cache=False)[source]

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]

set_state_from(other)[source]

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

_remove_node(node)[source]

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)[source]
get_node_var_inds(node)[source]

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

get_node_is_bright(node)[source]

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)[source]

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)[source]

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)[source]

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

get_cache_contrib(node)[source]
peak_size(log=None)[source]

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()[source]

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

reorder_sliced_inds()[source]
exact_multi_stats(configs)[source]