pub struct ColumnParallelLinear { /* private fields */ }Expand description
Column-parallel linear layer
Splits the weight matrix by columns (output dimension) across devices. This is typically used for the first linear layer in MLP blocks.
Implementations§
Source§impl ColumnParallelLinear
impl ColumnParallelLinear
pub fn new( in_features: usize, out_features: usize, bias: bool, mp_context: Arc<ModelParallelContext>, ) -> Result<Self>
pub fn forward(&self, input: &Tensor) -> Result<DistributedTensor>
pub fn backward( &mut self, grad_output: &DistributedTensor, input: &Tensor, ) -> Result<Tensor>
Auto Trait Implementations§
impl Freeze for ColumnParallelLinear
impl !RefUnwindSafe for ColumnParallelLinear
impl Send for ColumnParallelLinear
impl Sync for ColumnParallelLinear
impl Unpin for ColumnParallelLinear
impl UnsafeUnpin for ColumnParallelLinear
impl !UnwindSafe for ColumnParallelLinear
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more
Source§impl<T> Instrument for T
impl<T> Instrument for T
Source§fn instrument(self, span: Span) -> Instrumented<Self>
fn instrument(self, span: Span) -> Instrumented<Self>
Source§fn in_current_span(self) -> Instrumented<Self>
fn in_current_span(self) -> Instrumented<Self>
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
Converts
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
Converts
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more