Repository: google/TensorNetwork

Issue:
add switch_backend to BaseMPS



Discussion:
I'd like to take a look, good opportunity to get acquainted with latest development on MPS class.

Resolving patch:
diff --git a/tensornetwork/__init__.py b/tensornetwork/__init__.py
--- a/tensornetwork/__init__.py
+++ b/tensornetwork/__init__.py
@@ -1,7 +1,7 @@
 #pylint: disable=line-too-long
 from tensornetwork.network_components import Node, Edge, CopyNode, BaseNode, NodeCollection
 #pylint: disable=line-too-long
-from tensornetwork.network_operations import norm, conj, copy, transpose, split_node, split_node_qr, split_node_rq, split_node_full_svd, reachable, check_connected, check_correct, get_all_nodes, get_all_edges, remove_node, contract_trace_edges, get_subgraph_dangling, reduced_density
+from tensornetwork.network_operations import norm, conj, copy, transpose, split_node, split_node_qr, split_node_rq, split_node_full_svd, reachable, check_connected, check_correct, get_all_nodes, get_all_edges, remove_node, contract_trace_edges, get_subgraph_dangling, reduced_density, switch_backend
 #pylint: disable=line-too-long
 from tensornetwork.network_components import contract, contract_copy_node, contract_between, outer_product, outer_product_final_nodes, contract_parallel, flatten_edges, split_edge, get_all_nondangling, get_all_dangling, flatten_all_edges, flatten_edges_between, get_parallel_edges, get_shared_edges
 from tensornetwork.backends.base_backend import BaseBackend
diff --git a/tensornetwork/matrixproductstates/base_mps.py b/tensornetwork/matrixproductstates/base_mps.py
--- a/tensornetwork/matrixproductstates/base_mps.py
+++ b/tensornetwork/matrixproductstates/base_mps.py
@@ -16,11 +16,10 @@
 from __future__ import division
 from __future__ import print_function
 import numpy as np
-import functools
 # pylint: disable=line-too-long
 from tensornetwork.network_components import Node, contract, contract_between, BaseNode
 # pylint: disable=line-too-long
-from tensornetwork.network_operations import split_node_qr, split_node_rq, split_node_full_svd, norm, conj
+from tensornetwork.network_operations import split_node_qr, split_node_rq, split_node_full_svd, norm, conj, switch_backend
 from typing import Any, List, Optional, Text, Type, Union, Dict, Sequence
 from tensornetwork.backends.base_backend import BaseBackend
 Tensor = Any
@@ -199,6 +198,15 @@ def physical_dimensions(self) -> List:
 
     return [node.shape[1] for node in self.nodes]
 
+  def switch_backend(self, new_backend: Text) -> None:
+    """
+    Change the backend of all the nodes in the MPS.
+
+    Args:
+      new_backend (str): The new backend.
+    """
+    switch_backend(self.nodes, new_backend)
+
   def right_envs(self, sites: Sequence[int]) -> Dict:
     raise NotImplementedError()
 
diff --git a/tensornetwork/network_operations.py b/tensornetwork/network_operations.py
--- a/tensornetwork/network_operations.py
+++ b/tensornetwork/network_operations.py
@@ -789,3 +789,4 @@ def switch_backend(nodes: Iterable[BaseNode], new_backend: Text) -> None:
                                 "backend is 'numpy'. Current backend "
                                 "is '{}'".format(node.backend))
     node.tensor = backend.convert_to_tensor(node.tensor)
+    node.backend = backend


Test patch:
diff --git a/tensornetwork/matrixproductstates/base_mps_test.py b/tensornetwork/matrixproductstates/base_mps_test.py
--- a/tensornetwork/matrixproductstates/base_mps_test.py
+++ b/tensornetwork/matrixproductstates/base_mps_test.py
@@ -142,3 +142,13 @@ def test_apply_two_site_gate(backend_dtype_values):
   res = tn.contract_between(mps.nodes[5], mps.nodes[6])
   res.reorder_edges(order)
   np.testing.assert_allclose(res.tensor, actual)
+
+
+def test_mps_switch_backend(backend):
+  D, d, N = 10, 2, 10
+  tensors = [get_random_np((1, d, D), np.float64)] + [
+      get_random_np((D, d, D), np.float64) for _ in range(N - 2)
+  ] + [get_random_np((D, d, 1), np.float64)]
+  mps = BaseMPS(tensors, center_position=0, backend="numpy")
+  mps.switch_backend(backend)
+  assert mps.backend.name == backend
diff --git a/tensornetwork/tests/network_operations_test.py b/tensornetwork/tests/network_operations_test.py
--- a/tensornetwork/tests/network_operations_test.py
+++ b/tensornetwork/tests/network_operations_test.py
@@ -15,10 +15,6 @@
 import tensornetwork as tn
 import pytest
 import numpy as np
-import tensorflow as tf
-import torch
-import jax
-from jax.config import config
 
 
 def test_split_node_full_svd_names(backend):
@@ -329,3 +325,12 @@ def test_reduced_density_contraction(backend):
   tn.reduced_density([a[0]])
   result = tn.contractors.greedy(tn.reachable(a), ignore_edge_order=True)
   np.testing.assert_allclose(result.tensor, np.eye(2))
+
+
+def test_switch_backend(backend):
+  a = tn.Node(np.random.rand(3, 3, 3), name="A", backend="numpy")
+  b = tn.Node(np.random.rand(3, 3, 3), name="B", backend="numpy")
+  c = tn.Node(np.random.rand(3, 3, 3), name="C", backend="numpy")
+  nodes = [a, b, c]
+  tn.switch_backend(nodes, backend)
+  assert nodes[0].backend.name == backend