Struct ShardingSpec 

pub struct ShardingSpec {
    pub node_id: usize,
    pub shard_dim: usize,
    pub num_shards: usize,
    pub shard_to_device: Vec<Device>,
}

Tensor sharding specification for model parallelism.

Fields

node_id: usize

Node ID that this sharding applies to

shard_dim: usize

Dimension along which to shard

num_shards: usize

Number of shards

shard_to_device: Vec<Device>

Device assignment for each shard

Implementations

impl ShardingSpec

pub fn new(node_id: usize, shard_dim: usize, devices: Vec<Device>) -> Self

Create a new sharding specification.

pub fn device_for_shard(&self, shard_id: usize) -> Option<&Device>

Get the device for a specific shard.

pub fn is_valid_shard(&self, shard_id: usize) -> bool

Check if a shard ID is valid.

Trait Implementations

impl Clone for ShardingSpec

fn clone(&self) -> ShardingSpec

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for ShardingSpec

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

Formats the value using the given formatter.

impl PartialEq for ShardingSpec

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

Tests for self and other values to be equal, and is used by ==.

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

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.

impl Eq for ShardingSpec

impl StructuralPartialEq for ShardingSpec