From cae397ab298b8d06fad12c90023510f97a1928e2 Mon Sep 17 00:00:00 2001 From: Christoph Staudt Date: Mon, 29 Apr 2024 12:41:39 +0200 Subject: [PATCH 1/8] Add basic first version of random kahypar+ --- cotengra/__init__.py | 1 + cotengra/pathfinders/path_basic.py | 41 +- cotengra/pathfinders/path_kahypar_plus.py | 651 ++++++++++++++++++++++ 3 files changed, 690 insertions(+), 3 deletions(-) create mode 100644 cotengra/pathfinders/path_kahypar_plus.py diff --git a/cotengra/__init__.py b/cotengra/__init__.py index bf88f82..5d60f14 100644 --- a/cotengra/__init__.py +++ b/cotengra/__init__.py @@ -65,6 +65,7 @@ path_igraph, path_kahypar, path_labels, + path_kahypar_plus, ) from .pathfinders.path_basic import ( GreedyOptimizer, diff --git a/cotengra/pathfinders/path_basic.py b/cotengra/pathfinders/path_basic.py index c9c1d15..5a139c1 100644 --- a/cotengra/pathfinders/path_basic.py +++ b/cotengra/pathfinders/path_basic.py @@ -5,6 +5,7 @@ import bisect import functools import itertools +import random from ..oe import PathOptimizer from ..utils import get_rng, GumbelBatchedGenerator @@ -1271,7 +1272,40 @@ def __init__( self.rng = get_rng(seed) self.best_ssa_path = None self.best_flops = float("inf") - self._optimize_fn = get_optimize_random_greedy_track_flops(accel) + sub_optimize_fn = get_optimize_random_greedy_track_flops(accel) + + # TODO: Remove This was just hacked in for testing + def mini_optimizer( + inputs, + output, + size_dict, + ntrials=1, + seed=None, + simplify=True, + use_ssa=False, + ): + best_cost = math.inf + best_path = None + for i in range(ntrials): + temp = 2 ** random.uniform(math.log2(0.001), math.log2(1)) + path, cost = sub_optimize_fn( # type: ignore + inputs, + output, + size_dict, + ntrials=1, + seed=seed, + costmod=random.uniform(0, 50), + temperature=temp, + simplify=simplify, + use_ssa=use_ssa, + ) + + if cost < best_cost: + best_cost = cost + best_path = path + return best_path, best_cost + + self._optimize_fn = mini_optimizer if (parallel == "auto") and ( self._optimize_fn is not optimize_random_greedy_track_flops @@ -1287,8 +1321,9 @@ def __init__( def maybe_update_defaults(self, **kwargs): # allow overriding of defaults opts = { - "costmod": self.costmod, - "temperature": self.temperature, + # TODO: Remove This was just hacked in for testing + # "costmod": self.costmod, + # "temperature": self.temperature, "simplify": self.simplify, } opts.update(kwargs) diff --git a/cotengra/pathfinders/path_kahypar_plus.py b/cotengra/pathfinders/path_kahypar_plus.py new file mode 100644 index 0000000..c3d884b --- /dev/null +++ b/cotengra/pathfinders/path_kahypar_plus.py @@ -0,0 +1,651 @@ +import copy +import math +import random +import uuid +from collections import defaultdict +from dataclasses import dataclass, field +from functools import total_ordering +from queue import PriorityQueue +from typing import ( + Callable, + Hashable, + Optional, + Sequence, + Union, + List, + Tuple, + Dict, + Set, + FrozenSet, +) + +from cotengra.pathfinders.path_basic import ( + OptimalOptimizer, + RandomGreedyOptimizer, +) +from .path_kahypar import kahypar_subgraph_find_membership +from ..core import ContractionTree +from ..hyperoptimizers.hyper import register_hyper_function + + +Inputs = List[List[Hashable]] +Shape = Tuple[int, ...] +Shapes = List[Shape] +Output = List[Hashable] +SizeDict = Dict[Hashable, int] +Path = List[Tuple[int, ...]] +GreedyOptimizer = Callable[["TensorNetwork"], Path] + + +@dataclass +class BasicInputNode: + indices: List[Hashable] + shape: Shape + + +@dataclass +class OriginalInputNode(BasicInputNode): + id: int + + def get_id(self): + return str(self.id) + + def __repr__(self) -> str: + return f"Original Input({self.indices}, {self.shape})" + + +@dataclass +class SubNetworkInputNode(BasicInputNode): + sub_network: "SubTensorNetwork" + + def get_id(self): + return f"sn-{self.sub_network.name}" + + def __repr__(self) -> str: + return f"Sub network Input({self.sub_network.output_indices}, {self.sub_network.get_output_shape()})" + + +InputNode = Union[OriginalInputNode, SubNetworkInputNode] +InputNodes = List["InputNode"] + + +@dataclass +class IntermediateContractNode: + all_indices: Set[Hashable] # Union all indices of children + scale: int + indices: Output + children: List["ContractTreeNode"] + uuid: str = field(default_factory=lambda: str(uuid.uuid4())) + + def get_id(self): + return self.uuid + + def __repr__(self) -> str: + return f"IntermediateContractNode({self.uuid}), children: {[child.get_id() for child in self.children]}" + + +ContractTreeNode = Union[ + OriginalInputNode, SubNetworkInputNode, IntermediateContractNode +] + +ContractTree = List[ContractTreeNode] + + +def safe_log2(x): + if x < 1: + return 0 + return math.log2(x) + + +def get_contract_tree_and_cost_from_path( + tn: "TensorNetwork", ssa_path +) -> Tuple[ContractTree, int]: + contract_tree: ContractTree = [] + histogram = defaultdict(lambda: 0) + for input in tn.get_all_input_nodes(): + contract_tree.append(input) + for edge in input.indices: + histogram[edge] += 1 + + for index in tn.output_indices: + histogram[index] += 1 + + # If there is only one input thats the whole tree + if len(contract_tree) == 1: + return contract_tree, 1 + total_cost = 0 + + for pair in ssa_path: + if len(pair) == 1: + left_node: ContractTreeNode = contract_tree[pair[0]] + all_indices = set(left_node.indices) + cost = 1 + remove = set() + for index in left_node.indices: + cost = cost * tn.size_dict[index] + if histogram[index] == 0: + remove.add(index) + total_cost += cost + intermediate = all_indices - remove + for index in intermediate: + histogram[index] += 1 + contract_tree.append( + IntermediateContractNode( + all_indices, + int(safe_log2(cost)), + list(intermediate), + [contract_tree[pair[0]]], + ) + ) + if len(pair) == 2: + left_node: ContractTreeNode = contract_tree[pair[0]] + right_node: ContractTreeNode = contract_tree[pair[1]] + all_indices = set(left_node.indices).union(right_node.indices) + cost = 1 + remove = set() + for index in all_indices: + cost = cost * tn.size_dict[index] + + for index in left_node.indices: + histogram[index] -= 1 + if histogram[index] == 0: + remove.add(index) + for index in right_node.indices: + histogram[index] -= 1 + if histogram[index] == 0: + remove.add(index) + total_cost += cost + intermediate = all_indices - remove + + for index in intermediate: + histogram[index] += 1 + + contract_tree.append( + IntermediateContractNode( + all_indices, + int(safe_log2(cost)), + list(intermediate), + [contract_tree[pair[0]], contract_tree[pair[1]]], + ) + ) + total_cost = 2 * total_cost + + return contract_tree, total_cost + + +def greedy_optimizer(tn: "TensorNetwork") -> Path: + inputs = [input.indices for input in tn.get_all_input_nodes()] + output = tn.output_indices + size_dict = tn.size_dict + + if len(inputs) <= 15: + optimal_opt = OptimalOptimizer() + return optimal_opt.ssa_path(inputs, output, size_dict) + + greedy_opt = RandomGreedyOptimizer(max_repeats=512) + return greedy_opt.ssa_path(inputs, output, size_dict) + + +@dataclass +class SubTensorNetwork: + name: str + key: int + parent_name: str + inputs: InputNodes + indices: FrozenSet[Hashable] + size_dict: SizeDict + cut_indices: FrozenSet[Hashable] + output_indices: Output + + def get_all_input_nodes(self) -> InputNodes: + return self.inputs + + def get_all_networks(self): + return [self] + + def get_output_shape(self): + return tuple([self.size_dict[index] for index in self.output_indices]) + + def find_path(self): + """ + Finds the path for the sub-tensor network. + + Returns: + SubTensorNetworkWithContractTree: The sub-tensor network with the computed contract tree and its cost. + """ + path = greedy_optimizer(self) + + contract_tree, cost = get_contract_tree_and_cost_from_path(self, path) + + tn_with_tree = SubTensorNetworkWithContractTree( + name=self.name, + key=self.key, + parent_name=self.parent_name, + inputs=self.inputs, + indices=self.indices, + size_dict=self.size_dict, + cut_indices=self.cut_indices, + output_indices=self.output_indices, + cost=cost, + contract_tree=contract_tree, + ) + return tn_with_tree + + +@total_ordering +@dataclass +class SubTensorNetworkWithContractTree(SubTensorNetwork): + cost: int + contract_tree: ContractTree + + def get_total_cost(self): + return self.cost + + def get_contract_tree(self): + return self.contract_tree + + def __eq__(self, other): + if not isinstance(other, __class__): + return NotImplemented + return self.cost == other.cost + + def __lt__(self, other): + if not isinstance(other, __class__): + return NotImplemented + # Yes this might seem weird, but we want the one with the highest cost to be the first in the priority queue + return self.cost > other.cost + + def find_path(self): + """ + Since the path is already computed, this function just returns the current object. + + Returns: + The current object. + """ + return self + + +@dataclass +class SuperTensorNetwork(SubTensorNetwork): + parent_name: str + sub_networks: Sequence[SubTensorNetwork] + + def get_all_input_nodes(self) -> InputNodes: + sub_input_nodes = [ + SubNetworkInputNode( + sub_network.output_indices, + sub_network.get_output_shape(), + sub_network, + ) + for sub_network in self.sub_networks + ] + + return sub_input_nodes + self.inputs + + def get_all_networks(self): + return [self] + [sub_network for sub_network in self.sub_networks] + + def find_path(self): + sub_networks_with_path: List[SubTensorNetworkWithContractTree] = [] + for sub_network in self.sub_networks: + sub_networks_with_path.append(sub_network.find_path()) + + self.sub_networks = sub_networks_with_path + + path = greedy_optimizer(self) + contract_tree, cost = get_contract_tree_and_cost_from_path(self, path) + + tn_with_tree = SuperTensorNetworkWithTree( + self.name, + self.key, + self.parent_name, + self.inputs, + self.indices, + self.size_dict, + self.cut_indices, + self.output_indices, + cost, + contract_tree, + sub_networks_with_path, + ) + + return tn_with_tree + + +@dataclass +class SuperTensorNetworkWithTree( + SuperTensorNetwork, + SubTensorNetworkWithContractTree, +): + sub_networks: List[SubTensorNetworkWithContractTree] + + # def find_path(self, greedy_optimizer: GreedyOptimizer): + # """ + # Since the path is already computed, this function just returns the current object. + + # Returns: + # The current object. + # """ + # return self + + def get_total_cost(self): + return ( + sum([sub_network.cost for sub_network in self.sub_networks]) + + self.cost + ) + + def get_parent_tree(self): + super_tree = self.get_contract_tree() + + parent_tree = [] + sub_tree_root: Dict[str, ContractTreeNode] = {} + for node in super_tree: + if ( + isinstance(node, SubNetworkInputNode) + and node.sub_network.parent_name == self.name + ): + assert isinstance( + node.sub_network, SubTensorNetworkWithContractTree + ), "The subnetworks should have a contract tree, when calling get_parent_tree" + sub_tree = None + for sn in self.sub_networks: + if sn.name == node.sub_network.name: + sub_tree = sn.contract_tree + assert ( + sub_tree is not None + ), f"Sub tree {node.sub_network.name} not found in {self.name}" + for sub_node in sub_tree: + parent_tree.append(sub_node) + + sub_tree_root[node.sub_network.name] = sub_tree[-1] + elif isinstance(node, IntermediateContractNode): + for key, child in enumerate(node.children): + if ( + isinstance(child, SubNetworkInputNode) + and child.sub_network.parent_name == self.name + ): + node.children[key] = sub_tree_root[ + child.sub_network.name + ] + + parent_tree.append(node) + else: + parent_tree.append(node) + + return parent_tree + + def update_tree(self, name, tree: ContractTree): + if name == self.name: + self.contract_tree = tree + else: + for sub_network in self.sub_networks: + if sub_network.name == name: + sub_network.contract_tree = tree + return + raise Exception(f"name {name} not found in {self.name}") + + def find_path(self): + return self + + +TensorNetwork = Union[SubTensorNetwork, SuperTensorNetwork] +TensorNetworkWithTree = Union[ + SubTensorNetworkWithContractTree, SuperTensorNetworkWithTree +] + + +def get_sub_networks( + tensor_network_name: str, + input_nodes, + output, + size_dict: SizeDict, + imbalance: float, + weight_nodes: str = "const", +): + num_input_nodes = len(input_nodes) + assert ( + num_input_nodes > 2 + ), f"Not enough input nodes to split, pass at least two input nodes, {input_nodes}" + + inputs = [input.indices for input in input_nodes] + + if len(output) > 0: + inputs.append(output) + + block_ids = kahypar_subgraph_find_membership( + inputs, + set(), + size_dict, + weight_nodes=weight_nodes, + weight_edges="log", + fix_output_nodes=False, + parts=2, + imbalance=imbalance, + compress=0, + mode="recursive", + objective="cut", + quiet=True, + ) + + ## Noramlize block ids + + # Check if all input nodes were assigned to the same block + input_block_ids = block_ids[:num_input_nodes] + min_block_id = min(input_block_ids) + max_block_id = max(input_block_ids) + if min_block_id == max_block_id: + # If there is only one block just distribute them with modulo + block_ids = [i % 2 for i in range(num_input_nodes + 1)] + input_block_ids = block_ids[:num_input_nodes] + else: + if min_block_id != 0 or max_block_id != 1: + block_ids = [0 if id == min_block_id else 1 for id in block_ids] + + assert ( + len(set(input_block_ids)) == 2 + ), f"There should be two blocks, {input_block_ids}, {min_block_id}, {max_block_id}" + + # Group inputs by block id + block_inputs: list[InputNodes] = [[], []] + for block_id, input_node in zip(block_ids, input_nodes): + block_inputs[block_id].append(input_node) + + block_indices = [ + frozenset( + set.union(*[set(input_node.indices) for input_node in block]) + ) + for block in block_inputs + ] + + cut_indices = set() + cut_indices = cut_indices.union( + block_indices[0].intersection(block_indices[1]) + ) + + cut_indices = frozenset(cut_indices.union(output)) + output_indices: list[Output] = [ + cut_indices.intersection(block) for block in block_indices + ] + + output_block_id = 0 + if len(output) > 0: + output_block_id = block_ids[-1] + else: + output_block_id = random.choice([0, 1]) + + sub_networks = [ + SubTensorNetwork( + f"{tensor_network_name}.{key}", + key, + None, + block_inputs[key], + block_indices[key], + size_dict, + cut_indices, + output_indices[key], + ) + for key, _ in enumerate(block_inputs) + ] + + super_sub_network = sub_networks.pop(output_block_id) + super_network = build_super_network( + super_sub_network, + tensor_network_name, + cut_indices, + output, + sub_networks, + ) + + return super_network.find_path() + + +def build_super_network( + super_network: TensorNetwork, + parent_name: str, + cut_indices: FrozenSet[Hashable], + output: Output, + sub_networks: List[SubTensorNetwork], +): + # Set parent for sub_networks + for sub_network in sub_networks: + sub_network.parent_name = super_network.name + + super_network = SuperTensorNetwork( + name=super_network.name, + key=super_network.key, + parent_name=parent_name, + inputs=super_network.inputs, + indices=super_network.indices, + size_dict=super_network.size_dict, + cut_indices=cut_indices, + output_indices=output, + sub_networks=sub_networks, + ) + return super_network + + +def get_remapped_id(id, input_remap: Optional[Dict[str, int]]): + return input_remap[id] if input_remap != None and id in input_remap else id + + +def contract_tree_to_path( + tree: ContractTree, remap: Optional[Dict[str, int]] = None +): + root = tree[-1] + + if isinstance(root, BasicInputNode): + assert ( + len(tree) == 1 + ), "Tree should only contain one node, if root is basic input" + return [(int(get_remapped_id(root.get_id(), remap)),)] + + path = [] + + counter = (len(tree) + 1) // 2 + uuid_to_ssa_id = {} + for node in tree: + if isinstance(node, IntermediateContractNode): + uuid_to_ssa_id[node.get_id()] = counter + counter += 1 + pair = [] + if isinstance(node.children[0], BasicInputNode): + pair.append(get_remapped_id(node.children[0].get_id(), remap)) + else: + pair.append(uuid_to_ssa_id[node.children[0].get_id()]) + + if len(node.children) > 1: + if isinstance(node.children[1], BasicInputNode): + pair.append( + get_remapped_id(node.children[1].get_id(), remap) + ) + else: + pair.append(uuid_to_ssa_id[node.children[1].get_id()]) + path.append(tuple(pair)) + return path + + +def hybrid_hypercut_greedy( + inputs: Inputs, + output: Output, + size_dict: SizeDict, + imbalance, + weight_nodes="const", + cutoff=15, +): + # Problem setup, transform arguments to tanser network with path + shapes = [tuple([size_dict[i] for i in input]) for input in inputs] + inputs = [list(input) for input in inputs] + output = list(output) + indices = frozenset.union(*[frozenset(input) for input in inputs]) + + input_nodes: InputNodes = [ + OriginalInputNode(in_sh[0], in_sh[1], id) + for id, in_sh in enumerate(zip(inputs, shapes)) + ] + + tensor_network_without_path = SubTensorNetwork( + "tn", 0, None, input_nodes, indices, size_dict, frozenset(), output + ) + + tensor_network = tensor_network_without_path.find_path() + + root_name = tensor_network.name + + # Initialize queues + to_partition: PriorityQueue[TensorNetwork] = PriorityQueue() + to_partition.put(tensor_network) + + # Initialize dictoionary for finalized partitions + partitioned: Dict[str, TensorNetworkWithTree] = {} + + while not to_partition.empty(): + next_network = to_partition.get() + if len(next_network.inputs) <= cutoff: + with_path = next_network.find_path() # + partitioned[next_network.name] = with_path + continue + + super_network = get_sub_networks( + next_network.name, + next_network.get_all_input_nodes(), + next_network.output_indices, + next_network.size_dict, + imbalance=imbalance, + weight_nodes=weight_nodes, + ) + + partitioned[super_network.parent_name] = copy.deepcopy(super_network) + + to_partition.put(super_network) + + for sub_network in super_network.sub_networks: + to_partition.put(sub_network) + + def merge(network_name) -> ContractTree: + partitioned_network = partitioned[network_name] + # Check if we reached a leave + if partitioned_network.name == network_name or not ( + isinstance(partitioned_network, SuperTensorNetworkWithTree) + ): + with_tree = partitioned_network.find_path() + return with_tree.contract_tree + for sub_network in partitioned_network.get_all_networks(): + merged_sub_tree = merge(sub_network.name) + partitioned_network.update_tree(sub_network.name, merged_sub_tree) + + parent_tree = partitioned_network.get_parent_tree() + + return parent_tree + + parent_tree = merge(root_name) + path = contract_tree_to_path(parent_tree) + path = [tuple([int(i) for i in pair]) for pair in path] + return ContractionTree.from_path(inputs, output, size_dict, ssa_path=path) + + +hyper_space = { + "imbalance": {"type": "FLOAT", "min": 0.01, "max": 0.2}, + "weight_nodes": {"type": "STRING", "options": ["const", "log"]}, + "cutoff": {"type": "INT", "min": 40, "max": 200}, +} +register_hyper_function("random-kahypar+", hybrid_hypercut_greedy, hyper_space) From 9ab15ac5a0e69b878fdafa440c37781527da778c Mon Sep 17 00:00:00 2001 From: Christoph Staudt Date: Mon, 29 Apr 2024 12:48:43 +0200 Subject: [PATCH 2/8] Add link to original repo --- cotengra/pathfinders/path_kahypar_plus.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cotengra/pathfinders/path_kahypar_plus.py b/cotengra/pathfinders/path_kahypar_plus.py index c3d884b..0e07598 100644 --- a/cotengra/pathfinders/path_kahypar_plus.py +++ b/cotengra/pathfinders/path_kahypar_plus.py @@ -1,3 +1,7 @@ +# +# Based on https://github.com/ti2-group/hybrid_contraction_tree_optimizer/ +# + import copy import math import random From 008564bcb5275b9c62a0ada39b57cefd6db916e8 Mon Sep 17 00:00:00 2001 From: Christoph Staudt Date: Sun, 5 May 2024 00:15:16 +0200 Subject: [PATCH 3/8] Reduced iterative version with contract trees --- cotengra/pathfinders/path_kahypar_plus.py | 436 +++++----------------- 1 file changed, 89 insertions(+), 347 deletions(-) diff --git a/cotengra/pathfinders/path_kahypar_plus.py b/cotengra/pathfinders/path_kahypar_plus.py index 0e07598..e81d3e4 100644 --- a/cotengra/pathfinders/path_kahypar_plus.py +++ b/cotengra/pathfinders/path_kahypar_plus.py @@ -2,25 +2,20 @@ # Based on https://github.com/ti2-group/hybrid_contraction_tree_optimizer/ # -import copy import math import random import uuid from collections import defaultdict from dataclasses import dataclass, field -from functools import total_ordering -from queue import PriorityQueue from typing import ( Callable, Hashable, Optional, - Sequence, Union, List, Tuple, Dict, Set, - FrozenSet, ) from cotengra.pathfinders.path_basic import ( @@ -33,8 +28,7 @@ Inputs = List[List[Hashable]] -Shape = Tuple[int, ...] -Shapes = List[Shape] + Output = List[Hashable] SizeDict = Dict[Hashable, int] Path = List[Tuple[int, ...]] @@ -44,7 +38,6 @@ @dataclass class BasicInputNode: indices: List[Hashable] - shape: Shape @dataclass @@ -55,18 +48,18 @@ def get_id(self): return str(self.id) def __repr__(self) -> str: - return f"Original Input({self.indices}, {self.shape})" + return f"Original Input({self.indices})" @dataclass class SubNetworkInputNode(BasicInputNode): - sub_network: "SubTensorNetwork" + sub_network: "TensorNetwork" def get_id(self): return f"sn-{self.sub_network.name}" def __repr__(self) -> str: - return f"Sub network Input({self.sub_network.output_indices}, {self.sub_network.get_output_shape()})" + return f"Sub network Input({self.sub_network.output_indices}," InputNode = Union[OriginalInputNode, SubNetworkInputNode] @@ -85,7 +78,10 @@ def get_id(self): return self.uuid def __repr__(self) -> str: - return f"IntermediateContractNode({self.uuid}), children: {[child.get_id() for child in self.children]}" + return ( + f"IntermediateContractNode({self.uuid}), " + + f"children: {[child.get_id() for child in self.children]}" + ) ContractTreeNode = Union[ @@ -106,7 +102,7 @@ def get_contract_tree_and_cost_from_path( ) -> Tuple[ContractTree, int]: contract_tree: ContractTree = [] histogram = defaultdict(lambda: 0) - for input in tn.get_all_input_nodes(): + for input in tn.inputs: contract_tree.append(input) for edge in input.indices: histogram[edge] += 1 @@ -178,7 +174,7 @@ def get_contract_tree_and_cost_from_path( def greedy_optimizer(tn: "TensorNetwork") -> Path: - inputs = [input.indices for input in tn.get_all_input_nodes()] + inputs = [input.indices for input in tn.inputs] output = tn.output_indices size_dict = tn.size_dict @@ -191,225 +187,44 @@ def greedy_optimizer(tn: "TensorNetwork") -> Path: @dataclass -class SubTensorNetwork: +class TensorNetwork: name: str - key: int parent_name: str inputs: InputNodes - indices: FrozenSet[Hashable] size_dict: SizeDict - cut_indices: FrozenSet[Hashable] output_indices: Output - - def get_all_input_nodes(self) -> InputNodes: - return self.inputs - - def get_all_networks(self): - return [self] - - def get_output_shape(self): - return tuple([self.size_dict[index] for index in self.output_indices]) - - def find_path(self): - """ - Finds the path for the sub-tensor network. - - Returns: - SubTensorNetworkWithContractTree: The sub-tensor network with the computed contract tree and its cost. - """ - path = greedy_optimizer(self) - - contract_tree, cost = get_contract_tree_and_cost_from_path(self, path) - - tn_with_tree = SubTensorNetworkWithContractTree( - name=self.name, - key=self.key, - parent_name=self.parent_name, - inputs=self.inputs, - indices=self.indices, - size_dict=self.size_dict, - cut_indices=self.cut_indices, - output_indices=self.output_indices, - cost=cost, - contract_tree=contract_tree, - ) - return tn_with_tree - - -@total_ordering -@dataclass -class SubTensorNetworkWithContractTree(SubTensorNetwork): - cost: int - contract_tree: ContractTree - - def get_total_cost(self): - return self.cost + _contract_tree: ContractTree = None + _cost: Optional[int] = None + _ssa_id: Optional[int] = None def get_contract_tree(self): - return self.contract_tree - - def __eq__(self, other): - if not isinstance(other, __class__): - return NotImplemented - return self.cost == other.cost - - def __lt__(self, other): - if not isinstance(other, __class__): - return NotImplemented - # Yes this might seem weird, but we want the one with the highest cost to be the first in the priority queue - return self.cost > other.cost + if self._contract_tree is None: + self.find_path() + return self._contract_tree def find_path(self): """ - Since the path is already computed, this function just returns the current object. - - Returns: - The current object. + Finds the path for the tensor network. """ - return self - - -@dataclass -class SuperTensorNetwork(SubTensorNetwork): - parent_name: str - sub_networks: Sequence[SubTensorNetwork] - - def get_all_input_nodes(self) -> InputNodes: - sub_input_nodes = [ - SubNetworkInputNode( - sub_network.output_indices, - sub_network.get_output_shape(), - sub_network, - ) - for sub_network in self.sub_networks - ] - - return sub_input_nodes + self.inputs - - def get_all_networks(self): - return [self] + [sub_network for sub_network in self.sub_networks] - - def find_path(self): - sub_networks_with_path: List[SubTensorNetworkWithContractTree] = [] - for sub_network in self.sub_networks: - sub_networks_with_path.append(sub_network.find_path()) - - self.sub_networks = sub_networks_with_path - path = greedy_optimizer(self) contract_tree, cost = get_contract_tree_and_cost_from_path(self, path) - - tn_with_tree = SuperTensorNetworkWithTree( - self.name, - self.key, - self.parent_name, - self.inputs, - self.indices, - self.size_dict, - self.cut_indices, - self.output_indices, - cost, - contract_tree, - sub_networks_with_path, - ) - - return tn_with_tree - - -@dataclass -class SuperTensorNetworkWithTree( - SuperTensorNetwork, - SubTensorNetworkWithContractTree, -): - sub_networks: List[SubTensorNetworkWithContractTree] - - # def find_path(self, greedy_optimizer: GreedyOptimizer): - # """ - # Since the path is already computed, this function just returns the current object. - - # Returns: - # The current object. - # """ - # return self - - def get_total_cost(self): - return ( - sum([sub_network.cost for sub_network in self.sub_networks]) - + self.cost - ) - - def get_parent_tree(self): - super_tree = self.get_contract_tree() - - parent_tree = [] - sub_tree_root: Dict[str, ContractTreeNode] = {} - for node in super_tree: - if ( - isinstance(node, SubNetworkInputNode) - and node.sub_network.parent_name == self.name - ): - assert isinstance( - node.sub_network, SubTensorNetworkWithContractTree - ), "The subnetworks should have a contract tree, when calling get_parent_tree" - sub_tree = None - for sn in self.sub_networks: - if sn.name == node.sub_network.name: - sub_tree = sn.contract_tree - assert ( - sub_tree is not None - ), f"Sub tree {node.sub_network.name} not found in {self.name}" - for sub_node in sub_tree: - parent_tree.append(sub_node) - - sub_tree_root[node.sub_network.name] = sub_tree[-1] - elif isinstance(node, IntermediateContractNode): - for key, child in enumerate(node.children): - if ( - isinstance(child, SubNetworkInputNode) - and child.sub_network.parent_name == self.name - ): - node.children[key] = sub_tree_root[ - child.sub_network.name - ] - - parent_tree.append(node) - else: - parent_tree.append(node) - - return parent_tree - - def update_tree(self, name, tree: ContractTree): - if name == self.name: - self.contract_tree = tree - else: - for sub_network in self.sub_networks: - if sub_network.name == name: - sub_network.contract_tree = tree - return - raise Exception(f"name {name} not found in {self.name}") - - def find_path(self): - return self - - -TensorNetwork = Union[SubTensorNetwork, SuperTensorNetwork] -TensorNetworkWithTree = Union[ - SubTensorNetworkWithContractTree, SuperTensorNetworkWithTree -] + self._contract_tree = contract_tree + self._cost = cost def get_sub_networks( - tensor_network_name: str, - input_nodes, - output, - size_dict: SizeDict, + tensor_network: TensorNetwork, imbalance: float, weight_nodes: str = "const", ): + input_nodes = tensor_network.inputs + output = tensor_network.output_indices + size_dict = tensor_network.size_dict + tensor_network_name = tensor_network.name num_input_nodes = len(input_nodes) assert ( num_input_nodes > 2 - ), f"Not enough input nodes to split, pass at least two input nodes, {input_nodes}" + ), f"You need to pass at least two input nodes, {input_nodes}" inputs = [input.indices for input in input_nodes] @@ -468,104 +283,54 @@ def get_sub_networks( cut_indices = frozenset(cut_indices.union(output)) output_indices: list[Output] = [ - cut_indices.intersection(block) for block in block_indices + list(cut_indices.intersection(block)) for block in block_indices ] - output_block_id = 0 if len(output) > 0: output_block_id = block_ids[-1] else: output_block_id = random.choice([0, 1]) sub_networks = [ - SubTensorNetwork( + TensorNetwork( f"{tensor_network_name}.{key}", - key, - None, + tensor_network_name, block_inputs[key], block_indices[key], size_dict, - cut_indices, output_indices[key], ) for key, _ in enumerate(block_inputs) ] - super_sub_network = sub_networks.pop(output_block_id) - super_network = build_super_network( - super_sub_network, - tensor_network_name, - cut_indices, - output, - sub_networks, + parent_sub_network = sub_networks.pop(output_block_id) + child_sub_network = sub_networks.pop() + child_sub_network.parent_name = parent_sub_network.name + sub_network_node = SubNetworkInputNode( + child_sub_network.output_indices, + child_sub_network, ) + parent_sub_network.inputs.append(sub_network_node) + parent_sub_network.output_indices = output - return super_network.find_path() - + return parent_sub_network, child_sub_network -def build_super_network( - super_network: TensorNetwork, - parent_name: str, - cut_indices: FrozenSet[Hashable], - output: Output, - sub_networks: List[SubTensorNetwork], -): - # Set parent for sub_networks - for sub_network in sub_networks: - sub_network.parent_name = super_network.name - - super_network = SuperTensorNetwork( - name=super_network.name, - key=super_network.key, - parent_name=parent_name, - inputs=super_network.inputs, - indices=super_network.indices, - size_dict=super_network.size_dict, - cut_indices=cut_indices, - output_indices=output, - sub_networks=sub_networks, - ) - return super_network +def contract_tree_to_path(root: ContractTreeNode, last_id, path: Path): + if isinstance(root, IntermediateContractNode): + sub_ids = [] + for child in root.children: + sub_root = contract_tree_to_path(child, last_id, path) + last_id = max(last_id, sub_root) + sub_ids.append(sub_root) + path.append(tuple(sub_ids)) + new_root_id = max(last_id, max(sub_ids)) + 1 + return new_root_id -def get_remapped_id(id, input_remap: Optional[Dict[str, int]]): - return input_remap[id] if input_remap != None and id in input_remap else id - - -def contract_tree_to_path( - tree: ContractTree, remap: Optional[Dict[str, int]] = None -): - root = tree[-1] - - if isinstance(root, BasicInputNode): - assert ( - len(tree) == 1 - ), "Tree should only contain one node, if root is basic input" - return [(int(get_remapped_id(root.get_id(), remap)),)] - - path = [] - - counter = (len(tree) + 1) // 2 - uuid_to_ssa_id = {} - for node in tree: - if isinstance(node, IntermediateContractNode): - uuid_to_ssa_id[node.get_id()] = counter - counter += 1 - pair = [] - if isinstance(node.children[0], BasicInputNode): - pair.append(get_remapped_id(node.children[0].get_id(), remap)) - else: - pair.append(uuid_to_ssa_id[node.children[0].get_id()]) - - if len(node.children) > 1: - if isinstance(node.children[1], BasicInputNode): - pair.append( - get_remapped_id(node.children[1].get_id(), remap) - ) - else: - pair.append(uuid_to_ssa_id[node.children[1].get_id()]) - path.append(tuple(pair)) - return path + if isinstance(root, OriginalInputNode): + return int(root.get_id()) + if isinstance(root, SubNetworkInputNode): + return root.sub_network._ssa_id def hybrid_hypercut_greedy( @@ -576,80 +341,57 @@ def hybrid_hypercut_greedy( weight_nodes="const", cutoff=15, ): - # Problem setup, transform arguments to tanser network with path - shapes = [tuple([size_dict[i] for i in input]) for input in inputs] + # Noramlize parameters inputs = [list(input) for input in inputs] output = list(output) - indices = frozenset.union(*[frozenset(input) for input in inputs]) input_nodes: InputNodes = [ - OriginalInputNode(in_sh[0], in_sh[1], id) - for id, in_sh in enumerate(zip(inputs, shapes)) + OriginalInputNode(input, id) for id, input in enumerate(inputs) ] - tensor_network_without_path = SubTensorNetwork( - "tn", 0, None, input_nodes, indices, size_dict, frozenset(), output - ) - - tensor_network = tensor_network_without_path.find_path() - - root_name = tensor_network.name - - # Initialize queues - to_partition: PriorityQueue[TensorNetwork] = PriorityQueue() - to_partition.put(tensor_network) - - # Initialize dictoionary for finalized partitions - partitioned: Dict[str, TensorNetworkWithTree] = {} + tensor_network = TensorNetwork("tn", None, input_nodes, size_dict, output) - while not to_partition.empty(): - next_network = to_partition.get() - if len(next_network.inputs) <= cutoff: - with_path = next_network.find_path() # - partitioned[next_network.name] = with_path + stack = [tensor_network] + path = [] + last_id = len(inputs) - 1 + network_by_name = {tensor_network.name: tensor_network} + cost = 0 + while stack: + tn = stack.pop() + # print(f"Popped {tn.name}") + if len(tn.inputs) <= cutoff: + tree = tn.get_contract_tree() + cost += tn._cost + sub_id = contract_tree_to_path(tree[-1], last_id, path) + last_id = max(last_id, sub_id) + tn._ssa_id = sub_id + while tn.parent_name and len(tn.parent_name) < len(tn.name): + network_by_name[tn.parent_name]._ssa_id = sub_id + tn = network_by_name[tn.parent_name] continue - - super_network = get_sub_networks( - next_network.name, - next_network.get_all_input_nodes(), - next_network.output_indices, - next_network.size_dict, + parent_sub_network, child_sub_network = get_sub_networks( + tn, imbalance=imbalance, weight_nodes=weight_nodes, ) - - partitioned[super_network.parent_name] = copy.deepcopy(super_network) - - to_partition.put(super_network) - - for sub_network in super_network.sub_networks: - to_partition.put(sub_network) - - def merge(network_name) -> ContractTree: - partitioned_network = partitioned[network_name] - # Check if we reached a leave - if partitioned_network.name == network_name or not ( - isinstance(partitioned_network, SuperTensorNetworkWithTree) - ): - with_tree = partitioned_network.find_path() - return with_tree.contract_tree - for sub_network in partitioned_network.get_all_networks(): - merged_sub_tree = merge(sub_network.name) - partitioned_network.update_tree(sub_network.name, merged_sub_tree) - - parent_tree = partitioned_network.get_parent_tree() - - return parent_tree - - parent_tree = merge(root_name) - path = contract_tree_to_path(parent_tree) + stack.append(parent_sub_network) + network_by_name[parent_sub_network.name] = parent_sub_network + # print(f"Pushed parent {parent_sub_network.name}") + stack.append(child_sub_network) + network_by_name[child_sub_network.name] = child_sub_network + # print(f"Pushed child {child_sub_network.name}") + + # contract_tree = tensor_network.build_tree(imbalance, cutoff, weight_nodes) + # path = [] + # contract_tree_to_path(contract_tree[-1], len(inputs) - 1, path) path = [tuple([int(i) for i in pair]) for pair in path] + # print(f"{cost:.6e}", math.log10(cost / 2), imbalance, cutoff, weight_nodes) return ContractionTree.from_path(inputs, output, size_dict, ssa_path=path) hyper_space = { - "imbalance": {"type": "FLOAT", "min": 0.01, "max": 0.2}, - "weight_nodes": {"type": "STRING", "options": ["const", "log"]}, - "cutoff": {"type": "INT", "min": 40, "max": 200}, + "imbalance": {"type": "FLOAT", "min": 0.001, "max": 0.1}, + "weight_nodes": {"type": "STRING", "options": ["log"]}, + "cutoff": {"type": "INT", "min": 50, "max": 100}, } -register_hyper_function("random-kahypar+", hybrid_hypercut_greedy, hyper_space) +register_hyper_function("kahypar+", hybrid_hypercut_greedy, hyper_space) From 1d17e31690df562c705c34df4e54367cd6d9d085 Mon Sep 17 00:00:00 2001 From: Christoph Staudt Date: Sun, 5 May 2024 01:15:06 +0200 Subject: [PATCH 4/8] Remove contract trees --- cotengra/pathfinders/path_kahypar_plus.py | 242 ++++------------------ 1 file changed, 45 insertions(+), 197 deletions(-) diff --git a/cotengra/pathfinders/path_kahypar_plus.py b/cotengra/pathfinders/path_kahypar_plus.py index e81d3e4..1c958f7 100644 --- a/cotengra/pathfinders/path_kahypar_plus.py +++ b/cotengra/pathfinders/path_kahypar_plus.py @@ -2,20 +2,15 @@ # Based on https://github.com/ti2-group/hybrid_contraction_tree_optimizer/ # -import math import random -import uuid -from collections import defaultdict from dataclasses import dataclass, field from typing import ( - Callable, Hashable, Optional, Union, List, Tuple, Dict, - Set, ) from cotengra.pathfinders.path_basic import ( @@ -28,11 +23,9 @@ Inputs = List[List[Hashable]] - Output = List[Hashable] SizeDict = Dict[Hashable, int] Path = List[Tuple[int, ...]] -GreedyOptimizer = Callable[["TensorNetwork"], Path] @dataclass @@ -45,10 +38,7 @@ class OriginalInputNode(BasicInputNode): id: int def get_id(self): - return str(self.id) - - def __repr__(self) -> str: - return f"Original Input({self.indices})" + return self.id @dataclass @@ -56,7 +46,7 @@ class SubNetworkInputNode(BasicInputNode): sub_network: "TensorNetwork" def get_id(self): - return f"sn-{self.sub_network.name}" + return self.sub_network._ssa_id def __repr__(self) -> str: return f"Sub network Input({self.sub_network.output_indices}," @@ -66,119 +56,12 @@ def __repr__(self) -> str: InputNodes = List["InputNode"] -@dataclass -class IntermediateContractNode: - all_indices: Set[Hashable] # Union all indices of children - scale: int - indices: Output - children: List["ContractTreeNode"] - uuid: str = field(default_factory=lambda: str(uuid.uuid4())) - - def get_id(self): - return self.uuid - - def __repr__(self) -> str: - return ( - f"IntermediateContractNode({self.uuid}), " - + f"children: {[child.get_id() for child in self.children]}" - ) - - -ContractTreeNode = Union[ - OriginalInputNode, SubNetworkInputNode, IntermediateContractNode -] - -ContractTree = List[ContractTreeNode] - - -def safe_log2(x): - if x < 1: - return 0 - return math.log2(x) - - -def get_contract_tree_and_cost_from_path( - tn: "TensorNetwork", ssa_path -) -> Tuple[ContractTree, int]: - contract_tree: ContractTree = [] - histogram = defaultdict(lambda: 0) - for input in tn.inputs: - contract_tree.append(input) - for edge in input.indices: - histogram[edge] += 1 - - for index in tn.output_indices: - histogram[index] += 1 - - # If there is only one input thats the whole tree - if len(contract_tree) == 1: - return contract_tree, 1 - total_cost = 0 - - for pair in ssa_path: - if len(pair) == 1: - left_node: ContractTreeNode = contract_tree[pair[0]] - all_indices = set(left_node.indices) - cost = 1 - remove = set() - for index in left_node.indices: - cost = cost * tn.size_dict[index] - if histogram[index] == 0: - remove.add(index) - total_cost += cost - intermediate = all_indices - remove - for index in intermediate: - histogram[index] += 1 - contract_tree.append( - IntermediateContractNode( - all_indices, - int(safe_log2(cost)), - list(intermediate), - [contract_tree[pair[0]]], - ) - ) - if len(pair) == 2: - left_node: ContractTreeNode = contract_tree[pair[0]] - right_node: ContractTreeNode = contract_tree[pair[1]] - all_indices = set(left_node.indices).union(right_node.indices) - cost = 1 - remove = set() - for index in all_indices: - cost = cost * tn.size_dict[index] - - for index in left_node.indices: - histogram[index] -= 1 - if histogram[index] == 0: - remove.add(index) - for index in right_node.indices: - histogram[index] -= 1 - if histogram[index] == 0: - remove.add(index) - total_cost += cost - intermediate = all_indices - remove - - for index in intermediate: - histogram[index] += 1 - - contract_tree.append( - IntermediateContractNode( - all_indices, - int(safe_log2(cost)), - list(intermediate), - [contract_tree[pair[0]], contract_tree[pair[1]]], - ) - ) - total_cost = 2 * total_cost - - return contract_tree, total_cost - - -def greedy_optimizer(tn: "TensorNetwork") -> Path: +def greedy_optimizer(tn: "TensorNetwork") -> Tuple[Path, float]: inputs = [input.indices for input in tn.inputs] output = tn.output_indices size_dict = tn.size_dict - if len(inputs) <= 15: + if len(inputs) <= 12: optimal_opt = OptimalOptimizer() return optimal_opt.ssa_path(inputs, output, size_dict) @@ -193,23 +76,7 @@ class TensorNetwork: inputs: InputNodes size_dict: SizeDict output_indices: Output - _contract_tree: ContractTree = None - _cost: Optional[int] = None - _ssa_id: Optional[int] = None - - def get_contract_tree(self): - if self._contract_tree is None: - self.find_path() - return self._contract_tree - - def find_path(self): - """ - Finds the path for the tensor network. - """ - path = greedy_optimizer(self) - contract_tree, cost = get_contract_tree_and_cost_from_path(self, path) - self._contract_tree = contract_tree - self._cost = cost + _ssa_id: Optional[int] = field(default=None, init=False) def get_sub_networks( @@ -219,8 +86,6 @@ def get_sub_networks( ): input_nodes = tensor_network.inputs output = tensor_network.output_indices - size_dict = tensor_network.size_dict - tensor_network_name = tensor_network.name num_input_nodes = len(input_nodes) assert ( num_input_nodes > 2 @@ -234,7 +99,7 @@ def get_sub_networks( block_ids = kahypar_subgraph_find_membership( inputs, set(), - size_dict, + tensor_network.size_dict, weight_nodes=weight_nodes, weight_edges="log", fix_output_nodes=False, @@ -276,61 +141,52 @@ def get_sub_networks( for block in block_inputs ] - cut_indices = set() - cut_indices = cut_indices.union( - block_indices[0].intersection(block_indices[1]) - ) - - cut_indices = frozenset(cut_indices.union(output)) - output_indices: list[Output] = [ - list(cut_indices.intersection(block)) for block in block_indices - ] + cut_indices = block_indices[0].intersection(block_indices[1]) if len(output) > 0: - output_block_id = block_ids[-1] + parent_block_id = block_ids[-1] else: - output_block_id = random.choice([0, 1]) - - sub_networks = [ - TensorNetwork( - f"{tensor_network_name}.{key}", - tensor_network_name, - block_inputs[key], - block_indices[key], - size_dict, - output_indices[key], - ) - for key, _ in enumerate(block_inputs) - ] + parent_block_id = random.choice([0, 1]) + + child_block_id = 1 - parent_block_id + + parent_sub_network = TensorNetwork( + f"{tensor_network.name}.{parent_block_id}", + tensor_network.name, + block_inputs[parent_block_id], + tensor_network.size_dict, + output, + ) + + child_sub_network = TensorNetwork( + f"{tensor_network.name}.{child_block_id}", + parent_sub_network.name, + block_inputs[child_block_id], + tensor_network.size_dict, + cut_indices, + ) - parent_sub_network = sub_networks.pop(output_block_id) - child_sub_network = sub_networks.pop() - child_sub_network.parent_name = parent_sub_network.name sub_network_node = SubNetworkInputNode( child_sub_network.output_indices, child_sub_network, ) parent_sub_network.inputs.append(sub_network_node) - parent_sub_network.output_indices = output return parent_sub_network, child_sub_network -def contract_tree_to_path(root: ContractTreeNode, last_id, path: Path): - if isinstance(root, IntermediateContractNode): - sub_ids = [] - for child in root.children: - sub_root = contract_tree_to_path(child, last_id, path) - last_id = max(last_id, sub_root) - sub_ids.append(sub_root) - path.append(tuple(sub_ids)) - new_root_id = max(last_id, max(sub_ids)) + 1 - return new_root_id +def extend_path(tn: TensorNetwork, sub_path: Path, last_id, path: Path): + n = len(tn.inputs) + for pair in sub_path: + new_pair = [] + for element in pair: + if element < n: + new_pair.append(int(tn.inputs[element].get_id())) + else: + new_pair.append(last_id - n + element + 1) + path.append(tuple(new_pair)) - if isinstance(root, OriginalInputNode): - return int(root.get_id()) - if isinstance(root, SubNetworkInputNode): - return root.sub_network._ssa_id + return last_id + len(sub_path) def hybrid_hypercut_greedy( @@ -358,15 +214,13 @@ def hybrid_hypercut_greedy( cost = 0 while stack: tn = stack.pop() - # print(f"Popped {tn.name}") if len(tn.inputs) <= cutoff: - tree = tn.get_contract_tree() - cost += tn._cost - sub_id = contract_tree_to_path(tree[-1], last_id, path) - last_id = max(last_id, sub_id) - tn._ssa_id = sub_id + sub_path, sub_cost = greedy_optimizer(tn) + cost += sub_cost + last_id = extend_path(tn, sub_path, last_id, path) + tn._ssa_id = last_id while tn.parent_name and len(tn.parent_name) < len(tn.name): - network_by_name[tn.parent_name]._ssa_id = sub_id + network_by_name[tn.parent_name]._ssa_id = last_id tn = network_by_name[tn.parent_name] continue parent_sub_network, child_sub_network = get_sub_networks( @@ -376,22 +230,16 @@ def hybrid_hypercut_greedy( ) stack.append(parent_sub_network) network_by_name[parent_sub_network.name] = parent_sub_network - # print(f"Pushed parent {parent_sub_network.name}") stack.append(child_sub_network) network_by_name[child_sub_network.name] = child_sub_network - # print(f"Pushed child {child_sub_network.name}") - # contract_tree = tensor_network.build_tree(imbalance, cutoff, weight_nodes) - # path = [] - # contract_tree_to_path(contract_tree[-1], len(inputs) - 1, path) - path = [tuple([int(i) for i in pair]) for pair in path] # print(f"{cost:.6e}", math.log10(cost / 2), imbalance, cutoff, weight_nodes) return ContractionTree.from_path(inputs, output, size_dict, ssa_path=path) hyper_space = { - "imbalance": {"type": "FLOAT", "min": 0.001, "max": 0.1}, + "imbalance": {"type": "FLOAT", "min": 0.001, "max": 0.2}, "weight_nodes": {"type": "STRING", "options": ["log"]}, - "cutoff": {"type": "INT", "min": 50, "max": 100}, + "cutoff": {"type": "INT", "min": 60, "max": 100}, } register_hyper_function("kahypar+", hybrid_hypercut_greedy, hyper_space) From 697e93fc54db9a2ba653b35d1494df9aae9b6664 Mon Sep 17 00:00:00 2001 From: Christoph Staudt Date: Sun, 5 May 2024 01:15:28 +0200 Subject: [PATCH 5/8] Return flops from randomtopimizer --- cotengra/pathfinders/path_basic.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cotengra/pathfinders/path_basic.py b/cotengra/pathfinders/path_basic.py index 5a139c1..df52b74 100644 --- a/cotengra/pathfinders/path_basic.py +++ b/cotengra/pathfinders/path_basic.py @@ -1369,7 +1369,7 @@ def ssa_path(self, inputs, output, size_dict, **kwargs): self.best_ssa_path = ssa_path self.best_flops = flops - return self.best_ssa_path + return self.best_ssa_path, self.best_flops def search(self, inputs, output, size_dict, **kwargs): from ..core import ContractionTree @@ -1411,7 +1411,6 @@ def get_optimize_optimal(accel="auto"): if accel is False: return optimize_optimal - raise ValueError(f"Unrecognized value for `accel`: {accel}.") From e147c0a451a663b106253f974061542bb24cfe8a Mon Sep 17 00:00:00 2001 From: Christoph Staudt Date: Sun, 5 May 2024 02:03:48 +0200 Subject: [PATCH 6/8] Remove cost tracking, fix cut indices, fix pep --- cotengra/pathfinders/path_basic.py | 2 +- cotengra/pathfinders/path_kahypar_plus.py | 25 ++++++++++++----------- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/cotengra/pathfinders/path_basic.py b/cotengra/pathfinders/path_basic.py index df52b74..01217a1 100644 --- a/cotengra/pathfinders/path_basic.py +++ b/cotengra/pathfinders/path_basic.py @@ -1369,7 +1369,7 @@ def ssa_path(self, inputs, output, size_dict, **kwargs): self.best_ssa_path = ssa_path self.best_flops = flops - return self.best_ssa_path, self.best_flops + return self.best_ssa_path def search(self, inputs, output, size_dict, **kwargs): from ..core import ContractionTree diff --git a/cotengra/pathfinders/path_kahypar_plus.py b/cotengra/pathfinders/path_kahypar_plus.py index 1c958f7..00a5af6 100644 --- a/cotengra/pathfinders/path_kahypar_plus.py +++ b/cotengra/pathfinders/path_kahypar_plus.py @@ -56,7 +56,7 @@ def __repr__(self) -> str: InputNodes = List["InputNode"] -def greedy_optimizer(tn: "TensorNetwork") -> Tuple[Path, float]: +def greedy_optimizer(tn: "TensorNetwork") -> Tuple[Path]: inputs = [input.indices for input in tn.inputs] output = tn.output_indices size_dict = tn.size_dict @@ -65,7 +65,7 @@ def greedy_optimizer(tn: "TensorNetwork") -> Tuple[Path, float]: optimal_opt = OptimalOptimizer() return optimal_opt.ssa_path(inputs, output, size_dict) - greedy_opt = RandomGreedyOptimizer(max_repeats=512) + greedy_opt = RandomGreedyOptimizer(max_repeats=64) return greedy_opt.ssa_path(inputs, output, size_dict) @@ -111,7 +111,7 @@ def get_sub_networks( quiet=True, ) - ## Noramlize block ids + # Noramlize block ids # Check if all input nodes were assigned to the same block input_block_ids = block_ids[:num_input_nodes] @@ -127,7 +127,7 @@ def get_sub_networks( assert ( len(set(input_block_ids)) == 2 - ), f"There should be two blocks, {input_block_ids}, {min_block_id}, {max_block_id}" + ), f"There should be two blocks, {input_block_ids}" # Group inputs by block id block_inputs: list[InputNodes] = [[], []] @@ -141,7 +141,8 @@ def get_sub_networks( for block in block_inputs ] - cut_indices = block_indices[0].intersection(block_indices[1]) + # Include output indices in cut, since it is not in block indices + cut_indices = block_indices[0].intersection(block_indices[1]).union(output) if len(output) > 0: parent_block_id = block_ids[-1] @@ -149,6 +150,9 @@ def get_sub_networks( parent_block_id = random.choice([0, 1]) child_block_id = 1 - parent_block_id + child_output = list( + cut_indices.intersection(block_indices[child_block_id]) + ) parent_sub_network = TensorNetwork( f"{tensor_network.name}.{parent_block_id}", @@ -163,7 +167,7 @@ def get_sub_networks( parent_sub_network.name, block_inputs[child_block_id], tensor_network.size_dict, - cut_indices, + child_output, ) sub_network_node = SubNetworkInputNode( @@ -211,12 +215,10 @@ def hybrid_hypercut_greedy( path = [] last_id = len(inputs) - 1 network_by_name = {tensor_network.name: tensor_network} - cost = 0 while stack: tn = stack.pop() if len(tn.inputs) <= cutoff: - sub_path, sub_cost = greedy_optimizer(tn) - cost += sub_cost + sub_path = greedy_optimizer(tn) last_id = extend_path(tn, sub_path, last_id, path) tn._ssa_id = last_id while tn.parent_name and len(tn.parent_name) < len(tn.name): @@ -233,13 +235,12 @@ def hybrid_hypercut_greedy( stack.append(child_sub_network) network_by_name[child_sub_network.name] = child_sub_network - # print(f"{cost:.6e}", math.log10(cost / 2), imbalance, cutoff, weight_nodes) return ContractionTree.from_path(inputs, output, size_dict, ssa_path=path) hyper_space = { "imbalance": {"type": "FLOAT", "min": 0.001, "max": 0.2}, - "weight_nodes": {"type": "STRING", "options": ["log"]}, - "cutoff": {"type": "INT", "min": 60, "max": 100}, + "weight_nodes": {"type": "STRING", "options": ["log", "const"]}, + "cutoff": {"type": "INT", "min": 10, "max": 200}, } register_hyper_function("kahypar+", hybrid_hypercut_greedy, hyper_space) From 17133724a28e130f2191944b270a9e914ac58ef5 Mon Sep 17 00:00:00 2001 From: Christoph Staudt Date: Thu, 11 Jul 2024 14:27:48 +0200 Subject: [PATCH 7/8] Add kahypar+ to text matrix --- tests/test_optimizers.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py index 2d4d6ca..355bebc 100644 --- a/tests/test_optimizers.py +++ b/tests/test_optimizers.py @@ -72,6 +72,10 @@ def contraction_20_5(): functools.partial(ctg.UniformOptimizer, methods="kahypar-agglom"), "kahypar", ), + ( + functools.partial(ctg.UniformOptimizer, methods="kahypar+"), + "kahypar", + ), ( functools.partial(ctg.UniformOptimizer, methods="betweenness"), "igraph", From 4194112edf402f4e922eab3fbcf74e014b033825 Mon Sep 17 00:00:00 2001 From: Christoph Staudt Date: Thu, 11 Jul 2024 14:28:04 +0200 Subject: [PATCH 8/8] Add kahypar_plus to all --- cotengra/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/cotengra/__init__.py b/cotengra/__init__.py index 5d60f14..6b20732 100644 --- a/cotengra/__init__.py +++ b/cotengra/__init__.py @@ -171,6 +171,7 @@ "path_compressed_greedy", "path_igraph", "path_kahypar", + "path_kahypar_plus", "path_labels", "plot_contractions_alt", "plot_contractions",