Struct OpSharding 

pub struct OpSharding {

    pub type: i32,
    pub tile_shape: Option<ShapeProto>,
    pub tile_assignment_dimensions: Vec<i64>,
    pub tile_assignment_devices: Vec<i64>,
    pub tuple_shardings: Vec<OpSharding>,
    pub replicate_on_last_tile_dim: bool,
    pub metadata: Vec<OpMetadata>,
    pub last_tile_dims: Vec<i32>,
    pub iota_reshape_dims: Vec<i64>,
    pub iota_transpose_perm: Vec<i32>,
    pub is_shard_group: bool,
    pub shard_group_id: i64,
    pub shard_group_type: i32,
}

LINT.IfChange

Fields

type: i32

tile_shape: Option<ShapeProto>

The shape of the sharded tile.

tile_assignment_dimensions: Vec<i64>

The shape of the tile assignment tensor - this must be the same rank as tile_shape and the product of its dimensions must equal tile_assignment_devices.size().

tile_assignment_devices: Vec<i64>

Flattened list of device IDs. The order of flattening is the same as used by IndexUtil::MultiToLinearIndex(tile_assignment_shape). Only one of tile_assignment_devices and iota_dimensions shall be non-empty.

tuple_shardings: Vec<OpSharding>

If type == TUPLE, the sub-shardings, one per leaf node in the tuple shape, in pre-order. The tuple shape could be nested; here we store just a flattened list of all leaves in the tuple shape. Note that the tuple shape is not stored here; shardings do not store the shapes to which they are applied, this is inferred from the instruction this sharding gets attached to.

replicate_on_last_tile_dim: bool

Only used for OTHER type. If true, data is sharded according to other dimensions of tile_assignment(), but replicated across devices along the last dimension. (Experimental)

metadata: Vec<OpMetadata>

This field is used to track the source of this sharding, usually derived from instructions. Multple metadata may be populated if sharding is combined with other shardings. Metadata are to not be populated when type == TUPLE and instead metadata should be set on individual tuple elements.

last_tile_dims: Vec<i32>

This field is used to represented the sharding type of each subgroup. For example, sharding={devices=[2,2,2,2]0,1,2,…,15 last_tile_dims={ replicate, manual, unreduced}} means that each of the last 3 dimensions in [2,2,2,2] represents a subgrouping in replicate, manual, unreduced sharding type respectively.

iota_reshape_dims: Vec<i64>

Dimensions used to reshape the 1D iota array of device IDs. Only one of tile_assignment_devices and iota_reshape_dims shall be non-empty.

iota_transpose_perm: Vec<i32>

Dimension permutations to transposed the iota array reshaped to iota_reshape_dims. This must have the same size as iota_reshape_dims.

is_shard_group: bool

This field decides whether this op is in a shard group.

shard_group_id: i64

This field is used to store the unique id of the shard group.

shard_group_type: i32

Implementations

impl OpSharding

pub fn type(&self) -> Type

Returns the enum value of type, or the default if the field is set to an invalid enum value.

pub fn set_type(&mut self, value: Type)

Sets type to the provided enum value.

pub fn last_tile_dims( &self, ) -> FilterMap<Cloned<Iter<'_, i32>>, fn(i32) -> Option<Type>>

Returns an iterator which yields the valid enum values contained in last_tile_dims.

pub fn push_last_tile_dims(&mut self, value: Type)

Appends the provided enum value to last_tile_dims.

pub fn shard_group_type(&self) -> ShardGroupType

Returns the enum value of shard_group_type, or the default if the field is set to an invalid enum value.

pub fn set_shard_group_type(&mut self, value: ShardGroupType)

Sets shard_group_type to the provided enum value.

Trait Implementations

impl Clone for OpSharding

fn clone(&self) -> OpSharding

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for OpSharding

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

Formats the value using the given formatter.

impl Default for OpSharding

fn default() -> Self

Returns the “default value” for a type.

impl Message for OpSharding

fn encoded_len(&self) -> usize

Returns the encoded length of the message without a length delimiter.

fn clear(&mut self)

Clears the message, resetting all fields to their default.

fn encode(&self, buf: &mut impl BufMut) -> Result<(), EncodeError>

where Self: Sized,

Encodes the message to a buffer.

fn encode_to_vec(&self) -> Vec<u8>

where Self: Sized,

Encodes the message to a newly allocated buffer.

fn encode_length_delimited( &self, buf: &mut impl BufMut, ) -> Result<(), EncodeError>

where Self: Sized,

Encodes the message with a length-delimiter to a buffer.

fn encode_length_delimited_to_vec(&self) -> Vec<u8>

where Self: Sized,

Encodes the message with a length-delimiter to a newly allocated buffer.

fn decode(buf: impl Buf) -> Result<Self, DecodeError>

where Self: Default,

Decodes an instance of the message from a buffer.

fn decode_length_delimited(buf: impl Buf) -> Result<Self, DecodeError>

where Self: Default,

Decodes a length-delimited instance of the message from the buffer.

fn merge(&mut self, buf: impl Buf) -> Result<(), DecodeError>

where Self: Sized,

Decodes an instance of the message from a buffer, and merges it into self.

fn merge_length_delimited(&mut self, buf: impl Buf) -> Result<(), DecodeError>

where Self: Sized,

Decodes a length-delimited instance of the message from buffer, and merges it into self.

impl PartialEq for OpSharding

fn eq(&self, other: &OpSharding) -> 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 StructuralPartialEq for OpSharding