revised local shapes calculation flow

Author: tharittk

pylops_mpi/DistributedArray.py

@@ -340,7 +340,7 @@ def local_shapes(self):
         local_shapes : :obj:`list`
         """
         if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
-            return self._nccl_local_shapes(self.base_comm_nccl)
+            return self._nccl_local_shapes(False)
         else:
             return self._allgather(self.local_shape)
 
@@ -375,9 +375,8 @@ def asarray(self, masked: bool = False):
             return self.local_array
 
         if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
-            local_shapes = self._nccl_local_shapes(self.sub_comm if masked else self.base_comm_nccl)
             return nccl_asarray(self.sub_comm if masked else self.base_comm_nccl,
-                                self.local_array, local_shapes, self.axis)
+                                self.local_array, self._nccl_local_shapes(masked), self.axis)
         else:
             # Gather all the local arrays and apply concatenation.
             if masked:
@@ -550,14 +549,13 @@ def _recv(self, recv_buf=None, source=0, count=None, tag=None):
             self.base_comm.Recv(buf=recv_buf, source=source, tag=tag)
             return recv_buf
 
-    def _nccl_local_shapes(self, nccl_comm: NcclCommunicatorType):
+    def _nccl_local_shapes(self, masked: bool):
         """Get the the list of shapes of every GPU in the communicator
         """
         # gather tuple of shapes from every rank within thee communicator and copy from GPU to CPU
-        if nccl_comm == self.sub_comm:
+        if masked:
             all_tuples = self._allgather_subcomm(self.local_shape).get()
         else:
-            assert (nccl_comm == self.base_comm_nccl)
             all_tuples = self._allgather(self.local_shape).get()
         # NCCL returns the flat array that packs every tuple as 1-dimensional array
         # unpack each tuple from each rank