Final PR for block sparsity (#506)

* remove line

* add docstring

* add svd

* fix bug in svd and diag

bug appeared when tensors where empty

* fixed empyt-array bug in qrr

* added tests

* fix bug for empty arrays

* sone typing stuff

* tests

* uncomment tests

* added eig, sqrt, trace

* add routines, fix some typing issues

* formatting of exception

* linting

* added pinv, ones, rannd, randn, zeros

* fix typing issues

* add pinv

remove conj and T from subclassed BlockSparseTensor again

* added conj and transpose back

typing with __new__ is a pain ...

* some more tests

* fix transpose arguments

add docstring to conj

* more complete docstring for diag

* typo

* space

* docstring

* fix return type in docstring

* improve docstring of outerproduct

* set numpy seed

* disable annoying linting error

* dito

* ditto

* add dtpye fixture to test_tensordot_reshape

* fake line

* remove fake line

* numpy>=1.17

* added decompositions

* add files

* rename rand to random

* rand -> random, fix bug

* add `relative` argument to svd_decomposition

* add `relative` argument to svd_decomposition

* add `relative`

* remove ,

* fix svd_decomposition

* added 'symmetric' backend

* adding symmetric backend

* remove comment

* added init

* fix __init__

* test tensordot raises

* more tests

* added test files

* added symmetric backend support

* added sparse_shape

* fixing the last pieces

* added sparse_shape to TestNode definition

* final touches

* change size dtypes to np.unit64

* changed size dtypes to np.int64

* add size_t to change index-array size types

* capitalize size_t

* newline

* typo-bug

* add unit tests

* added tests

* fix dtype issue (again)

* fix edge case of empty arrays

* add tests

* more comprehensive testing

* fix tests

* fix bug in _find_diagonal_sparse_blocks for empty tensors

* remove print statements

* fix edge case of empty charges

* fix tests

* reduce_charges: edge case of empty charge

* linting

* remove comment

* linting

* fix edge case of empty input in BaseCharge.__init__

* raise if input to __eq__ is empty

* fix edge case of empty charges in __add__

* fix edge case of empty charges in `isin`

* fix bug in cahrge_test.py

* fix pytest warning

* fix warning

* more comprehensive testing

* test svd for empty tensors

* more comprehensive testing

* more comprehensive testing

* more comprehensive testing of edge cases

* edge case tests for eig

* shorten code

* edge case fixed

* fix bug _find_transposed_diagonal_sparse_blocks

* fix edge case in diag

* testing edge cases of `diag`

* raise error when reshaping empty tensors

* fix tests

* fix reshape

* linting

* better testing

* linting

* more comprehensive testing

* remove redundant test

* more comprehensive testing

* linting

* remove lines

* shorten import code

* change import

* fix bug in tensordot for inner product mode

add outerproduct mode to tensordot

* fix tests

* more comprehensive testing

* return np.ndarray for tensordot in inner mode

* fix import

* fix import

* fix import

* fix bug

* fix warning

* more checks for input to tensordot

* conform to numpy convention when using tensordot for inner product on empty tensors

* fix tests

* remove print

* fix outerproduct in tensordot

* raise ValueError in pinv

* add raise-tests

* fix typo

* comment

* added two more checks to tensordot

* more tests

* fix return type of tensordot

* restructuring files

* added files

* remove unused type

* remove unused types

* remove unused types

* remove unused types

* import outerproduct

* remove unused imports

* fix imports

* fix `relative` case

* formatting

* import norm

* fix imports

* remove unneeded import

* add TODO

* fix import

* tensordot returns BlockSparseTensor

for inner product.

* fix test

* fix imports in tests

* fix issue with divison by 0

for empty tensors

* remove comment

* fix typo

* add LABEL_SIZE_T variable

to set the size type of charge_labels

* remove hard coded dtypes for charges

* fix overflow bug when using > 3 symmetries

* remove comment

* remove comment

* rename variable

* remove sparse_shape

* remove sparse_shape

* more comprehensive testing

* fix tests

* remove unused functions

* extend tests to include more than 1 symmetry

* linting

* more tests

* typo in docstring

* adding broadcast_right_multiplication to backends

* fix test

* bradcasting support added

* remove backend specific code

Co-authored-by: Chase Roberts <chaseriley@google.com>

Author: mganahl

tensornetwork/__init__.py

@@ -14,6 +14,7 @@
 from tensornetwork.matrixproductstates.finite_mps import FiniteMPS
 from tensornetwork.matrixproductstates.infinite_mps import InfiniteMPS
 from tensornetwork.backend_contextmanager import DefaultBackend, set_default_backend
+from tensornetwork import block_sparse
 from tensornetwork.block_sparse.blocksparsetensor import BlockSparseTensor, ChargeArray
 from tensornetwork.block_sparse.index import Index
 from tensornetwork.block_sparse.charge import U1Charge, BaseCharge

tensornetwork/block_sparse/blocksparsetensor.py

@@ -766,7 +766,7 @@ def outerproduct(tensor1: BlockSparseTensor,
 def tensordot(tensor1: BlockSparseTensor,
               tensor2: BlockSparseTensor,
               axes: Optional[Union[Sequence[Sequence[int]], int]] = 2
-             ) -> Union[BlockSparseTensor, np.number]:
+             ) -> BlockSparseTensor:
   """
   Contract two `BlockSparseTensor`s along `axes`.
   Args:
@@ -856,9 +856,20 @@ def tensordot(tensor1: BlockSparseTensor,
   if (len(axes1) == tensor1.ndim) and (len(axes2) == tensor2.ndim):
     t1 = tensor1.transpose(axes1).transpose_data()
     t2 = tensor2.transpose(axes2).transpose_data()
-    #NOTE (mganahl): for t1.data=[] and t2.data=[] this returns 0.0,
-    #is consistent with numpy behaviour.
-    return np.dot(t1.data, t2.data)
+    data = np.dot(t1.data, t2.data)
+    charge = tensor1._charges[0]
+    final_charge = charge.__new__(type(charge))
+
+    final_charge.__init__(
+        np.empty((charge.num_symmetries, 0), dtype=np.int16),
+        charge_labels=np.empty(0, dtype=np.int16),
+        charge_types=charge.charge_types)
+    return BlockSparseTensor(
+        data=data,
+        charges=[final_charge],
+        flows=[False],
+        order=[[0]],
+        check_consistency=False)
 
   #in all other cases we perform a regular tensordot
   free_axes1 = sorted(set(np.arange(tensor1.ndim)) - set(axes1))

tensornetwork/block_sparse/charge.py

@@ -59,7 +59,8 @@ def __next__(self):
   def __init__(self,
                charges: np.ndarray,
                charge_labels: Optional[np.ndarray] = None,
-               charge_types: Optional[List[Type["BaseCharge"]]] = None) -> None:
+               charge_types: Optional[List[Type["BaseCharge"]]] = None,
+               charge_dtype: Optional[Type[np.number]] = np.int16) -> None:
     if charges.ndim == 1:
       charges = charges[None, :]
 
@@ -68,20 +69,25 @@ def __init__(self,
           "`len(charge_types) = {}` does not match `charges.shape[0]={}`"
           .format(len(charge_types), charges.shape[0]))
 
+    if charges.shape[0] <= 3:
+      label_dtype = np.int16
+    else:
+      label_dtype = np.int32
     self.charge_types = charge_types
     if charge_labels is None:
       if charges.shape[1] > 0:
         self.unique_charges, self.charge_labels = np.unique(
-            charges.astype(np.int16), return_inverse=True, axis=1)
-        self.charge_labels = self.charge_labels.astype(np.int16)
+            charges.astype(charge_dtype), return_inverse=True, axis=1)
+        self.charge_labels = self.charge_labels.astype(label_dtype)
       else:
-        self.unique_charges = np.empty((charges.shape[0], 0), dtype=np.int16)
-        self.charge_labels = np.empty(0, dtype=np.int16)
+        self.unique_charges = np.empty((charges.shape[0], 0),
+                                       dtype=charge_dtype)
+        self.charge_labels = np.empty(0, dtype=label_dtype)
     else:
-      self.charge_labels = np.asarray(charge_labels, dtype=np.int16)
+      self.charge_labels = np.asarray(charge_labels, dtype=label_dtype)
 
-      self.unique_charges = charges.astype(np.int16)
-      self.charge_labels = charge_labels.astype(np.int16)
+      self.unique_charges = charges.astype(charge_dtype)
+      self.charge_labels = charge_labels.astype(label_dtype)
 
   @staticmethod
   def fuse(charge1, charge2):
@@ -142,10 +148,15 @@ def charges(self):
   def dtype(self):
     return self.unique_charges.dtype
 
+  @property
+  def label_dtype(self):
+    return self.charge_labels.dtype
+
   @property
   def degeneracies(self):
     exp1 = self.charge_labels[:, None]
-    exp2 = np.arange(self.unique_charges.shape[1], dtype=np.int16)[None, :]
+    exp2 = np.arange(
+        self.unique_charges.shape[1], dtype=self.label_dtype)[None, :]
     return np.sum(exp1 == exp2, axis=0)
 
   def __repr__(self):
@@ -194,8 +205,8 @@ def identity_charges(self) -> "BaseCharge":
     """
     unique_charges = np.asarray(
         [ct.identity_charge() for ct in self.charge_types],
-        dtype=np.int16)[:, None]
-    charge_labels = np.zeros(1, dtype=np.int16)
+        dtype=self.dtype)[:, None]
+    charge_labels = np.zeros(1, dtype=self.label_dtype)
     obj = self.__new__(type(self))
     obj.__init__(unique_charges, charge_labels, self.charge_types)
     return obj
