mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-19 00:04:41 +08:00
In place all-reduce and forgiving init (#1178)
This commit is contained in:

committed by
GitHub

parent
4d485fca24
commit
3de8ce3f3c
@@ -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(
|
||||
|
Reference in New Issue
Block a user