Trait Encoder

Source
pub trait Encoder {
    // Required methods
    fn encode(
        &self,
        input: &Tensor,
        attention_mask: Option<&Tensor>,
        train: bool,
    ) -> Result<Vec<LayerOutput>, TransformerError>;
    fn n_layers(&self) -> i64;
}
Expand description

Encoder networks.

Required Methods§

Source

fn encode( &self, input: &Tensor, attention_mask: Option<&Tensor>, train: bool, ) -> Result<Vec<LayerOutput>, TransformerError>

Apply the encoder.

Returns the output and attention per layer. The (optional) attention mask of shape [batch_size, time_steps] indicates which tokens should be included (true) and excluded (false) from attention. This can be used to mask inactive timesteps.

Source

fn n_layers(&self) -> i64

Get the number of layers that is returned by the encoder.

Implementors§