Repository: google/TensorNetwork

Issue:
Unintended lazy transposition in blocksparse 
Hi @mganahl , I have encountered some weird behavior with the transposition in blocksparse. Consider the following code:
```
ham = BlockSparseTensor.random([i0,i0,i0,i1,i1,i1]) 
hamT = ham.conj().transpose([3,4,5,0,1,2]) 
```
If I check the orders then the following is returned:
```
# ham._order == [[0],[1],[2],[3],[4],[5]]
# hamT._order == [[3],[4],[5],[0],[1],[2]]
```
As expected, this indicates that `hamT` is subject to a lazy transposition (and `ham` is not). However if I now do:
```
matrix = ham + hamT
```
Then when I recheck the order I now get:
```
# ham._order == [[3],[4],[5],[0],[1],[2]]
```
So, in when adding the two tensors together, blocksparse has done a transposition of `ham` in memory, even though `ham` was never a lazily transposed tensor. I think I understand why this is the case, but do not think that this should be allowed behavior from blocksparse. One reason is that it could result in the code transposing tensors in memory much more than is necessary. And on a conceptual level, I do not think blocksparse should be doing transpositions (in memory) on tensors where no transposition call was ever issued (i.e. the lazy transpose should delay evaluation of a `.transpose` operation until necessary, but should not decide on its own to start shifting in memory tensors where `.transpose` was never called to begin with).

I think that binary operations on tensors, e.g. `ham + hamT`, should instead ask: has any tensor been subjected to a lazy transposition that has not yet been evaluated? If yes, then do the transposition explicitly before performing the binary operation.









Discussion:
Thanks @gevenbly! That is my bad. I'll fix this such that `hamT` is transposed in memory instead of `ham`
@gevenbly there is the case though where both of the two operands have been lazily transposed. I think it is not desired in this case that both operands get transposed since one transposition is enough. I think in this we should pick one of the two and transpose only this one.

Resolving patch:
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
@@ -626,13 +626,29 @@ def _sub_add_protection(self, other):
       raise ValueError(
           "cannot add or subtract tensors non-matching sparse shapes")
 
+  def _align_storage_layout(self, other):
+    """
+    Align storage layouts of self and other.
+    """
+    self_is_ordered = np.array_equal(self.flat_order,
+                                     np.arange(len(self.flat_order)))
+    other_is_ordered = np.array_equal(other.flat_order,
+                                      np.arange(len(other.flat_order)))
+    both_unordered = (not self_is_ordered) and (not other_is_ordered)
+    if self_is_ordered and (not other_is_ordered):
+      #bring other into the same storage layout as other
+      perm = np.empty(len(other.flat_order), dtype=np.int32)
+      perm[self.flat_order] = other.flat_order
+      other.transpose_data(perm, inplace=True)
+    elif ((not self_is_ordered) and other_is_ordered) or both_unordered:
+      #bring self into the same storage layout as other
+      perm = np.empty(len(self.flat_order), dtype=np.int32)
+      perm[other.flat_order] = self.flat_order
+      self.transpose_data(perm, inplace=True)
+
   def __sub__(self, other: "BlockSparseTensor") -> "BlockSparseTensor":
-    self._sub_add_protection(other)
-    #bring self into the same storage layout as other
-    perm = np.empty(len(self.flat_order), dtype=np.int32)
-    perm[other.flat_order] = self.flat_order
-    self.transpose_data(perm, inplace=True)
-    #now subtraction is save
+    self._sub_add_protection(other)  #perform checks
+    self._align_storage_layout(other)  #align storage layout of self and other
     return BlockSparseTensor(
         data=self.data - other.data,
         charges=self._charges,
@@ -641,11 +657,8 @@ def __sub__(self, other: "BlockSparseTensor") -> "BlockSparseTensor":
         check_consistency=False)
 
   def __add__(self, other: "BlockSparseTensor") -> "BlockSparseTensor":
-    self._sub_add_protection(other)
-    #bring self into the same storage layout as other
-    perm = np.empty(len(self.flat_order), dtype=np.int32)
-    perm[other.flat_order] = self.flat_order
-    self.transpose_data(perm, inplace=True)
+    self._sub_add_protection(other)  #perform checks
+    self._align_storage_layout(other)  #align storage layout of self and other
     #now addition is save
     return BlockSparseTensor(
         data=self.data + other.data,


Test patch: