| 17 | using namespace nb::literals; |
| 18 | |
| 19 | void init_distributed(nb::module_& parent_module) { |
| 20 | auto m = parent_module.def_submodule( |
| 21 | "distributed", "mlx.core.distributed: Communication operations"); |
| 22 | |
| 23 | nb::class_<mx::distributed::Group>( |
| 24 | m, |
| 25 | "Group", |
| 26 | R"pbcopy( |
| 27 | An :class:`mlx.core.distributed.Group` represents a group of independent mlx |
| 28 | processes that can communicate. |
| 29 | )pbcopy") |
| 30 | .def( |
| 31 | "rank", &mx::distributed::Group::rank, "Get the rank of this process") |
| 32 | .def("size", &mx::distributed::Group::size, "Get the size of the group") |
| 33 | .def( |
| 34 | "split", |
| 35 | &mx::distributed::Group::split, |
| 36 | "color"_a, |
| 37 | "key"_a = -1, |
| 38 | nb::sig("def split(self, color: int, key: int = -1) -> Group"), |
| 39 | R"pbdoc( |
| 40 | Split the group to subgroups based on the provided color. |
| 41 | |
| 42 | Processes that use the same color go to the same group. The ``key`` |
| 43 | argument defines the rank in the new group. The smaller the key the |
| 44 | smaller the rank. If the key is negative then the rank in the |
| 45 | current group is used. |
| 46 | |
| 47 | Args: |
| 48 | color (int): A value to group processes into subgroups. |
| 49 | key (int, optional): A key to optionally change the rank ordering |
| 50 | of the processes. |
| 51 | )pbdoc"); |
| 52 | |
| 53 | m.def( |
| 54 | "is_available", |
| 55 | [](const std::string& backend) { |
| 56 | return mx::distributed::is_available(backend); |
| 57 | }, |
| 58 | "backend"_a = "any", |
| 59 | nb::sig("def is_available(backend: str = 'any') -> bool"), |
| 60 | R"pbdoc( |
| 61 | Check if a communication backend is available. |
| 62 | |
| 63 | Note, this function returns whether MLX has the capability of |
| 64 | instantiating that distributed backend not whether it is possible to |
| 65 | create a communication group. For that purpose one should use |
| 66 | ``init(strict=True)``. |
| 67 | |
| 68 | Args: |
| 69 | backend (str, optional): The name of the backend to check for availability. |
| 70 | It takes the same values as :func:`init()`. Default: ``"any"``. |
| 71 | |
| 72 | Returns: |
| 73 | bool: Whether the distributed backend is available. |
| 74 | )pbdoc"); |
| 75 | |
| 76 | m.def( |
no test coverage detected