Change param name to buffer_chunk_shape

Author: rohanbabbar04

pylops_mpi/utils/_common.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 
-def _unroll_allgather_recv(recv_buf, chunk_shape, send_buf_shapes, displs=None) -> list:
+def _unroll_allgather_recv(recv_buf, buffer_chunk_shape, send_buf_shapes, displs=None) -> list:
     r"""Unroll recv_buf after Buffered Allgather (MPI and NCCL)
 
     Depending on the provided parameters, the function:
@@ -18,14 +18,14 @@ def _unroll_allgather_recv(recv_buf, chunk_shape, send_buf_shapes, displs=None)
     Parameters
     ----------
     recv_buf: :obj:`cupy.ndarray` or array-like
-        The data buffer returned from nccl_allgather call
+        The data buffer returned from the allgather call
     send_buf_shapes: :obj:`list`
         A list of original shapes of each rank's send_buf before any padding.
-    chunk_shape : tuple
-        Shape of each gathered chunk in recv_buf. This must match the shape
-        used to construct the gathered buffer: use the padded send buffer shape
-        when padding is required (e.g., nccl_allgather with padding), or the original send buffer
-        shape when no padding is used.
+    buffer_chunk_shape : tuple
+        Shape of each rank’s data as stored in ``recv_buf``. This should match
+        the layout used during allgather: use the padded send buffer shape when
+        padding is applied (e.g., NCCL), or the original send buffer shape when
+        no padding is used.
     displs : list, optional
         Starting offsets in recv_buf for each rank's data, used when chunks have
         variable sizes (e.g., mpi_allgather with displacements).
@@ -44,7 +44,8 @@ def _unroll_allgather_recv(recv_buf, chunk_shape, send_buf_shapes, displs=None)
             for i in range(ndev)
         ]
     else:
-        chunk_size = np.prod(chunk_shape)
+        # extract an individual array from each device
+        chunk_size = np.prod(buffer_chunk_shape)
         chunks = [
             recv_buf[i * chunk_size:(i + 1) * chunk_size]
             for i in range(ndev)
@@ -53,5 +54,5 @@ def _unroll_allgather_recv(recv_buf, chunk_shape, send_buf_shapes, displs=None)
         # in the middle of the flat array and thus the reshape and slicing for each dimension is required
         for i in range(ndev):
             slicing = tuple(slice(0, end) for end in send_buf_shapes[i])
-            chunks[i] = chunks[i].reshape(chunk_shape)[slicing]
+            chunks[i] = chunks[i].reshape(buffer_chunk_shape)[slicing]
         return chunks