feat: moved _send to Distributed

Author: mrava87

pylops_mpi/Distributed.py

@@ -2,7 +2,7 @@
 
 from mpi4py import MPI
 from pylops.utils import deps as pylops_deps  # avoid namespace crashes with pylops_mpi.utils
-from pylops_mpi.utils._mpi import mpi_allreduce
+from pylops_mpi.utils._mpi import mpi_allreduce, mpi_send
 from pylops_mpi.utils import deps
 
 cupy_message = pylops_deps.cupy_import("the DistributedArray module")
@@ -43,3 +43,15 @@ def _allreduce_subcomm(self, send_buf, recv_buf=None, op: MPI.Op = MPI.SUM):
         else:
             return mpi_allreduce(self.sub_comm, send_buf, 
                                  recv_buf, self.engine, op)
+
+    def _send(self, send_buf, dest, count=None, tag=0):
+        """Send operation
+        """
+        if deps.nccl_enabled and self.base_comm_nccl:
+            if count is None:
+                count = send_buf.size
+            nccl_send(self.base_comm_nccl, send_buf, dest, count)
+        else:
+            mpi_send(self.base_comm,
+                     send_buf, dest, count, tag=tag,
+                     engine=self.engine)

pylops_mpi/DistributedArray.py

@@ -507,18 +507,6 @@ def _allgather_subcomm(self, send_buf, recv_buf=None):
                 return self.sub_comm.allgather(send_buf)
             self.sub_comm.Allgather(send_buf, recv_buf)
 
-    def _send(self, send_buf, dest, count=None, tag=0):
-        """Send operation
-        """
-        if deps.nccl_enabled and self.base_comm_nccl:
-            if count is None:
-                count = send_buf.size
-            nccl_send(self.base_comm_nccl, send_buf, dest, count)
-        else:
-            mpi_send(self.base_comm,
-                     send_buf, dest, count, tag=tag,
-                     engine=self.engine)
-
     def _recv(self, recv_buf=None, source=0, count=None, tag=0):
         """Receive operation
         """