Repository: google/TensorNetwork

Issue:
switch to using a single instantiation of backend objects
We're currently using a new backend object every time we create e.g. a `Node`. 
To facilitate implementation of DMRG into backends like jax or tensorflow we should
switch to instantiating only one object per backend, and then keep reusing it .


Discussion:
I'm already using this experimentally. Could submit a PR quickly 
+1 sgtm. All of the nodes should share the same backend object

Resolving patch:
diff --git a/tensornetwork/backends/backend_factory.py b/tensornetwork/backends/backend_factory.py
--- a/tensornetwork/backends/backend_factory.py
+++ b/tensornetwork/backends/backend_factory.py
@@ -29,6 +29,9 @@
     "symmetric": symmetric_backend.SymmetricBackend
 }
 
+#we instantiate each backend only once and store it here
+_INSTANTIATED_BACKENDS = dict()
+
 
 def get_backend(
     backend: Union[Text, base_backend.BaseBackend]) -> base_backend.BaseBackend:
@@ -36,4 +39,9 @@ def get_backend(
     return backend
   if backend not in _BACKENDS:
     raise ValueError("Backend '{}' does not exist".format(backend))
-  return _BACKENDS[backend]()
+
+  if backend in _INSTANTIATED_BACKENDS:
+    return _INSTANTIATED_BACKENDS[backend]
+
+  _INSTANTIATED_BACKENDS[backend] = _BACKENDS[backend]()
+  return _INSTANTIATED_BACKENDS[backend]
diff --git a/tensornetwork/block_sparse/blocksparsetensor.py b/tensornetwork/block_sparse/blocksparsetensor.py
--- a/tensornetwork/block_sparse/blocksparsetensor.py
+++ b/tensornetwork/block_sparse/blocksparsetensor.py
@@ -54,7 +54,7 @@ class ChargeArray:
   def __init__(self,
                data: np.ndarray,
                charges: List[BaseCharge],
-               flows: List[bool],
+               flows: Union[np.ndarray, List[bool]],
                order: Optional[List[List[int]]] = None,
                check_consistency: Optional[bool] = False) -> None:
     """
@@ -163,7 +163,7 @@ def flat_flows(self) -> List:
     return list(self._flows)
 
   @property
-  def flat_order(self) -> List:
+  def flat_order(self) -> np.ndarray:
     """
     The flattened `ChargeArray._oder`.
     """
@@ -191,10 +191,8 @@ def todense(self) -> np.ndarray:
     """
     return np.reshape(self.data, self.shape)
 
-  def reshape(
-      self,
-      shape: Union[List[Index], Tuple[Index, ...], List[int], Tuple[int, ...]]
-  ) -> "ChargeArray":
+  def reshape(self, shape: Union[np.ndarray, List[Index], Tuple[Index, ...],
+                                 List[int], Tuple[int, ...]]) -> "ChargeArray":
     """
     Reshape `tensor` into `shape.
     `ChargeArray.reshape` works the same as the dense 
@@ -329,10 +327,10 @@ def transpose_data(self) -> "ChargeArray":
         check_consistency=False)
     return result
 
-  def transpose(self,
-                order: Optional[Union[List[int], np.ndarray]] = np.asarray(
-                    [1, 0]),
-                shuffle: Optional[bool] = False) -> "ChargeArray":
+  def transpose(
+      self,
+      order: Union[Tuple[int, ...], List[int], np.ndarray] = np.asarray([1, 0]),
+      shuffle: Optional[bool] = False) -> "ChargeArray":
     """
     Transpose the tensor into the new order `order`. If `shuffle=False`
     no data-reshuffling is done.
@@ -418,7 +416,7 @@ class BlockSparseTensor(ChargeArray):
   def __init__(self,
                data: np.ndarray,
                charges: List[BaseCharge],
-               flows: List[bool],
+               flows: Union[np.ndarray, List[bool]],
                order: Optional[List[Union[List, np.ndarray]]] = None,
                check_consistency: Optional[bool] = False) -> None:
     """
@@ -646,9 +644,10 @@ def __truediv__(self, number: np.number) -> "BlockSparseTensor":
         check_consistency=False)
 
   # pylint: disable=arguments-differ
-  def transpose_data(self,
-                     flat_order: Optional[Union[List, np.ndarray]] = None,
-                     inplace: Optional[bool] = False) -> Any:
+  def transpose_data(
+      self,
+      flat_order: Optional[Union[Tuple, List, np.ndarray]] = None,
+      inplace: Optional[bool] = False) -> Any:
     """
     Transpose the tensor data in place such that the linear order 
     of the elements in `BlockSparseTensor.data` corresponds to the 
diff --git a/tensornetwork/block_sparse/utils.py b/tensornetwork/block_sparse/utils.py
--- a/tensornetwork/block_sparse/utils.py
+++ b/tensornetwork/block_sparse/utils.py
@@ -83,7 +83,7 @@ def fuse_stride_arrays(dims: Union[List[int], np.ndarray],
 
 
 def compute_sparse_lookup(
-    charges: List[BaseCharge], flows: List[bool],
+    charges: List[BaseCharge], flows: Union[np.ndarray, List[bool]],
     target_charges: BaseCharge) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
   """
   Compute lookup table for how dense index positions map 
@@ -153,7 +153,7 @@ def _find_best_partition(dims: Union[List[int], np.ndarray]) -> int:
 
 def compute_fused_charge_degeneracies(
     charges: List[BaseCharge],
-    flows: List[bool]) -> Tuple[BaseCharge, np.ndarray]:
+    flows: Union[np.ndarray, List[bool]]) -> Tuple[BaseCharge, np.ndarray]:
   """
   For a list of charges, computes all possible fused charges resulting
   from fusing `charges` and their respective degeneracies
@@ -193,8 +193,9 @@ def compute_fused_charge_degeneracies(
   return accumulated_charges, accumulated_degeneracies
 
 
-def compute_unique_fused_charges(charges: List[BaseCharge],
-                                 flows: List[bool]) -> BaseCharge:
+def compute_unique_fused_charges(
+    charges: List[BaseCharge],
+    flows: Union[np.ndarray, List[bool]]) -> BaseCharge:
   """
   For a list of charges, compute all possible fused charges resulting
   from fusing `charges`.
@@ -219,7 +220,8 @@ def compute_unique_fused_charges(charges: List[BaseCharge],
   return accumulated_charges
 
 
-def compute_num_nonzero(charges: List[BaseCharge], flows: List[bool]) -> int:
+def compute_num_nonzero(charges: List[BaseCharge],
+                        flows: Union[np.ndarray, List[bool]]) -> int:
   """
   Compute the number of non-zero elements, given the meta-data of 
   a symmetric tensor.
@@ -246,7 +248,7 @@ def compute_num_nonzero(charges: List[BaseCharge], flows: List[bool]) -> int:
 
 
 def reduce_charges(charges: List[BaseCharge],
-                   flows: List[bool],
+                   flows: Union[np.ndarray, List[bool]],
                    target_charges: np.ndarray,
                    return_locations: Optional[bool] = False,
                    strides: Optional[np.ndarray] = None) -> Any:
@@ -360,7 +362,7 @@ def reduce_charges(charges: List[BaseCharge],
 
 
 def _find_diagonal_sparse_blocks(
-    charges: List[BaseCharge], flows: List[bool],
+    charges: List[BaseCharge], flows: Union[np.ndarray, List[bool]],
     partition: int) -> Tuple[List, BaseCharge, np.ndarray]:
   """
   Find the location of all non-trivial symmetry blocks from the data vector of
@@ -456,9 +458,10 @@ def _find_diagonal_sparse_blocks(
 
 def _find_transposed_diagonal_sparse_blocks(
     charges: List[BaseCharge],
-    flows: List[bool],
+    flows: Union[np.ndarray, List[bool]],
     tr_partition: int,
-    order: Optional[np.ndarray] = None) -> Tuple[List, BaseCharge, np.ndarray]:
+    order: Optional[Union[List, np.ndarray]] = None
+) -> Tuple[List, BaseCharge, np.ndarray]:
   """
   Find the diagonal blocks of a transposed tensor with 
   meta-data `charges` and `flows`. `charges` and `flows` 


Test patch:
diff --git a/tensornetwork/backends/backend_test.py b/tensornetwork/backends/backend_test.py
--- a/tensornetwork/backends/backend_test.py
+++ b/tensornetwork/backends/backend_test.py
@@ -5,6 +5,7 @@
 import numpy as np
 from tensornetwork import connect, contract, Node
 from tensornetwork.backends.base_backend import BaseBackend
+from tensornetwork.backends import backend_factory
 
 
 def clean_tensornetwork_modules():
@@ -32,6 +33,8 @@ def mocked_import(name, globals, locals, fromlist, level):
     return import_orig(name, globals, locals, fromlist, level)
 
   monkeypatch.setattr(builtins, '__import__', mocked_import)
+  # Nuke the cache.
+  backend_factory._INSTANTIATED_BACKENDS = dict()
 
 
 @pytest.mark.usefixtures('no_backend_dependency')
@@ -397,4 +400,4 @@ def test_base_backend_broadcast_right_multiplication_not_implemented():
 def test_base_backend_broadcast_left_multiplication_not_implemented():
   backend = BaseBackend()
   with pytest.raises(NotImplementedError):
-    backend.broadcast_left_multiplication(np.ones((2, 2)), np.ones((2, 2)))
\ No newline at end of file
+    backend.broadcast_left_multiplication(np.ones((2, 2)), np.ones((2, 2)))
diff --git a/tensornetwork/block_sparse/blocksparsetensor_test.py b/tensornetwork/block_sparse/blocksparsetensor_test.py
--- a/tensornetwork/block_sparse/blocksparsetensor_test.py
+++ b/tensornetwork/block_sparse/blocksparsetensor_test.py
@@ -582,7 +582,8 @@ def test_matmul(dtype, num_charges, chargetype):
   tensor2 = BlockSparseTensor.random(is2, dtype=dtype)
   result = tensor1 @ tensor2
   assert result.dtype == dtype
-  dense_result = tensor1.todense() @ tensor2.todense()
+
+  dense_result = tensor1.todense() @ tensor2.todense()  #pytype: disable=unsupported-operands
   np.testing.assert_allclose(dense_result, result.todense())