def get_shared_edges( self, node1: network_components.BaseNode, node2: network_components.BaseNode) -> Set[network_components.Edge]: """Get all edges shared between two nodes. Args: node1: The first node. node2: The second node. Returns: A (possibly empty) `set` of `Edge`s shared by the nodes. """ return network_components.get_shared_edges(node1, node2)
def ncon( tensors: Sequence[Union[network_components.BaseNode, Tensor]], network_structure: Sequence[Sequence], con_order: Optional[Sequence] = None, out_order: Optional[Sequence] = None, backend: Optional[Text] = None ) -> Union[network_components.BaseNode, Tensor]: r"""Contracts a list of tensors or nodes according to a tensor network specification. The network is provided as a list of lists, one for each tensor, specifying labels for the edges connected to that tensor. If a contraction order `con_order` and an output order `out_order` are both provided, the edge labels can be anything. Otherwise (`con_order == None or out_order == None`), the edge labels must be nonzero integers and edges will be contracted in ascending order. Negative integers denote the (dangling) indices of the output tensor, which will be in descending order, e.g. [-1,-2,-3,...]. For example, matrix multiplication: ```python A = np.array([[1.0, 2.0], [3.0, 4.0]]) B = np.array([[1.0, 1.0], [0.0, 1.0]]) ncon([A,B], [(-1, 1), (1, -2)]) ``` Matrix trace: ```python A = np.array([[1.0, 2.0], [3.0, 4.0]]) ncon([A], [(1, 1)]) # 5.0 ``` Note: The reason `0` is not allowed as an edge label without manually specifying the contraction order is to maintain compatibility with the [original NCON implementation](https://arxiv.org/abs/1402.0939). However, the use of `0` in `con_order` to denote outer products is not (currently) supported in this implementation. Args: tensors: List of `Tensor`s or `BaseNode`s. network_structure: List of lists specifying the tensor network structure. con_order: List of edge labels specifying the contraction order. out_order: List of edge labels specifying the output order. backend: String specifying the backend to use. Defaults to `tensornetwork.config.default_backend`. Returns: The result of the contraction. The result is returned as a `Node` if all elements of `tensors` are `BaseNode` objects, else it is returned as a `Tensor` object. """ if backend and (backend not in backend_factory._BACKENDS): raise ValueError("Backend '{}' does not exist".format(backend)) if backend is None: backend = config.default_backend are_nodes = [isinstance(t, network_components.BaseNode) for t in tensors] nodes = {t for t in tensors if isinstance(t, network_components.BaseNode)} if not all([n.backend.name == backend for n in nodes]): raise ValueError( "Some nodes have backends different from '{}'".format(backend)) _tensors = [] for t in tensors: if isinstance(t, network_components.BaseNode): _tensors.append(t.tensor) else: _tensors.append(t) nodes, con_edges, out_edges = ncon_network(_tensors, network_structure, con_order=con_order, out_order=out_order, backend=backend) nodes = set(nodes) # we don't need the ordering here # Reverse the list so we can pop from the end: O(1). con_edges = con_edges[::-1] while con_edges: nodes_to_contract = con_edges[-1].get_nodes() edges_to_contract = network_components.get_shared_edges( *nodes_to_contract) # Eat up all parallel edges that are adjacent in the ordering. adjacent_parallel_edges = set() for edge in reversed(con_edges): if edge in edges_to_contract: adjacent_parallel_edges.add(edge) else: break con_edges = con_edges[:-len(adjacent_parallel_edges)] # In an optimal ordering, all edges connecting a given pair of nodes are # adjacent in con_order. If this is not the case, warn the user. leftovers = edges_to_contract - adjacent_parallel_edges if leftovers: warnings.warn( "Suboptimal ordering detected. Edges {} are not adjacent in the " "contraction order to edges {}, connecting nodes {}. Deviating from " "the specified ordering!".format( list(map(str, leftovers)), list(map(str, adjacent_parallel_edges)), list(map(str, nodes_to_contract)))) con_edges = [e for e in con_edges if e not in edges_to_contract] if set(nodes_to_contract) == nodes: # This contraction produces the final output, so order the edges # here to avoid transposes in some cases. contraction_output_order = out_edges else: contraction_output_order = None nodes = nodes - set(nodes_to_contract) nodes.add( network_components.contract_between( *nodes_to_contract, name="con({},{})".format(*nodes_to_contract), output_edge_order=contraction_output_order)) # TODO: More efficient ordering of products based on out_edges res_node = network_components.outer_product_final_nodes(nodes, out_edges) if all(are_nodes): return res_node return res_node.tensor
def _jittable_ncon(tensors, network_structure, con_order, out_order, backend): """Jittable Ncon function. Args: tensors: List of tensors. network_structure: List of list of integers that descripes the network structure. con_order: Order of the contraction. out_order: Order of the final axis order. backend: A backend object. Returns: The final tensor after contraction. """ nodes, con_edges, out_edges = ncon_network(tensors, network_structure, con_order=con_order, out_order=out_order, backend=backend) nodes = set(nodes) # we don't need the ordering here # Reverse the list so we can pop from the end: O(1). con_edges = con_edges[::-1] while con_edges: nodes_to_contract = con_edges[-1].get_nodes() edges_to_contract = network_components.get_shared_edges( *nodes_to_contract) # Eat up all parallel edges that are adjacent in the ordering. adjacent_parallel_edges = set() for edge in reversed(con_edges): if edge in edges_to_contract: adjacent_parallel_edges.add(edge) else: break con_edges = con_edges[:-len(adjacent_parallel_edges)] # In an optimal ordering, all edges connecting a given pair of nodes are # adjacent in con_order. If this is not the case, warn the user. leftovers = edges_to_contract - adjacent_parallel_edges if leftovers: warnings.warn( "Suboptimal ordering detected. Edges {} are not adjacent in the " "contraction order to edges {}, connecting nodes {}. Deviating from " "the specified ordering!".format( list(map(str, leftovers)), list(map(str, adjacent_parallel_edges)), list(map(str, nodes_to_contract)))) con_edges = [e for e in con_edges if e not in edges_to_contract] if set(nodes_to_contract) == nodes: # This contraction produces the final output, so order the edges # here to avoid transposes in some cases. contraction_output_order = out_edges else: contraction_output_order = None nodes = nodes - set(nodes_to_contract) nodes.add( network_components.contract_between( *nodes_to_contract, name="con({},{})".format(*nodes_to_contract), output_edge_order=contraction_output_order)) # TODO: More efficient ordering of products based on out_edges res_node = network_components.outer_product_final_nodes(nodes, out_edges) return res_node.tensor