Trait DecoderLayer

pub trait DecoderLayer {
    type Cache;

    // Required method
    fn forward_t(
        &self,
        input: &Tensor,
        attention_mask: &AttentionMask,
        cache: &mut Self::Cache,
        positions: Option<&Tensor>,
        train: bool,
    ) -> Result<Tensor, BoxedError>;
}

Trait for decoder layers.

Required Associated Types

type Cache

Cache type for the decoder.

The cache can store the intermediate values of the decoder layer, avoiding recomputation when calling the decoder again for generating another output.

Required Methods

fn forward_t( &self, input: &Tensor, attention_mask: &AttentionMask, cache: &mut Self::Cache, positions: Option<&Tensor>, train: bool, ) -> Result<Tensor, BoxedError>

Apply the decoder layer to the given hidden representations.

• input - Hidden representations to apply the layer to. Shape: (batch_size, seq_len, width)
• attention_mask - Attention mask. Sequence elements for which the corresponding mask element is set to false are ignored during attention calculation. Shape: (batch_size, seq_len)
• cache - Cache to avoid recomputing intermediate values.
• positions - Input positions. Shape: (batch_size, seq_len)
• train - Whether to train the layer.

Returns layer output and the cache. Shape: (batch_size, seq_len, width)

Implementors

impl DecoderLayer for TransformerDecoderLayer

type Cache = LayerKeyValueCache