update in slicing history shape

danlkv · Dec 1, 2023 · 5f5a15c · 5f5a15c
1 parent 957ddb5
commit 5f5a15c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/qtensor/contraction_backends/common.py b/qtensor/contraction_backends/common.py
@@ -44,8 +44,11 @@ def slice_numpy_tensor(data:np.ndarray, indices_in, indices_out, slice_dict):
         i for sl, i in zip(slice_bounds, indices_in) if not isinstance(sl, int)
     ]
     indices_sized = [v.copy(size=size) for v, size in zip(indices_sliced, s_data.shape)]
+    #print("indices_sized", indices_sized)
+    #print("Slice bounds", slice_bounds)
+    #print("Slice dict", slice_dict)
+    #print("data shape, sliced data shape", data.shape, s_data.shape)
     indices_out = [v for v in indices_out if not isinstance(slice_dict.get(v, None), int)]
-    assert len(indices_sized) == len(s_data.shape)
     assert len(indices_sliced) == len(s_data.shape)
     st_data = permute_np_tensor_data(s_data, indices_sliced, indices_out)
     return st_data, indices_out

diff --git a/qtensor/optimisation/Optimizer.py b/qtensor/optimisation/Optimizer.py
@@ -318,7 +318,7 @@ def _split_graph(self, p_graph, max_tw):
         peo_ints = self.peo_ints
         tw = self.treewidth
         self._slice_hist = []
-        self._slice_hist.append([0, tw])
+        self._slice_hist.append([0, tw, peo_ints])
         log.info('Treewidth: {}', tw)
         log.info('Target treewidth: {}', max_tw)
         result = []
@@ -354,7 +354,7 @@ def _split_graph(self, p_graph, max_tw):
 
             peo_ints, path = self._update_peo_after_slice(p_graph, result)
             tw = max(path)
-            self._slice_hist.append([pv_cnt, tw])
+            self._slice_hist.append([pv_cnt, tw, peo_ints])
             delta = tw - max_tw
 
         return peo_ints, result