minor fixes based on recently merged PR

Author: tharittk

pylops_mpi/DistributedArray.py

@@ -380,7 +380,7 @@ def asarray(self, masked: bool = False):
             return self.local_array
 
         if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
-            return nccl_asarray(self.sub_comm if masked else self.base_comm,
+            return nccl_asarray(self.sub_comm if masked else self.base_comm_nccl,
                                 self.local_array, self.local_shapes, self.axis)
         else:
             # Gather all the local arrays and apply concatenation.
@@ -640,9 +640,9 @@ def dot(self, dist_array):
         self._check_mask(dist_array)
         ncp = get_module(self.engine)
         # Convert to Partition.SCATTER if Partition.BROADCAST
-        x = DistributedArray.to_dist(x=self.local_array, base_comm_nccl=self.base_comm_nccl) \
+        x = DistributedArray.to_dist(x=self.local_array, base_comm=self.base_comm, base_comm_nccl=self.base_comm_nccl) \
             if self.partition in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST] else self
-        y = DistributedArray.to_dist(x=dist_array.local_array, base_comm_nccl=self.base_comm_nccl) \
+        y = DistributedArray.to_dist(x=dist_array.local_array, base_comm=self.base_comm, base_comm_nccl=self.base_comm_nccl) \
             if self.partition in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST] else dist_array
         # Flatten the local arrays and calculate dot product
         return self._allreduce_subcomm(ncp.dot(x.local_array.flatten(), y.local_array.flatten()))
@@ -716,7 +716,7 @@ def norm(self, ord: Optional[int] = None,
             Axis along which vector norm needs to be computed. Defaults to ``-1``
         """
         # Convert to Partition.SCATTER if Partition.BROADCAST
-        x = DistributedArray.to_dist(x=self.local_array, base_comm_nccl=self.base_comm_nccl) \
+        x = DistributedArray.to_dist(x=self.local_array, base_comm=self.base_comm, base_comm_nccl=self.base_comm_nccl) \
             if self.partition in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST] else self
         if axis == -1:
             # Flatten the local arrays and calculate norm

tutorials/poststack_nccl.py

@@ -80,7 +80,7 @@
 BDiag = pylops_mpi.basicoperators.MPIBlockDiag(ops=[Top.H @ PPop @ Top, ])
 
 # This computation will be done in GPU. The call asarray() trigger the NCCL communication (gather result from each GPU).
-# But array `d` and `d_0` still lives in GPU memory
+# But array `d` and `d_0` still live in GPU memory
 d_dist = BDiag @ m3d_dist
 d_local = d_dist.local_array.reshape((ny_i, nx, nz))
 d = d_dist.asarray().reshape((ny, nx, nz))
@@ -89,8 +89,8 @@
 
 # ###############################################################################
 
-# Inversion using CGLS solver - There is no code change to have run on NCCL
-# In this particular case, the local computation will be done in GPU. And the collective communication calls
+# Inversion using CGLS solver - There is no code change to have run on NCCL (it handles though MPI operator and DistributedArray)
+# In this particular case, the local computation will be done in GPU. Collective communication calls
 # will be carried through NCCL GPU-to-GPU.
 minv3d_iter_dist = pylops_mpi.optimization.basic.cgls(BDiag, d_dist, x0=mback3d_dist, niter=1, show=True)[0]
 minv3d_iter = minv3d_iter_dist.asarray().reshape((ny, nx, nz))