Enum TensorParallelLayer 

pub enum TensorParallelLayer {
    RowParallelLinear {
        input_size: usize,
        output_size: usize,
        bias: bool,
        input_is_parallel: bool,
    },
    ColumnParallelLinear {
        input_size: usize,
        output_size: usize,
        bias: bool,
        gather_output: bool,
    },
    ParallelEmbedding {
        num_embeddings: usize,
        embedding_dim: usize,
        padding_idx: Option<usize>,
    },
    ParallelAttention {
        hidden_size: usize,
        num_attention_heads: usize,
        dropout_prob: f32,
    },
}

Tensor parallel layer types

Variants

RowParallelLinear

Row-parallel linear layer

Fields

input_size: usize

output_size: usize

bias: bool

input_is_parallel: bool

ColumnParallelLinear

Column-parallel linear layer

Fields

input_size: usize

output_size: usize

bias: bool

gather_output: bool

ParallelEmbedding

Parallel embedding layer

Fields

num_embeddings: usize

embedding_dim: usize

padding_idx: Option<usize>

ParallelAttention

Parallel attention layer

Fields

hidden_size: usize

num_attention_heads: usize

dropout_prob: f32

Trait Implementations

impl Clone for TensorParallelLayer

fn clone(&self) -> TensorParallelLayer

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for TensorParallelLayer

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

Formats the value using the given formatter.