Struct TensorShard 

pub struct TensorShard {
    pub rank: usize,
    pub world_size: usize,
    pub dim: ShardDim,
    pub original_shape: [usize; 2],
}

Metadata for one weight shard.

Fields

rank: usize

Rank that owns this shard.

world_size: usize

Total number of ranks.

dim: ShardDim

Shard dimension.

original_shape: [usize; 2]

Original (unsharded) shape as [rows, cols].

Implementations

impl TensorShard

pub fn start(&self) -> usize

Start index of this shard along the sharded dimension.

pub fn end(&self) -> usize

End index (exclusive) of this shard.

pub fn shard_size(&self) -> usize

Number of elements along the sharded dimension for this rank.

Trait Implementations

impl Clone for TensorShard

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 TensorShard

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

Formats the value using the given formatter.

impl<'de> Deserialize<'de> for TensorShard

fn deserialize<__D>(__deserializer: __D) -> Result<Self, __D::Error>

where __D: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl Serialize for TensorShard

fn serialize<__S>(&self, __serializer: __S) -> Result<__S::Ok, __S::Error>

where __S: Serializer,

Serialize this value into the given Serde serializer.