Clean up snapshots

xarray-contrib · Aug 2, 2024 · f9f4f0f · f9f4f0f
1 parent 7c2c1c2
commit f9f4f0f
Show file tree

Hide file tree

Showing 2 changed files with 18,104 additions and 14,905 deletions.
diff --git a/flox/core.py b/flox/core.py
@@ -394,7 +394,9 @@ def find_group_cohorts(
         chunks_per_label = chunks_per_label[present_labels_mask]
 
     label_chunks = {
-        present_labels[idx]: bitmask.indices[slice(bitmask.indptr[idx], bitmask.indptr[idx + 1])]
+        present_labels[idx].item(): bitmask.indices[
+            slice(bitmask.indptr[idx], bitmask.indptr[idx + 1])
+        ]
         for idx in range(bitmask.shape[LABEL_AXIS])
     }
 
@@ -495,11 +497,11 @@ def invert(x) -> tuple[np.ndarray, ...]:
             slice(containment.indptr[rowidx], containment.indptr[rowidx + 1])
         ]
         cohort_ = present_labels[cohidx]
-        cohort = [elem for elem in cohort_ if elem not in merged_keys]
+        cohort = [elem.item() for elem in cohort_ if elem not in merged_keys]
         if not cohort:
             continue
         merged_keys.update(cohort)
-        allchunks = (label_chunks[member] for member in cohort)
+        allchunks = (label_chunks[member].tolist() for member in cohort)
         chunk = tuple(set(itertools.chain(*allchunks)))
         merged_cohorts[chunk] = cohort