| 23 | } // namespace |
| 24 | |
| 25 | array all_sum( |
| 26 | const array& x, |
| 27 | std::optional<Group> group_ /* = std::nullopt */, |
| 28 | StreamOrDevice s /* = {} */) { |
| 29 | auto group = to_group(group_); |
| 30 | |
| 31 | if (group.size() == 1) { |
| 32 | return x; |
| 33 | } |
| 34 | auto stream = detail::communication_stream(group, s); |
| 35 | |
| 36 | return array( |
| 37 | x.shape(), |
| 38 | x.dtype(), |
| 39 | std::make_shared<AllReduce>(stream, group, AllReduce::Sum), |
| 40 | {x}); |
| 41 | } |
| 42 | |
| 43 | array all_max( |
| 44 | const array& x, |
no test coverage detected