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-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
