pub fn sum_scatter(x: &Array, group: Option<&Group>) -> ArrayExpand description
Rank r gets the sum over group of every rank’s rows r·k..(r+1)·k of x, whose
first axis is size·k, as MLX’s sum_scatter (the reduce-scatter half of an
all-reduce). A group of size 1 returns x.
use nuvec_core::{array, distributed};
// Without a configured group, the group is this process alone.
let x = array![1.0, 2.0];
assert_eq!(distributed::sum_scatter(&x, None).to_vec::<f32>().unwrap(), [1.0, 2.0]);