zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-03 09:58:17 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Add docs for the distributed namespace (#1184)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-06-06 11:37:00 -07:00
committed by GitHub
parent 578842954c
commit 0163a8e57a
12 changed files with 202 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/src/distributed.cpp
View File

@@ -69,13 +69,13 @@ void init_distributed(nb::module_& parent_module) {
)pbdoc");
m.def(
"all_reduce_sum",
&distributed::all_reduce_sum,
"all_sum",
&distributed::all_sum,
"x"_a,
nb::kw_only(),
"group"_a = nb::none(),
nb::sig(
"def all_reduce_sum(x: array, *, group: Optional[Group] = None) -> array"),
"def all_sum(x: array, *, group: Optional[Group] = None) -> array"),
R"pbdoc(
All reduce sum.