@@ -217,19 +228,19 @@ def __add__(self, other: "BaseCharge") -> "BaseCharge":
         other.charge_labels) == 0):
       obj = self.__new__(type(self))
       obj.__init__(
-          np.empty((self.num_symmetries, 0), dtype=np.int16),
-          np.empty(0, dtype=np.int16), self.charge_types)
+          np.empty((self.num_symmetries, 0), dtype=self.dtype),
+          np.empty(0, dtype=self.label_dtype), self.charge_types)
       return obj
     unique_charges, charge_labels = np.unique(
         comb_charges, return_inverse=True, axis=1)
     charge_labels = charge_labels.reshape(self.unique_charges.shape[1],
                                           other.unique_charges.shape[1]).astype(
-                                              np.int16)
+                                              self.label_dtype)
     # find new labels using broadcasting
     left_labels = self.charge_labels[:, None] + np.zeros([1, len(other)],
-                                                         dtype=np.int16)
-    right_labels = other.charge_labels[None, :] + np.zeros([len(self), 1],
-                                                           dtype=np.int16)
+                                                         dtype=self.label_dtype)
+    right_labels = other.charge_labels[None, :] + np.zeros(
+        [len(self), 1], dtype=self.label_dtype)
     charge_labels = charge_labels[np.ravel(left_labels), np.ravel(right_labels)]
 
     obj = self.__new__(type(self))
