Struct CollectiveShape 

pub struct CollectiveShape {
    pub input: Shape,
    pub output: Shape,
    pub input_elements: usize,
    pub output_elements: usize,
}

Logical input and rank-ordered output of an equal-size collective.

Fields

input: Shape

Each rank’s original tensor shape.

output: Shape

Gathered shape or this rank’s contiguous reduced shard.

input_elements: usize

Logical input element count, independent of view strides.

output_elements: usize

Logical output element count.

Implementations

impl CollectiveShape

pub fn all_gather( input: Shape, world_size: usize, ) -> Result<CollectiveShape, MetadataError>

Concatenate each rank’s tensor along axis zero, in rank order.

pub fn reduce_scatter( input: Shape, world_size: usize, ) -> Result<CollectiveShape, MetadataError>

Partition a reduced tensor into equal contiguous axis-zero shards. The leading axis, not just the total volume, must divide by world size.

Trait Implementations

impl Clone for CollectiveShape

fn clone(&self) -> CollectiveShape

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for CollectiveShape

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.

impl Eq for CollectiveShape

impl PartialEq for CollectiveShape

fn eq(&self, other: &CollectiveShape) -> bool

Equality operator ==.

fn ne(&self, other: &Rhs) -> bool

Inequality operator !=.

impl StructuralPartialEq for CollectiveShape