Skip to main content

sum_scatter

Function sum_scatter 

Source
pub fn sum_scatter(x: &Array, group: Option<&Group>) -> Array
Expand 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]);