minor: fix flake8

Author: mrava87SW

pylops_mpi/Distributed.py

@@ -265,17 +265,45 @@ def _send(self,
                      send_buf, dest, count, tag=tag,
                      engine=engine)
 
-    def _recv(self, recv_buf=None, source=0, count=None, tag=0):
+    def _recv(self,
+              base_comm: MPI.Comm,
+              base_comm_nccl: NcclCommunicatorType,
+              recv_buf=None, source=0, count=None, tag=0,
+              engine: str = "numpy",
+              ) -> NDArray:
         """Receive operation
+
+        Parameters
+        ----------
+        base_comm : :obj:`MPI.Comm`
+            Base MPI Communicator.
+        base_comm_nccl : :obj:`cupy.cuda.nccl.NcclCommunicator`
+            NCCL Communicator.
+        recv_buf : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`, optional
+            The buffered array to receive data.
+        source : :obj:`int`
+            The rank of the sending CPU/GPU device.
+        count : :obj:`int`
+            Number of elements to receive.
+        tag : :obj:`int`
+            Tag of the message to be sent.
+        engine : :obj:`str`, optional
+            Engine used to store array (``numpy`` or ``cupy``)
+
+        Returns
+        -------
+        recv_buf : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`
+            The buffer containing the received data.
+
         """
-        if deps.nccl_enabled and self.base_comm_nccl:
+        if deps.nccl_enabled and base_comm_nccl is not None:
             if recv_buf is None:
                 raise ValueError("recv_buf must be supplied when using NCCL")
             if count is None:
                 count = recv_buf.size
-            nccl_recv(self.base_comm_nccl, recv_buf, source, count)
+            nccl_recv(base_comm_nccl, recv_buf, source, count)
             return recv_buf
         else:
-            return mpi_recv(self.base_comm,
+            return mpi_recv(base_comm,
                             recv_buf, source, count, tag=tag,
-                            engine=self.engine)
+                            engine=engine)

pylops_mpi/DistributedArray.py

@@ -797,7 +797,9 @@ def add_ghost_cells(self, cells_front: Optional[int] = None,
                 # Transfer of ghost cells can be skipped if len(recv_buf) = 0
                 # Additionally, NCCL will hang if the buffer size is 0 so this optimization is somewhat mandatory
                 if len(recv_buf) != 0:
-                    ghosted_array = ncp.concatenate([self._recv(recv_buf, source=self.rank - 1, tag=1), ghosted_array], axis=self.axis)
+                    ghosted_array = ncp.concatenate([self._recv(self.base_comm, self.base_comm_nccl,
+                                                                recv_buf, source=self.rank - 1, tag=1,
+                                                                engine=self.engine), ghosted_array], axis=self.axis)
             # The skip in sender is to match with what described in receiver
             if self.rank != self.size - 1 and len(send_buf) != 0:
                 if cells_front > self.local_shape[self.axis]:
@@ -831,7 +833,9 @@ def add_ghost_cells(self, cells_front: Optional[int] = None,
                 recv_shape[self.axis] = total_cells_back[self.rank]
                 recv_buf = ncp.zeros(recv_shape, dtype=ghosted_array.dtype)
                 if len(recv_buf) != 0:
-                    ghosted_array = ncp.append(ghosted_array, self._recv(recv_buf, source=self.rank + 1, tag=0),
+                    ghosted_array = ncp.append(ghosted_array, self._recv(self.base_comm, self.base_comm_nccl,
+                                                                         recv_buf, source=self.rank + 1, tag=0,
+                                                                         engine=self.engine),
                                                axis=self.axis)
         return ghosted_array
 

pylops_mpi/utils/_mpi.py

@@ -8,7 +8,6 @@
 
 from typing import Optional, Union
 
-import numpy as np
 from mpi4py import MPI
 from pylops.utils import NDArray
 from pylops.utils.backend import get_module
@@ -37,12 +36,12 @@ def mpi_allgather(base_comm: MPI.Comm,
         a new buffer will be allocated with the appropriate shape.
     engine : :obj:`str`, optional
         Engine used to store array (``numpy`` or ``cupy``)
-    
+
     Returns
     -------
     recv_buf : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`
         A buffer containing the gathered data from all ranks.
-    
+
     """
     if deps.cuda_aware_mpi_enabled or engine == "numpy":
         send_shapes = base_comm.allgather(send_buf.shape)