Enum CollectiveOp 

#[non_exhaustive]
pub enum CollectiveOp {
    Reduce(AllReduce),
    Gather,
    SumScatter,
    Send {
        dst: usize,
    },
    Recv {
        src: usize,
        shape: Vec<usize>,
        dtype: DType,
    },
}

A collective operation’s kind, with its parameters (Primitive::Collective).

Variants (Non-exhaustive)

Non-exhaustive enums could have additional variants added in future. Therefore, when matching against variants of non-exhaustive enums, an extra wildcard arm must be added to account for any future variants.

Reduce(AllReduce)

AllReduce over the group.

Gather

Every rank’s array, concatenated along the first axis in rank order.

SumScatter

The sum over the group of every rank’s rows of this rank’s slice of the first axis.

Send

The input sent to rank dst; the output is the input.

Fields

dst: usize

The receiving rank.

Recv

An array of shape and dtype received from rank src.

Fields

src: usize

The sending rank.

shape: Vec<usize>

The received shape.

dtype: DType

The received dtype.

Implementations

impl CollectiveOp

pub fn name(&self) -> &'static str

The operation’s name (MLX’s).

Trait Implementations

impl Clone for CollectiveOp

fn clone(&self) -> Self

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for CollectiveOp

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

Formats the value using the given formatter.

impl Eq for CollectiveOp

impl Hash for CollectiveOp

fn hash<__H: Hasher>(&self, state: &mut __H)

Feeds this value into the given Hasher.

fn hash_slice<H>(data: &[Self], state: &mut H)

where H: Hasher, Self: Sized,

Feeds a slice of this type into the given Hasher.

impl PartialEq for CollectiveOp

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

Equality operator ==.

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

Inequality operator !=.

impl StructuralPartialEq for CollectiveOp