def test_eigh_vs_backend(backend, dtype): shape = (3, 6, 4, 4) dtype = testing_utils.np_dtype_to_backend(backend, dtype) tensor = tensornetwork.ones(shape, backend=backend, dtype=dtype) tn_result = linalg.eigh(tensor) if backend is None: backend = backend_contextmanager.get_default_backend() backend_obj = backends.backend_factory.get_backend(backend) backend_result = backend_obj.eigh(tensor.array) tn_arrays = [t.array for t in tn_result] for tn_arr, backend_arr in zip(tn_arrays, backend_result): np.testing.assert_allclose(tn_arr, backend_arr)
def test_eigh_vs_backend(backend, dtype): np.random.seed(10) shape = (4, 4) dtype = testing_utils.np_dtype_to_backend(backend, dtype) tensor = initialize_hermitian_matrix(backend, shape, dtype) tn_result = linalg.eigh(tensor) if backend is None: backend = backend_contextmanager.get_default_backend() backend_obj = backends.backend_factory.get_backend(backend) backend_result = backend_obj.eigh(tensor.array) tn_arrays = [t.array for t in tn_result] for tn_arr, backend_arr in zip(tn_arrays, backend_result): testing_utils.assert_allclose(tn_arr, backend_arr, backend_obj)
def test_rq_vs_backend(backend, dtype): shape = (3, 6, 4, 2) dtype = testing_utils.np_dtype_to_backend(backend, dtype) tensor = tensornetwork.ones(shape, backend=backend, dtype=dtype) split_axis = 1 tn_result = linalg.rq(tensor, split_axis, non_negative_diagonal=False) if backend is None: backend = backend_contextmanager.get_default_backend() backend_obj = backends.backend_factory.get_backend(backend) backend_result = backend_obj.rq(tensor.array, split_axis) tn_arrays = [t.array for t in tn_result] for tn_arr, backend_arr in zip(tn_arrays, backend_result): np.testing.assert_allclose(tn_arr, backend_arr)
def test_rq_vs_backend(backend, dtype): np.random.seed(10) shape = (3, 6, 4, 2) dtype = testing_utils.np_dtype_to_backend(backend, dtype) tensor = initialize_tensor('randn', backend, shape, dtype) split_axis = 1 tn_result = linalg.rq(tensor, split_axis, non_negative_diagonal=False) if backend is None: backend = backend_contextmanager.get_default_backend() backend_obj = backends.backend_factory.get_backend(backend) backend_result = backend_obj.rq(tensor.array, split_axis) tn_arrays = [t.array for t in tn_result] for tn_arr, backend_arr in zip(tn_arrays, backend_result): testing_utils.assert_allclose(tn_arr, backend_arr, backend_obj)
def set_tensornetwork_backend(backend: Optional[str] = None) -> None: """ set the runtime backend of tensornetwork :param backend: numpy, tensorflow, jax, pytorch :return: """ if not backend: backend = get_default_backend() backend_obj = get_backend(backend) for module in modules: if module in sys.modules: setattr(sys.modules[module], "backend", backend_obj) tn.set_default_backend(backend)
def test_init_tensor_default_backend(dtype): """ Creates a numpy array, initializes a Tensor from it, and checks that all its members have been correctly initialized. """ backend = backend_contextmanager.get_default_backend() backend_obj = backends.backend_factory.get_backend(backend) shape = (3, 5, 2) testA = backend_obj.zeros(shape, dtype=dtype) init = np.zeros(shape, dtype=dtype) A = tensornetwork.Tensor(init) assert A.backend.name == backend np.testing.assert_allclose(A.array, testA) assert A.shape == testA.shape assert A.size == testA.size assert A.ndim == testA.ndim
def test_svd_vs_backend(backend, dtype): np.random.seed(10) shape = (3, 6, 4, 6) dtype = testing_utils.np_dtype_to_backend(backend, dtype) tensor = initialize_tensor('randn', backend, shape, dtype) split_axis = 1 max_singular_values = 5 max_trunc_error = 0.1 relative = True tn_result = linalg.svd(tensor, split_axis, max_singular_values=max_singular_values, max_truncation_error=max_trunc_error, relative=relative) if backend is None: backend = backend_contextmanager.get_default_backend() backend_obj = backends.backend_factory.get_backend(backend) backend_result = backend_obj.svd(tensor.array, split_axis, max_singular_values=max_singular_values, max_truncation_error=max_trunc_error, relative=relative) tn_arrays = [t.array for t in tn_result] for tn_arr, backend_arr in zip(tn_arrays, backend_result): testing_utils.assert_allclose(tn_arr, backend_arr, backend_obj)
def initialize_node(fname: Text, *fargs: Any, name: Optional[Text] = None, axis_names: Optional[List[Text]] = None, backend: Optional[Union[Text, BaseBackend]] = None, **fkwargs: Any) -> Tensor: """Return a Node wrapping data obtained by an initialization function implemented in a backend. The Node will have the same shape as the underlying array that function generates, with all Edges dangling. This function is not intended to be called directly, but doing so should be safe enough. Args: fname: Name of the method of backend to call (a string). *fargs: Positional arguments to the initialization method. name: Optional name of the Node. axis_names: Optional names of the Node's dangling edges. backend: The backend or its name. **fkwargs: Keyword arguments to the initialization method. Returns: node: A Node wrapping data generated by (the_backend).fname(*fargs, **fkwargs), with one dangling edge per axis of data. """ if backend is None: backend_obj = backend_contextmanager.get_default_backend() else: backend_obj = backends.backend_factory.get_backend(backend) func = getattr(backend_obj, fname) data = func(*fargs, **fkwargs) node = network_components.Node(data, name=name, axis_names=axis_names, backend=backend) return node
def __init__( self, tensors: List[Tensor], center_position: Optional[int] = None, connector_matrix: Optional[Tensor] = None, backend: Optional[Union[Text, AbstractBackend]] = None) -> None: """Initialize a BaseMPS. Args: tensors: A list of `Tensor` objects. center_position: The initial position of the center site. connector_matrix: A `Tensor` of rank 2 connecting different unitcells. A value `None` is equivalent to an identity `connector_matrix`. backend: The name of the backend that should be used to perform contractions. Available backends are currently 'numpy', 'tensorflow', 'pytorch', 'jax' """ if (center_position is not None) and (center_position < 0 or center_position >= len(tensors)): raise ValueError( "`center_position = {}` is different from `None` and " "not between 0 <= center_position < {}".format( center_position, len(tensors))) if backend is None: backend = get_default_backend() if isinstance(backend, AbstractBackend): self.backend = backend else: self.backend = backend_factory.get_backend(backend) # the dtype is deduced from the tensor object. self.tensors = [self.backend.convert_to_tensor(t) for t in tensors] if not all( [self.tensors[0].dtype == tensor.dtype for tensor in self.tensors]): raise TypeError('not all dtypes in BaseMPS.tensors are the same') self.connector_matrix = connector_matrix self.center_position = center_position ######################################################################## ########## define functions for jitted operations ########## ######################################################################## @partial(jit, backend=self.backend, static_argnums=(1, )) def svd(tensor, max_singular_values=None): return self.backend.svd(tensor=tensor, pivot_axis=2, max_singular_values=max_singular_values) self.svd = svd @partial(jit, backend=self.backend) def qr(tensor): return self.backend.qr(tensor, 2) self.qr = qr @partial(jit, backend=self.backend) def rq(tensor): return self.backend.rq(tensor, 1) self.rq = rq self.norm = self.backend.jit(self.backend.norm)
def __init__(self, tensors: List[Tensor], center_position: Optional[int] = None, connector_matrix: Optional[Tensor] = None, backend: Optional[Union[Text, BaseBackend]] = None) -> None: """Initialize a BaseMPS. Args: tensors: A list of `Tensor` objects. center_position: The initial position of the center site. connector_matrix: A `Tensor` of rank 2 connecting different unitcells. A value `None` is equivalent to an identity `connector_matrix`. backend: The name of the backend that should be used to perform contractions. Available backends are currently 'numpy', 'tensorflow', 'pytorch', 'jax' """ if (center_position is not None) and (center_position < 0 or center_position >= len(tensors)): raise ValueError("`center_position = {}` is different from `None` and " "not between 0 <= center_position < {}".format( center_position, len(tensors))) if backend is None: backend = get_default_backend() if isinstance(backend, BaseBackend): self.backend = backend else: self.backend = backend_factory.get_backend(backend) # the dtype is deduced from the tensor object. self.tensors = [self.backend.convert_to_tensor(t) for t in tensors] if not all( [self.tensors[0].dtype == tensor.dtype for tensor in self.tensors]): raise TypeError('not all dtypes in BaseMPS.tensors are the same') self.connector_matrix = connector_matrix self.center_position = center_position ######################################################################## ########## define functions for jitted operations ########## ######################################################################## def qr_decomposition(tensor): return self.backend.qr_decomposition(tensor, 2) self.qr_decomposition = self.backend.jit(qr_decomposition) def rq_decomposition(tensor): return self.backend.rq_decomposition(tensor, 1) self.rq_decomposition = self.backend.jit(rq_decomposition) def left_transfer_operator(A, l, Abar): return ncon([A, l, Abar], [[1, 2, -1], [1, 3], [3, 2, -2]], backend=self.backend.name) self.left_transfer_operator = self.backend.jit(left_transfer_operator) def right_transfer_operator(B, r, Bbar): return ncon([B, r, Bbar], [[-1, 2, 1], [1, 3], [-2, 2, 3]], backend=self.backend.name) self.right_transfer_operator = self.backend.jit(right_transfer_operator) def lcontract(R, tensor): return ncon([R, tensor], [[-1, 1], [1, -2, -3]], backend=self.backend.name) self.lcontract = self.backend.jit(lcontract) def rcontract(tensor, R): return ncon([tensor, R], [[-1, -2, 1], [1, -3]], backend=self.backend.name) self.rcontract = self.backend.jit(rcontract) self.norm = self.backend.jit(self.backend.norm)
def jit(fun: Callable, backend: Union[Text, AbstractBackend] = None, backend_argnum: Optional[int] = None, static_argnums: Union[int, Iterable[int]] = (), device=None, xla_backend: Optional[str] = None) -> Callable: """ Return a jitted or graph-compiled version of `fun` for JAX backend. For all other backends returns `fun`. Args: fun: Callable backend: The backend. backend_argnum: Labels the argument of the decorated function which specifies the backend. This argument will be treated as static in the sense of static_argnums. If backend_argnum is specified, backend must be None. static_argnums: Label the arguments which will be statically compiled against. xla_backend: Specifies the backend ('gpu', 'cpu'...) against which XLA is to run. donate_argnums: Labels arguments that Jit is allowed to overwrite. args: Arguments to `fun`. kwargs: Keyword arguments to `fun`. Raises: ValueError: If backend_argnum is specified but backend is not None. If backend_argnum is specified but the corresponding argument neither is nor labels a backend. Returns: Callable: jitted/graph-compiled version of `fun`, or just `fun`. """ argnum_mode = False if backend_argnum is not None: if backend is not None: raise ValueError("backend must be None if backend_argnum is specified.") argnum_mode = True static_argnums = tuple(list(static_argnums) + [backend_argnum,]) if not argnum_mode: if backend is None: backend = backend_contextmanager.get_default_backend() backend_obj = backends.backend_factory.get_backend(backend) @functools.wraps(fun) def wrapper(*args, **kwargs): jitted = backend_obj.jit(fun, static_argnums=static_argnums, device=device, backend=xla_backend) return jitted(*args, **kwargs) else: @functools.wraps(fun) def wrapper(*args, **kwargs): backend = args[backend_argnum] try: backend_obj = backends.backend_factory.get_backend(backend) except ValueError as error: errstr = (f"backend_argnum={backend_argnum} was specified" f"but the corresponding argument {args[backend_argnum]}" f"did not specify a backend.") raise ValueError(errstr) from error jitted = backend_obj.jit(fun, static_argnums=static_argnums, device=device, backend=xla_backend) return jitted(*args, **kwargs) return wrapper
def ncon( tensors: Sequence[Union[tn_tensor.Tensor, Tensor]], network_structure: Sequence[Sequence[Union[str, int]]], con_order: Optional[Sequence] = None, out_order: Optional[Sequence] = None, check_network: bool = True, backend: Optional[Union[Text, AbstractBackend]] = None ) -> Union[tn_tensor.Tensor, Tensor]: r"""Contracts a list of backend-tensors or `Tensor`s according to a tensor network specification. The network is provided as a list of lists, one for each tensor, specifying the labels for the edges connected to that tensor. Labels can be any numbers or strings. Negative number-type labels and string-type labels with a prepended hyphen ('-') are open labels and remain uncontracted. Positive number-type labels and string-type labels with no prepended hyphen ('-') are closed labels and are contracted. Any open label appearing more than once is treated as an open batch label. Any closed label appearing more than once is treated as a closed batch label. Upon finishing the contraction, all open batch labels will have been collapsed into a single dimension, and all closed batch labels will have been summed over. If `out_order = None`, output labels are ordered according to descending number ordering and ascending ASCII ordering, with number labels always appearing before string labels. Example: network_structure = [[-1, 1, '-rick', '2',-2], [-2, '2', 1, '-morty']] results in an output order of [-1, -2, '-morty', '-rick']. If `out_order` is given, the indices of the resulting tensor will be transposed into this order. If `con_order = None`, `ncon` will first contract all number labels in ascending order followed by all string labels in ascending ASCII order. If `con_order` is given, `ncon` will contract according to this order. For example, matrix multiplication: .. code-block:: 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: .. code-block:: python A = np.array([[1.0, 2.0], [3.0, 4.0]]) ncon([A], [(1, 1)]) # 5.0 Note: Disallowing `0` as an edge label is legacy behaviour, see `original NCON implementation`_. .. _original NCON implementation: https://arxiv.org/abs/1402.0939 Args: tensors: List of backend-tensors or `Tensor`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. check_network: Boolean flag. If `True` check the network. backend: String specifying the backend to use. Defaults to `tensornetwork.backend_contextmanager.get_default_backend`. Returns: The result of the contraction: * A backend-tensor: If all elements of `tensors` are backend-tensors. * A `Tensor`: If all elements of `tensors` are `Tensor` objects. """ # TODO (mganahl): for certain cases np.einsum is still faster than ncon: # - contractions containing batched outer products with small dimensions # This should eventually be fixed, but it's not a priority. if backend is None: backend = get_default_backend() if isinstance(backend, AbstractBackend): backend_obj = backend else: backend_obj = backend_factory.get_backend(backend) if out_order == []: #allow empty list as input out_order = None if con_order == []: #allow empty list as input con_order = None are_tensors = [isinstance(t, tn_tensor.Tensor) for t in tensors] tensors_set = {t for t in tensors if isinstance(t, tn_tensor.Tensor)} if not all([n.backend.name == backend_obj.name for n in tensors_set]): raise ValueError( "Some tensors have backends different from '{}'".format( backend_obj.name)) _tensors = [] for t in tensors: if isinstance(t, tn_tensor.Tensor): _tensors.append(t.array) else: _tensors.append(t) _tensors = [backend_obj.convert_to_tensor(t) for t in _tensors] if check_network: _check_network(network_structure, [t.shape for t in _tensors], con_order, out_order) network_structure, mapping = _canonicalize_network_structure( network_structure) flat_labels = [l for sublist in network_structure for l in sublist] unique_flat_labels = list(set(flat_labels)) if out_order is None: # negative batch labels (negative labels appearing more than once) # are subject to the same output ordering as regular output labels out_order = sorted([l for l in unique_flat_labels if l < 0], reverse=True) else: out_order = [mapping[o] for o in out_order] if con_order is None: # canonicalization of network structure takes care of appropriate # contraction ordering (i.e. use ASCII ordering for str and # regular ordering for int) # all positive labels appearing are considered proper contraction labels. con_order = sorted([l for l in unique_flat_labels if l > 0]) else: con_order = [mapping[o] for o in con_order] if backend not in _CACHED_JITTED_NCONS: _CACHED_JITTED_NCONS[backend] = backend_obj.jit(_jittable_ncon, static_argnums=(1, 2, 3, 4, 5)) sizes = tuple([len(l) for l in network_structure]) res_tensor = _CACHED_JITTED_NCONS[backend](_tensors, tuple(flat_labels), sizes, tuple(con_order), tuple(out_order), backend_obj) if all(are_tensors): return tn_tensor.Tensor(res_tensor, backend=backend_obj) return res_tensor
def ncon( tensors: Sequence[Union[network_components.BaseNode, Tensor]], network_structure: Sequence[Sequence], con_order: Optional[Sequence] = None, out_order: Optional[Sequence] = None, check_network: bool = True, backend: Optional[Union[Text, BaseBackend]] = 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. Labels appearing only once in `network_structure` (open labels) remain uncontracted, labels appearing twice (contracted labels) are contracted over. If `out_order = None`, output labels can either be negative numbers or strings with a hyphen character ('-') prepended, e.g. '-out_label_1'. If `out_order = None` output labels are ordered according to descending number ordering and ascending ASCII ordering, with number labels always appearing before string labels. Example: network_structure = [[-1, 1, '-3', '2'], [-2, '2', 1, '-33']] results in an output order of [-1, -2, '-3', '-33']. If `out_order` is given, the indices of the resulting tensor will be transposed into this order. In this case output labels can be arbitrary numbers and arbitrary strings (no minus or hyphen necessary). If `con_order = None`, `ncon` will first contract all number labels in ascending order followed by all string labels in ascending ASCII order. If `con_order` is given, `ncon` will contract according to this order. For example, matrix multiplication: .. code-block:: 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: .. code-block:: 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`_. However, the use of `0` in `con_order` to denote outer products is not (currently) supported in this implementation. .. _original NCON implementation: https://arxiv.org/abs/1402.0939 Args: tensors: List of `Tensors` or `BaseNodes`. 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. check_network: Boolean flag. If `True` check the network. backend: String specifying the backend to use. Defaults to `tensornetwork.backend_contextmanager.get_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 is None: backend = get_default_backend() if isinstance(backend, BaseBackend): backend_obj = backend else: backend_obj = backend_factory.get_backend(backend) if out_order == []: #allow empty list as input out_order = None if con_order == []: #allow empty list as input con_order = None # convert to lists network_structure = [list(l) for l in network_structure] 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_obj.name for n in nodes]): raise ValueError("Some nodes have backends different from '{}'".format( backend_obj.name)) _tensors = [] for t in tensors: if isinstance(t, network_components.BaseNode): _tensors.append(t.tensor) else: _tensors.append(t) _tensors = [backend_obj.convert_to_tensor(t) for t in _tensors] if check_network: cont_labels, out_labels = _check_network(network_structure, [t.shape for t in _tensors], con_order, out_order) else: # map the network strcuture to integers; if any of the labels is a `str` # type, the ordering defaults to string-ordering, i.e. # [[1, 2, '12'], [1, 9]] -> [[1, -2, -1],[1, -3]] # pylint: disable=line-too-long int_cont_labels, str_cont_labels, int_out_labels, str_out_labels = _get_cont_out_labels( network_structure) cont_labels = int_cont_labels + str_cont_labels out_labels = int_out_labels + str_out_labels network_structure, mapping = _canonicalize_network_structure( cont_labels, out_labels, network_structure) network_structure = [np.array(l) for l in network_structure] flat_connections = np.concatenate(network_structure) if out_order is None: out_order = np.sort(flat_connections[flat_connections < 0])[::-1] else: l = [] for o in out_order: l.append(mapping[o]) out_order = np.array(l) if con_order is None: #canonicalization of network structure takes care of appropriate #contraction ordering (i.e. use ASCII ordering for str and #regular ordering for int) con_order = np.unique(flat_connections[flat_connections > 0]) else: l = [] for o in con_order: l.append(mapping[o]) con_order = np.array(l) if backend not in _CACHED_JITTED_NCONS: _CACHED_JITTED_NCONS[backend] = backend_obj.jit( _jittable_ncon, static_argnums=(1, 2, 3, 4)) res_tensor = _CACHED_JITTED_NCONS[backend](_tensors, network_structure, con_order, out_order, backend_obj) if all(are_nodes): return network_components.Node(res_tensor, backend=backend_obj) return res_tensor
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[Union[Text, BaseBackend]] = 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: .. code-block:: 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: .. code-block:: 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`_. However, the use of `0` in `con_order` to denote outer products is not (currently) supported in this implementation. .. _original NCON implementation: https://arxiv.org/abs/1402.0939 Args: tensors: List of `Tensors` or `BaseNodes`. 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.backend_contextmanager.get_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 is None: backend = get_default_backend() if isinstance(backend, BaseBackend): backend_obj = backend else: backend_obj = backend_factory.get_backend(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_obj.name for n in nodes]): raise ValueError("Some nodes have backends different from '{}'".format( backend_obj.name)) _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_obj) 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 ncon( tensors: Sequence[Union[network_components.BaseNode, Tensor]], network_structure: Sequence[Sequence], con_order: Optional[Sequence] = None, out_order: Optional[Sequence] = None, backend: Optional[Union[Text, BaseBackend]] = 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: .. code-block:: 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: .. code-block:: 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`_. However, the use of `0` in `con_order` to denote outer products is not (currently) supported in this implementation. .. _original NCON implementation: https://arxiv.org/abs/1402.0939 Args: tensors: List of `Tensors` or `BaseNodes`. 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.backend_contextmanager.get_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 is None: backend = get_default_backend() if isinstance(backend, BaseBackend): backend_obj = backend else: backend_obj = backend_factory.get_backend(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_obj.name for n in nodes]): raise ValueError("Some nodes have backends different from '{}'".format( backend_obj.name)) _tensors = [] for t in tensors: if isinstance(t, network_components.BaseNode): _tensors.append(t.tensor) else: _tensors.append(t) if backend not in _CACHED_JITTED_NCONS: _CACHED_JITTED_NCONS[backend] = backend_obj.jit(_jittable_ncon, static_argnums=(1, 2, 3, 4)) res_tensor = _CACHED_JITTED_NCONS[backend](_tensors, network_structure, con_order, out_order, backend_obj) if all(are_nodes): return network_components.Node(res_tensor, backend=backend_obj) return res_tensor