In place all-reduce and forgiving init (#1178)

This commit is contained in:
Angelos Katharopoulos
2024-06-03 16:47:47 -07:00
committed by GitHub
parent 4d485fca24
commit 3de8ce3f3c
5 changed files with 31 additions and 7 deletions

View File

@@ -54,8 +54,18 @@ void init_distributed(nb::module_& parent_module) {
m.def(
"init",
&distributed::init,
"strict"_a = false,
nb::sig("def init(strict: bool = False) -> Group"),
R"pbdoc(
Initialize the communication backend and create the global communication group.
Args:
strict (bool, optional): If set to False it returns a singleton group
in case ``mx.distributed.is_available()`` returns False otherwise
it throws a runtime error. Default: ``False``
Returns:
Group: The group representing all the launched processes.
)pbdoc");
m.def(