fix: only perform a view when the distributed axis has changed

radiocosmology · Feb 23, 2021 · 6f14021 · 6f14021
1 parent 7d0aff8
commit 6f14021
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/caput/mpiarray.py b/caput/mpiarray.py
@@ -220,11 +220,15 @@ def __getitem__(self, slobj):
                 index for index, sl in enumerate(slobj) if not isinstance(sl, int)
             ].index(self.axis)
 
-            return MPIArray.wrap(
-                self.array.view(np.ndarray)[slobj],
-                axis=dist_axis,
-                comm=self.array._comm,
-            )
+            if dist_axis != self.axis:
+
+                return MPIArray.wrap(
+                    self.array.view(np.ndarray)[slobj],
+                    axis=dist_axis,
+                    comm=self.array._comm,
+                )
+            else:
+                return MPIArray.wrap(self.array[slobj], axis=dist_axis, comm=self.array._comm)
 
     def __setitem__(self, slobj, value):