Struct dfdx::nn::TransformerDecoderBlock

pub struct TransformerDecoderBlock<const MODEL_DIM: usize, const NUM_HEADS: usize, const FF_DIM: usize> {
    pub self_attn: MultiHeadAttention<MODEL_DIM, NUM_HEADS>,
    pub norm1: LayerNorm1D<MODEL_DIM>,
    pub mh_attn: MultiHeadAttention<MODEL_DIM, NUM_HEADS>,
    pub norm2: LayerNorm1D<MODEL_DIM>,
    pub ff: Residual<(Linear<M, F>, ReLU, Linear<F, M>)>,
    pub norm3: LayerNorm1D<MODEL_DIM>,
}

Requires Nightly A transformer decoder block. Different than the normal transformer block as this self attention accepts an additional sequence from the encoder.

Generics

• MODEL_DIM: The size of query/key/value tensors. Given to MultiHeadAttention.
• NUM_HEADS: The number of heads in MultiHeadAttention.
• FF_DIM: The size of the hidden layer in the feedforward network.

Pytorch equivalent:

decoder = torch.nn.TransformerDecoderLayer(
   EMBED_DIM, NUM_HEADS, dim_feedforward=FF_DIM, batch_first=True, dropout=0.0
)

TODO: Doctests

Fields

self_attn: MultiHeadAttention<MODEL_DIM, NUM_HEADS>

norm1: LayerNorm1D<MODEL_DIM>

mh_attn: MultiHeadAttention<MODEL_DIM, NUM_HEADS>

norm2: LayerNorm1D<MODEL_DIM>

ff: Residual<(Linear<M, F>, ReLU, Linear<F, M>)>

norm3: LayerNorm1D<MODEL_DIM>

Trait Implementations

impl<const MODEL_DIM: usize, const NUM_HEADS: usize, const FF_DIM: usize> CanUpdateWithGradients for TransformerDecoderBlock<MODEL_DIM, NUM_HEADS, FF_DIM>

fn update<G: GradientProvider>(
    &mut self,
    grads: &mut G,
    unused: &mut UnusedTensors
)

Updates self given the GradientProvider. When any parameters that are NOT present in G, then this function should add the tensor’s UniqueId to UnusedTensors.

impl<const MODEL_DIM: usize, const NUM_HEADS: usize, const FF_DIM: usize> Clone for TransformerDecoderBlock<MODEL_DIM, NUM_HEADS, FF_DIM>

fn clone(&self) -> TransformerDecoderBlock<MODEL_DIM, NUM_HEADS, FF_DIM>

Returns a copy of the value.

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

Performs copy-assignment from source.

impl<const MODEL_DIM: usize, const NUM_HEADS: usize, const FF_DIM: usize> Debug for TransformerDecoderBlock<MODEL_DIM, NUM_HEADS, FF_DIM>

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

Formats the value using the given formatter.

impl<const MODEL_DIM: usize, const NUM_HEADS: usize, const FF_DIM: usize> Default for TransformerDecoderBlock<MODEL_DIM, NUM_HEADS, FF_DIM>

fn default() -> TransformerDecoderBlock<MODEL_DIM, NUM_HEADS, FF_DIM>

Returns the “default value” for a type.

impl<const M: usize, const H: usize, const F: usize> LoadFromNpz for TransformerDecoderBlock<M, H, F>

fn read<R: Read + Seek>(
    &mut self,
    pre: &str,
    r: &mut ZipArchive<R>
) -> Result<(), NpzError>

Reads this object from a ZipArchive. r with a base filename of filename_prefix.

fn load<P: AsRef<Path>>(&mut self, path: P) -> Result<(), NpzError>

Loads data from a .npz zip archive at the specified path.

impl<const M: usize, const H: usize, const F: usize, Tgt, Mem> Module<(Tgt, Mem)> for TransformerDecoderBlock<M, H, F>where
    Tgt: Tensor<Dtype = f32>,
    Mem: Tensor<Dtype = f32, NoTape = Mem> + Clone,
    MultiHeadAttention<M, H>: Module<(Tgt, Tgt::NoTape, Tgt::NoTape), Output = Tgt> + Module<(Tgt, Mem, Mem), Output = Tgt>,
    LayerNorm1D<M>: Module<Tgt, Output = Tgt>,
    Residual<(Linear<M, F>, ReLU, Linear<F, M>)>: Module<Tgt, Output = Tgt>,

type Output = Tgt

The type that this unit produces given Input.

fn forward(&self, (tgt, mem): (Tgt, Mem)) -> Self::Output

Forward Input through the module and produce Module::Output.

impl<const MODEL_DIM: usize, const NUM_HEADS: usize, const FF_DIM: usize> ResetParams for TransformerDecoderBlock<MODEL_DIM, NUM_HEADS, FF_DIM>

fn reset_params<R: Rng>(&mut self, rng: &mut R)

Mutate the unit’s parameters using rand::Rng. Each implementor of this trait decides how the parameters are initialized. In fact, some impls may not even use the rng.

impl<const M: usize, const H: usize, const F: usize> SaveToNpz for TransformerDecoderBlock<M, H, F>

fn write<W: Write + Seek>(&self, p: &str, w: &mut ZipWriter<W>) -> ZipResult<()>

Write this object into ZipWriter w with a base filename of filename_prefix.

fn save<P: AsRef<Path>>(&self, path: P) -> ZipResult<()>

Save this object into the .npz file determined located at path.