Fix docs and comment complex case in test

Author: astroC86

pylops_mpi/Distributed.py

@@ -1,4 +1,4 @@
-from typing import Any, NewType, Optional, Union
+from typing import Any, NewType, Optional
 
 from mpi4py import MPI
 from pylops.utils import NDArray
@@ -198,12 +198,25 @@ def _bcast(self,
                root: int = 0,
                engine: str = "numpy",
                ) -> NDArray:
-        """BCast operation.
+        """Broadcast operation
 
-        Notes
-        -----
-        Any root-only assignment (e.g., setting a value prior to broadcast) must
-        be done outside this method.
+        Parameters
+        ----------
+        base_comm : :obj:`MPI.Comm`
+            Base MPI Communicator.
+        base_comm_nccl : :obj:`cupy.cuda.nccl.NcclCommunicator`
+            NCCL Communicator.
+        send_buf : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`
+            A buffer containing the data to be broadcast from the root rank.
+        root : :obj:`int`, optional
+            The rank of the process that holds the source data.
+        engine : :obj:`str`, optional
+            Engine used to store array (``numpy`` or ``cupy``)
+
+        Returns
+        -------
+        send_buf : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`
+            The buffer containing the broadcasted data.
         """
         if deps.nccl_enabled and base_comm_nccl is not None:
             nccl_bcast(base_comm_nccl, send_buf, root=root)

pylops_mpi/DistributedArray.py

@@ -205,7 +205,7 @@ def __setitem__(self, index, value):
         """
         if self.partition is Partition.BROADCAST:
             view = self.local_array[index]
-            if self.rank == 0: 
+            if self.rank == 0:
                 view[...] = value
             view = self._bcast(self.base_comm, self.base_comm_nccl,
                                view, root=0, engine=self.engine)

pylops_mpi/utils/_mpi.py

@@ -6,7 +6,7 @@
     "mpi_recv",
 ]
 
-from typing import Optional, Union
+from typing import Optional
 
 from mpi4py import MPI
 from pylops.utils import NDArray
@@ -113,10 +113,22 @@ def mpi_bcast(base_comm: MPI.Comm,
     Dispatch bcast routine based on type of input and availability of
     CUDA-Aware MPI.
 
-    Notes
-    -----
-    Any root-only assignment (e.g., setting a value prior to broadcast) must be
-    done outside this function.
+    Parameters
+    ----------
+    base_comm : :obj:`MPI.Comm`
+        Base MPI Communicator.
+    send_buf : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`
+        The data buffer to be broadcasted to the other ranks from the broadcasting root rank.
+    root : :obj:`int`, optional
+        The rank of the broadcasting process.
+    engine : :obj:`str`, optional
+        Engine used to store array (``numpy`` or ``cupy``)
+
+    Returns
+    -------
+    send_buf : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`
+        The buffer containing the broadcasted data.
+
     """
     if deps.cuda_aware_mpi_enabled or engine == "numpy":
         base_comm.Bcast(send_buf, root=root)

pylops_mpi/utils/_nccl.py

@@ -240,11 +240,14 @@ def nccl_allreduce(nccl_comm, send_buf, recv_buf=None, op: MPI.Op = MPI.SUM) ->
 
 def nccl_bcast(nccl_comm, send_buf, root: int = 0) -> None:
     """NCCL equivalent of MPI_Bcast for an array buffer.
-
-    Notes
-    -----
-    Any root-only assignment (e.g., setting a value prior to broadcast) must be
-    done outside this function.
+    Parameters
+    ----------
+    nccl_comm : :obj:`cupy.cuda.nccl.NcclCommunicator`
+        The NCCL communicator used for collective communication.
+    send_buf: :obj:`numpy.ndarray` or :obj:`cupy.ndarray`
+        The data buffer to be broadcasted to the other ranks from the broadcasting root rank.
+    root: :obj:`int`
+        The rank of the broadcasting process.
     """
     send_buf = send_buf if isinstance(send_buf, cp.ndarray) else cp.asarray(send_buf)
     nccl_comm.bcast(

tests/test_matrixmult.py

@@ -41,7 +41,7 @@
     pytest.param(37, 37, 37, "float64", id="f32_37_37_37"),
     pytest.param(50, 30, 40, "float64", id="f64_50_30_40"),
     # temporarely removed as sometimes crashed CI... to be investigated
-    pytest.param(22, 20, 16, "complex64", id="c64_22_20_16"),
+    # pytest.param(22, 20, 16, "complex64", id="c64_22_20_16"),
     pytest.param(3, 4, 5, "float32", id="f32_3_4_5"),
     pytest.param(1, 2, 1, "float64", id="f64_1_2_1",),
     pytest.param(2, 1, 3, "float32", id="f32_2_1_3",),