feat(mpiarray): experimental, convert the local array into ndarray

- must try with MPI
radiocosmology · Feb 4, 2021 · 951d27a · 951d27a
1 parent ccc6508
commit 951d27a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/caput/mpiarray.py b/caput/mpiarray.py
@@ -1030,7 +1030,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                         distr_axis == input_.axis
                     ), "The distributed axis for all MPIArrays in an exp should be the same"
                 input_mpi.append(input_)
-                args.append(input_.view(np.ndarray))
+                args.append(input_.local_array.view(np.ndarray))
             else:
                 args.append(input_)
 
@@ -1045,7 +1045,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             out_args = []
             for output in outputs:
                 if isinstance(output, MPIArray):
-                    out_args.append(output.view(np.ndarray))
+                    out_args.append(output.local_array.view(np.ndarray))
                 else:
                     out_args.append(output)
             kwargs["out"] = tuple(out_args)