[GitHub] [tvm] LeshengJin commented on a diff in pull request #15764: [Disco] Add AllGather

["LeshengJin (via GitHub)" <gi...@apache.org>]

LeshengJin commented on code in PR #15764:
URL: https://github.com/apache/tvm/pull/15764#discussion_r1336456333


##########
python/tvm/relax/op/ccl/ccl.py:
##########
@@ -44,6 +47,27 @@ def allreduce(x, op_type: str = "sum"):  # pylint: disable=invalid-name
     return _ffi_api.allreduce(x, op_type)  # type: ignore # pylint: disable=no-member
 
 
+def allgather(x, num_workers: Union[int, PrimExpr, PrimValue]):  # pylint: disable=invalid-name
+    """AllGather operator
+
+    Parameters
+    ----------
+    x : relax.Expr
+      The input tensor.
+
+    num_worker : Union[int, PrimExpr, PrimValue]
+      The number of workers, i.e. the number of parts the given buffer should be chunked into.

Review Comment:
   ```suggestion
         The number of workers to gather data from.
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org