@@ -392,8 +403,9 @@ def unique(self,
       unique_charges = tmp_charge.unique_charges[:, tmp]
     obj.__init__(
         charges=unique_charges,
-        charge_labels=np.arange(unique_charges.shape[1], dtype=np.int16),
-        charge_types=tmp_charge.charge_types)
+        charge_labels=np.arange(
+            unique_charges.shape[1], dtype=self.label_dtype),
+        charge_types=self.charge_types)
     out = [obj]
     if return_index or return_inverse or return_counts:
       for n in range(1, len(tmp)):
@@ -421,23 +433,24 @@ def reduce(self,
         of target values.
     """
     if isinstance(target_charges, (np.integer, int)):
-      target_charges = np.asarray([target_charges], dtype=np.int16)
+      target_charges = np.asarray([target_charges], dtype=self.dtype)
     if target_charges.ndim == 1:
       target_charges = target_charges[None, :]
-    target_charges = np.asarray(target_charges, dtype=np.int16)
+    target_charges = np.asarray(target_charges, dtype=self.dtype)
     # find intersection of index charges and target charges
     reduced_charges, label_to_unique, _ = intersect(
         self.unique_charges, target_charges, axis=1, return_indices=True)
     num_unique = len(label_to_unique)
 
     # construct the map to the reduced charges
-    map_to_reduced = np.full(self.dim, fill_value=-1, dtype=np.int16)
-    map_to_reduced[label_to_unique] = np.arange(num_unique, dtype=np.int16)
+    map_to_reduced = np.full(self.dim, fill_value=-1, dtype=self.label_dtype)
+    map_to_reduced[label_to_unique] = np.arange(
+        num_unique, dtype=self.label_dtype)
 
     # construct the map to the reduced charges
     reduced_ind_labels = map_to_reduced[self.charge_labels]
     reduced_locs = reduced_ind_labels >= 0
-    new_ind_labels = reduced_ind_labels[reduced_locs].astype(np.int16)
+    new_ind_labels = reduced_ind_labels[reduced_locs].astype(self.label_dtype)
     obj = self.__new__(type(self))
     obj.__init__(reduced_charges, new_ind_labels, self.charge_types)
 
@@ -517,8 +530,13 @@ class U1Charge(BaseCharge):
   def __init__(self,
                charges: np.ndarray,
                charge_labels: Optional[np.ndarray] = None,
-               charge_types: Optional[List[Type["BaseCharge"]]] = None) -> None:
-    super().__init__(charges, charge_labels, charge_types=[type(self)])
+               charge_types: Optional[List[Type["BaseCharge"]]] = None,
+               charge_dtype: Optional[Type[np.number]] = np.int16) -> None:
+    super().__init__(
+        charges,
+        charge_labels,
+        charge_types=[type(self)],
+        charge_dtype=charge_dtype)
 
   @staticmethod
   def fuse(charge1, charge2) -> np.ndarray:

tensornetwork/block_sparse/linalg_test.py

@@ -472,7 +472,7 @@ def test_trace_matrix(dtype, num_charges, D):
                                     dtype=dtype)
   res = trace(matrix)
   res_dense = np.trace(matrix.todense())
-  np.testing.assert_allclose(res, res_dense)
+  np.testing.assert_allclose(res.data, res_dense)
 
 
 @pytest.mark.parametrize("dtype", [np.float64, np.complex128])

tensornetwork/block_sparse/utils.py

@@ -107,7 +107,9 @@ def compute_sparse_lookup(
       return_inverse=True, sort=False)
   _, label_to_unique, _ = unique_charges.intersect(
       target_charges, return_indices=True)
-  tmp = np.full(len(unique_charges), fill_value=-1, dtype=np.int16)
+  tmp = np.full(
+      len(unique_charges), fill_value=-1, dtype=charges[0].label_dtype)
+
   tmp[label_to_unique] = label_to_unique
   lookup = tmp[inverse]
   lookup = lookup[lookup >= 0]
@@ -297,8 +299,8 @@ def reduce_charges(charges: List[BaseCharge],
       right_ind.charge_labels) == 0):
     obj = charges[0].__new__(type(charges[0]))
     obj.__init__(
-        np.empty((charges[0].num_symmetries, 0), dtype=np.int16),
-        np.empty(0, dtype=np.int16), charges[0].charge_types)
+        np.empty((charges[0].num_symmetries, 0), dtype=charges[0].dtype),
+        np.empty(0, dtype=charges[0].label_dtype), charges[0].charge_types)
     if return_locations:
       return obj, np.empty(0, dtype=SIZE_T)
     return obj
@@ -310,7 +312,7 @@ def reduce_charges(charges: List[BaseCharge],
   # intersect combined qnums and target_charges
   reduced_qnums, label_to_unique, _ = intersect(
       unique_comb_qnums, target_charges, axis=1, return_indices=True)
-  map_to_kept = -np.ones(num_unique, dtype=np.int16)
+  map_to_kept = -np.ones(num_unique, dtype=charges[0].label_dtype)
   map_to_kept[label_to_unique] = np.arange(len(label_to_unique))
   #new_comb_labels is a matrix of shape (left_ind.num_unique, right_ind.num_unique)
   #each row new_comb_labels[n,:] contains integers values. Positions where values > 0
@@ -386,14 +388,14 @@ def _find_diagonal_sparse_blocks(
     # special cases (matrix of trivial height or width)
     num_nonzero = compute_num_nonzero(charges, flows)
     block_maps = [np.arange(0, num_nonzero, dtype=SIZE_T).ravel()]
-    block_qnums = np.zeros([num_syms, 1], dtype=np.int16)
+    block_qnums = np.zeros([num_syms, 1], dtype=charges[0].dtype)
     block_dims = np.array([[1], [num_nonzero]])
 
     if partition == len(flows):
       block_dims = np.flipud(block_dims)
 
     obj = charges[0].__new__(type(charges[0]))
-    obj.__init__(block_qnums, np.arange(0, dtype=np.int16),
+    obj.__init__(block_qnums, np.arange(0, dtype=charges[0].label_dtype),
                  charges[0].charge_types)
 
     return block_maps, obj, block_dims
@@ -412,8 +414,8 @@ def _find_diagonal_sparse_blocks(
   if num_blocks == 0:
     obj = charges[0].__new__(type(charges[0]))
     obj.__init__(
-        np.zeros((charges[0].num_symmetries, 0), dtype=np.int16),
-        np.arange(0, dtype=np.int16), charges[0].charge_types)
+        np.zeros((charges[0].num_symmetries, 0), dtype=charges[0].dtype),
+        np.arange(0, dtype=charges[0].label_dtype), charges[0].charge_types)
 
     return [], obj, np.empty((2, 0), dtype=SIZE_T)
 
@@ -445,7 +447,8 @@ def _find_diagonal_sparse_blocks(
                np.arange(block_dims[1, n])[None, :]) for n in range(num_blocks)
   ]
   obj = charges[0].__new__(type(charges[0]))
-  obj.__init__(block_qnums, np.arange(block_qnums.shape[1], dtype=np.int16),
+  obj.__init__(block_qnums,
+               np.arange(block_qnums.shape[1], dtype=charges[0].label_dtype),
                charges[0].charge_types)
 
   return block_maps, obj, block_dims
@@ -508,8 +511,8 @@ def _find_transposed_diagonal_sparse_blocks(
     # special case: trivial number of non-zero elements
     obj = charges[0].__new__(type(charges[0]))
     obj.__init__(
-        np.empty((charges[0].num_symmetries, 0), dtype=np.int16),
-        np.arange(0, dtype=np.int16), charges[0].charge_types)
+        np.empty((charges[0].num_symmetries, 0), dtype=charges[0].dtype),
+        np.arange(0, dtype=charges[0].label_dtype), charges[0].charge_types)
 
     return [], obj, np.empty((2, 0), dtype=SIZE_T)
 
@@ -518,8 +521,9 @@ def _find_transposed_diagonal_sparse_blocks(
                               np.logical_not(flows[orig_partition:]))
 
   inv_row_map = -np.ones(
-      orig_unique_row_qnums.unique_charges.shape[1], dtype=np.int16)
-  inv_row_map[row_map] = np.arange(len(row_map), dtype=np.int16)
+      orig_unique_row_qnums.unique_charges.shape[1],
+      dtype=charges[0].label_dtype)
+  inv_row_map[row_map] = np.arange(len(row_map), dtype=charges[0].label_dtype)
 
   all_degens = np.append(orig_col_degen[col_map],
                          0)[inv_row_map[orig_row_ind.charge_labels]]
@@ -564,7 +568,8 @@ def _find_transposed_diagonal_sparse_blocks(
     block_maps = [(all_cumul_degens[orig_row_posR] +
                    dense_to_sparse[orig_col_posR]).ravel()]
     obj = charges[0].__new__(type(charges[0]))
-    obj.__init__(block_qnums, np.arange(block_qnums.shape[1], dtype=np.int16),
+    obj.__init__(block_qnums,
+                 np.arange(block_qnums.shape[1], dtype=charges[0].label_dtype),
                  charges[0].charge_types)
 
   elif tr_partition == len(charges):
@@ -595,7 +600,8 @@ def _find_transposed_diagonal_sparse_blocks(
     block_maps = [(all_cumul_degens[orig_row_posL] +
                    dense_to_sparse[orig_col_posL]).ravel()]
     obj = charges[0].__new__(type(charges[0]))
-    obj.__init__(block_qnums, np.arange(block_qnums.shape[1], dtype=np.int16),
+    obj.__init__(block_qnums,
+                 np.arange(block_qnums.shape[1], dtype=charges[0].label_dtype),
                  charges[0].charge_types)
   else:
 
@@ -639,7 +645,8 @@ def _find_transposed_diagonal_sparse_blocks(
           all_cumul_degens[np.add.outer(orig_row_posL, orig_row_posR)] +
           dense_to_sparse[np.add.outer(orig_col_posL, orig_col_posR)]).ravel()
     obj = charges[0].__new__(type(charges[0]))
-    obj.__init__(block_qnums, np.arange(block_qnums.shape[1], dtype=np.int16),
+    obj.__init__(block_qnums,
+                 np.arange(block_qnums.shape[1], dtype=charges[0].label_dtype),
                  charges[0].charge_types)
 
   return block_maps, obj, block_dims