cotengra.slicer#

Module Contents#

Classes#

ContractionCosts

A simplified struct for tracking the contraction costs of a path only.

SliceFinder

An object to help find the best indices to slice over in order to reduce

SlicedContractor

A contraction where certain indices are explicitly summed over,

Functions#

create_size_dict(inputs, arrays)

Attributes#

class cotengra.slicer.ContractionCosts(contractions, size_dict, nslices=1, original_flops=None)#

A simplified struct for tracking the contraction costs of a path only.

Parameters:
  • contractions (sequence of Contraction) – The set of individual operations that make up a full contraction path.

  • size_dict (dict[str, int]) – The sizes of the indices involved.

  • nslices (int, optional) – For keeping track of the ‘multiplicity’ of this set of contractions if part of a sliced contration where indices have been removed.

__slots__ = ['size_dict', 'contractions', 'nslices', 'original_flops', '_flops', '_sizes',...#
_set_state_from(other)#

Copy all internal structure from another ContractionCosts.

copy()#

Get a copy of this ContractionCosts.

property size#
property flops#
property total_flops#
property overhead#
classmethod from_info(info, **kwargs)#

Generate a set of contraction costs from a PathInfo object.

classmethod from_contraction_tree(contraction_tree, **kwargs)#

Generate a set of contraction costs from a ContractionTree object.

score(ix, minimize='flops', factor='')#

The ‘score’ for assessing whether to remove ix.

remove(ix, inplace=False)#
__repr__()#

Return repr(self).

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

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 PathInfo) – Object describing the target full contraction to slice, generated for example from a call to contract_path().

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

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

plot_slicings#
plot_slicings_alt#
_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.

SlicedContractor(arrays, target_size=None, target_overhead=None, target_slices=None, **kwargs)#

Generate a sliced contraction using the best indices found by this SliceFinder and by default the original contraction path as well.

cotengra.slicer.create_size_dict(inputs, arrays)#
class cotengra.slicer.SlicedContractor(eq, arrays, sliced, optimize='auto', size_dict=None)#

A contraction where certain indices are explicitly summed over, corresponding to taking different ‘slices’ of the input arrays, each of which can be contracted independently with hopefully a lower memory requirement. The recommended way of instantiating this is from a directly from SliceFinder which already.

Parameters:
  • eq (str) – The overall contraction to perform.

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

  • sliced (sequence of str) – Which indices in eq to slice over.

  • optimize (str or path or PathOptimizer, optional) – How to optimize the sliced contraction path - the contraction with sliced indices removed. If these sliced indices were found automatically is it generally best to supply the full path they were found with respect to rather than trying to re-optimize the path.

  • size_dict (dict[str, int], optional) – If already known, the sizes of each index.

classmethod from_info(info, arrays, sliced, optimize=None, **kwargs)#

Creat a SlicedContractor directly from a PathInfo object.

property individual_flops#

FLOP cost of a single contraction slice.

property total_flops#

FLOP cost of performing all sliced contractions.

property max_size#

The largest size tensor produced in an individual contraction.

get_sliced_arrays(i)#

Generate the tuple of array inputs corresponding to slice i.

contract_slice(i, **kwargs)#

Contraction of just slice i.

gather_slices(slices)#

Gather all the output contracted slices into the single full result.

contract_all(**kwargs)#

Contract (and sum) all slices at once.

get_dask_chunked(**kwargs)#
get_mars_chunked(**kwargs)#