cotengra.pathfinders.path_basic

Basic optimization routines.

Attributes

DEFAULT_MAX_NEIGHBORS

Classes

ContractionProcessor	A helper class for combining bottom up simplifications, greedy, and
EnsureInputsOutputAreSequence
GreedyOptimizer	Class interface to the greedy 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.
OptimalOptimizer	Class interface to the optimal optimizer which can be instantiated with

Functions

is_simplifiable(legs, appearances)	Check if legs contains any diag (repeated) or reduced (appears
compute_simplified(legs, appearances)	Compute the diag and reduced legs of a term. This function assumes that
compute_contracted(ilegs, jlegs, appearances)	Compute the contracted legs of two terms.
compute_size(legs, sizes)	Compute the size of a term.
compute_flops(ilegs, jlegs, sizes)	Compute the flops cost of contracting two terms.
compute_con_cost_flops(temp_legs, appearances, sizes, ...)	Compute the total flops cost of a contraction given by temporary legs,
compute_con_cost_max(temp_legs, appearances, sizes, ...)	Compute the max flops cost of a contraction given by temporary legs,
compute_con_cost_size(temp_legs, appearances, sizes, ...)	Compute the max size of a contraction given by temporary legs, also
compute_con_cost_write(temp_legs, appearances, sizes, ...)	Compute the total write cost of a contraction given by temporary legs,
compute_con_cost_combo(temp_legs, appearances, sizes, ...)	Compute the combined total flops and write cost of a contraction given
compute_con_cost_limit(temp_legs, appearances, sizes, ...)	Compute the combined total flops and write cost of a contraction given
parse_minimize_for_optimal(minimize)	Given a string, parse it into a function that computes the cost of a
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
edge_path_to_ssa(edge_path, inputs)	Convert a path specified by a sequence of edge indices to a path with
edge_path_to_linear(edge_path, inputs)	Convert a path specified by a sequence of edge indices to a path with
is_ssa_path(path, nterms)	Check if an explicitly given path is in 'static single assignment' form.
optimize_simplify(inputs, output, size_dict[, use_ssa])	Find the (likely only partial) contraction path corresponding to
optimize_greedy(inputs, output, size_dict[, costmod, ...])	Find a contraction path using a greedy algorithm.
optimize_random_greedy_track_flops(inputs, output, ...)	Perform a batch of random greedy optimizations, simulteneously tracking
optimize_optimal(inputs, output, size_dict[, ...])	Find the optimal contraction path using a dynamic programming
get_optimize_greedy([accel])
get_optimize_random_greedy_track_flops([accel])
get_optimize_optimal([accel])

Module Contents

cotengra.pathfinders.path_basic.DEFAULT_MAX_NEIGHBORS = 16

cotengra.pathfinders.path_basic.is_simplifiable(legs, appearances)

Check if legs contains any diag (repeated) or reduced (appears nowhere else) indices.

cotengra.pathfinders.path_basic.compute_simplified(legs, appearances)

Compute the diag and reduced legs of a term. This function assumes that the legs are already sorted. It handles the case where a index is both diag and reduced (i.e. traced).

cotengra.pathfinders.path_basic.compute_contracted(ilegs, jlegs, appearances)

Compute the contracted legs of two terms.

cotengra.pathfinders.path_basic.compute_size(legs, sizes)

Compute the size of a term.

cotengra.pathfinders.path_basic.compute_flops(ilegs, jlegs, sizes)

Compute the flops cost of contracting two terms.

cotengra.pathfinders.path_basic.compute_con_cost_flops(temp_legs, appearances, sizes, iscore, jscore)

Compute the total flops cost of a contraction given by temporary legs, also removing any contracted indices from the temporary legs.

cotengra.pathfinders.path_basic.compute_con_cost_max(temp_legs, appearances, sizes, iscore, jscore)

Compute the max flops cost of a contraction given by temporary legs, also removing any contracted indices from the temporary legs.

cotengra.pathfinders.path_basic.compute_con_cost_size(temp_legs, appearances, sizes, iscore, jscore)

Compute the max size of a contraction given by temporary legs, also removing any contracted indices from the temporary legs.

cotengra.pathfinders.path_basic.compute_con_cost_write(temp_legs, appearances, sizes, iscore, jscore)

Compute the total write cost of a contraction given by temporary legs, also removing any contracted indices from the temporary legs.

cotengra.pathfinders.path_basic.compute_con_cost_combo(temp_legs, appearances, sizes, iscore, jscore, factor)

Compute the combined total flops and write cost of a contraction given by temporary legs, also removing any contracted indices from the temporary legs. The combined cost is given by:

cost = flops + factor * size

cotengra.pathfinders.path_basic.compute_con_cost_limit(temp_legs, appearances, sizes, iscore, jscore, factor)

Compute the combined total flops and write cost of a contraction given by temporary legs, also removing any contracted indices from the temporary legs. The combined cost is given by:

cost = max(flops, factor * size)

I.e. assuming one or another to be the limiting factor.

cotengra.pathfinders.path_basic.parse_minimize_for_optimal(minimize)

Given a string, parse it into a function that computes the cost of a contraction. The string can be one of the following:

• “flops”: compute_con_cost_flops

• “max”: compute_con_cost_max

• “size”: compute_con_cost_size

• “write”: compute_con_cost_write

• “combo”: compute_con_cost_combo

• “combo-{factor}”: compute_con_cost_combo with specified factor

• “limit”: compute_con_cost_limit

• “limit-{factor}”: compute_con_cost_limit with specified factor

This function is cached for speed.

class cotengra.pathfinders.path_basic.ContractionProcessor(inputs, output, size_dict, track_flops=False, flops_limit=float('inf'))

A helper class for combining bottom up simplifications, greedy, and optimal contraction path optimization.

__slots__ = ('nodes', 'edges', 'indmap', 'appearances', 'sizes', 'ssa', 'ssa_path', 'track_flops', 'flops',...

nodes

edges

indmap

appearances = []

sizes = []

ssa

ssa_path = []

track_flops = False

flops = 0

flops_limit

copy()

neighbors(i)

Get all neighbors of node i.

Parameters:

i (int) – The node index to get neighbors of.

Yields:

j (int) – The neighboring node index.

neighbors_limit(i, max_neighbors)

Get all neighbors of node i, ignoring any index that connects to more than max_neighbors nodes. This is for greedy optimization where it is useful to avoid combinatorial explosions caused by essentially batch indices.

Parameters:

• i (int) – The node index to get neighbors of.

• max_neighbors (int) – If non-zero, skip any index that connects to more than this many nodes. This is useful to avoid combinatorial explosions when dealing with essentially batch indices.

Yields:

j (int) – The neighboring node index.

print_current_terms()

remove_ix(ix)

Drop the index ix, simply removing it from all nodes and the edgemap.

pop_node(i)

Remove node i from the graph, updating the edgemap and returning the legs of the node.

add_node(legs)

Add a new node to the graph, updating the edgemap and returning the node index of the new node.

check()

Check that the current graph is valid, useful for debugging.

contract_nodes(i, j, new_legs=None)

Contract the nodes i and j, adding a new node to the graph and returning its index.

simplify_batch()

Find any indices that appear in all terms and remove them, since they simply add an constant factor to the cost of the contraction, but create a fully connected graph if left.

simplify_single_terms()

Take any diags, reductions and traces of single terms.

simplify_scalars()

Remove all scalars, contracting them into the smallest remaining node, if there is one.

simplify_hadamard()

simplify()

subgraphs()

optimize_greedy(costmod=1.0, temperature=0.0, max_neighbors=DEFAULT_MAX_NEIGHBORS, seed=None)

optimize_optimal_connected(where, minimize='flops', cost_cap=2, search_outer=False)

optimize_optimal(minimize='flops', cost_cap=2, search_outer=False)

optimize_remaining_by_size()

This function simply contracts remaining terms in order of size, and is meant to handle the disconnected terms left after greedy or optimal optimization.

cotengra.pathfinders.path_basic.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.pathfinders.path_basic.ssa_to_linear(ssa_path, N=None)

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

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

cotengra.pathfinders.path_basic.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.pathfinders.path_basic.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.pathfinders.path_basic.is_ssa_path(path, nterms)

Check if an explicitly given path is in ‘static single assignment’ form.

cotengra.pathfinders.path_basic.optimize_simplify(inputs, output, size_dict, use_ssa=False)

Find the (likely only partial) contraction path corresponding to simplifications only. Those simplifiactions 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)

Parameters:

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

• output (tuple[str]) – The indices of the output tensor.

• size_dict (dict[str, int]) – A dictionary mapping indices to their dimension.

• use_ssa (bool, optional) – Whether to return the contraction path in ‘SSA’ format (i.e. as if each intermediate is appended to the list of inputs, without removals).

Returns:

path – The contraction path, given as a sequence of pairs of node indices.

Return type:

list[list[int]]

cotengra.pathfinders.path_basic.optimize_greedy(inputs, output, size_dict, costmod=1.0, temperature=0.0, max_neighbors=DEFAULT_MAX_NEIGHBORS, simplify=True, use_ssa=False)

Find a contraction path using a greedy algorithm.

Parameters:

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

• output (tuple[str]) – The indices of the output tensor.

• size_dict (dict[str, int]) – A dictionary mapping indices to their dimension.

• costmod (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
  

  This can be a useful hyper-parameter to tune.

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

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

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

• use_ssa (bool, optional) – Whether to return the contraction path in ‘single static assignment’ (SSA) format (i.e. as if each intermediate is appended to the list of inputs, without removals). This can be quicker and easier to work with than the ‘linear recycled’ format that numpy and opt_einsum use.

Returns:

path – The contraction path, given as a sequence of pairs of node indices.

Return type:

list[list[int]]

cotengra.pathfinders.path_basic.optimize_random_greedy_track_flops(inputs, output, size_dict, ntrials=1, costmod=(0.1, 4.0), temperature=(0.001, 1.0), seed=None, max_neighbors=DEFAULT_MAX_NEIGHBORS, simplify=True, use_ssa=False)

Perform a batch of random greedy optimizations, simulteneously tracking the best contraction path in terms of flops, so as to avoid constructing a separate contraction tree.

Parameters:

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

• output (tuple[str]) – The indices of the output tensor.

• size_dict (dict[str, int]) – A dictionary mapping indices to their dimension.

• ntrials (int, optional) – The number of random greedy trials to perform. The default is 1.

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

• seed (int, optional) – The seed for the random number generator.

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

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

• use_ssa (bool, optional) – Whether to return the contraction path in ‘single static assignment’ (SSA) format (i.e. as if each intermediate is appended to the list of inputs, without removals). This can be quicker and easier to work with than the ‘linear recycled’ format that numpy and opt_einsum use.

Returns:

• path (list[list[int]]) – The best contraction path, given as a sequence of pairs of node indices.

• flops (float) – The flops (/ contraction cost / number of multiplications), of the best contraction path, given log10.

cotengra.pathfinders.path_basic.optimize_optimal(inputs, output, size_dict, minimize='flops', cost_cap=2, search_outer=False, simplify=True, use_ssa=False)

Find the optimal contraction path using a dynamic programming algorithm (by default excluding outer products).

The algorithm is an optimized version of Phys. Rev. E 90, 033315 (2014) (preprint: https://arxiv.org/abs/1304.6112), adapted from the opt_einsum implementation.

Parameters:

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

• output (tuple[str]) – The indices of the output tensor.

• size_dict (dict[str, int]) – A dictionary mapping indices to their dimension.

• minimize (str, optional) –

  The cost function to minimize. The options are:

  • ”flops”: minimize with respect to total operation count only (also known as contraction cost)

  • ”size”: minimize with respect to maximum intermediate size only (also known as contraction width)

  • ’max’: minimize the single most expensive contraction, i.e. the asymptotic (in index size) scaling of the contraction

  • ’write’ : minimize the sum of all tensor sizes, i.e. memory written

  • ’combo’ or ‘combo={factor}` : minimize the sum of FLOPS + factor * WRITE, with a default factor of 64.

  • ’limit’ or ‘limit={factor}` : minimize the sum of MAX(FLOPS, alpha * WRITE) for each individual contraction, with a default factor of 64.

  ’combo’ is generally a good default in term of practical hardware performance, where both memory bandwidth and compute are limited.

• cost_cap (float, optional) – The maximum cost of a contraction to initially consider. This acts like a sieve and is doubled at each iteration until the optimal path can be found, but supplying an accurate guess can speed up the algorithm.

• search_outer (bool, optional) – Whether to allow outer products in the contraction path. The default is False. Especially when considering write costs, the fastest path is very unlikely to include outer products.

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

• use_ssa (bool, optional) – Whether to return the contraction path in ‘single static assignment’ (SSA) format (i.e. as if each intermediate is appended to the list of inputs, without removals). This can be quicker and easier to work with than the ‘linear recycled’ format that numpy and opt_einsum use.

Returns:

path – The contraction path, given as a sequence of pairs of node indices.

Return type:

list[list[int]]

class cotengra.pathfinders.path_basic.EnsureInputsOutputAreSequence(f)

f

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

cotengra.pathfinders.path_basic.get_optimize_greedy(accel='auto')

cotengra.pathfinders.path_basic.get_optimize_random_greedy_track_flops(accel='auto')

class cotengra.pathfinders.path_basic.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.pathfinders.path_basic.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.pathfinders.path_basic.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.pathfinders.path_basic.get_optimize_optimal(accel='auto')

class cotengra.pathfinders.path_basic.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)