BugFix: Got rid of rechunking in distributed loading

hyperspy · Apr 6, 2023 · f16dacf · f16dacf
1 parent a1c8506
commit f16dacf
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/rsciio/utils/tools.py b/rsciio/utils/tools.py
@@ -227,14 +227,14 @@ def get_chunk_index(
     """
     nav_shape = np.delete(np.array(shape), signal_axes)
     num_frames = np.prod(nav_shape)
-    indexes = da.arange(num_frames)
-    indexes = da.reshape(indexes, nav_shape)
+    indexes = np.arange(num_frames)
+    indexes = np.reshape(indexes, nav_shape)
     chunks = da.core.normalize_chunks(
         chunks=chunks, shape=shape, limit=block_size_limit, dtype=dtype
     )
 
     nav_chunks = tuple(np.delete(np.array(chunks, dtype=object), signal_axes))
-    indexes = da.rechunk(indexes, chunks=nav_chunks)
+    indexes = da.from_array(indexes, chunks=nav_chunks)
     return indexes