Enum LayerOutput

pub enum LayerOutput {
    Embedding(Tensor),
    EncoderWithAttention(HiddenLayer),
}

Output of a BERT layer.

Variants

Embedding(Tensor)

Embedding layer output.

EncoderWithAttention(HiddenLayer)

Encoder layer output.

Implementations

impl LayerOutput

pub fn attention(&self) -> Option<&Tensor>

Get the layer attention.

Return a Some value if the layer output is from an encoder layer, or None otherwise.

pub fn embedding(&self) -> Option<&Tensor>

Get the embedding.

Returns Some if the layer output is an embedding or None otherwise.

pub fn map_output<F>(&self, f: F) -> Result<Self, TransformerError>

where F: Fn(&Tensor) -> Result<Tensor, TransformerError>,

Map the output representation of this layer.

pub fn output(&self) -> &Tensor

Get the layer output.

pub fn output_mut(&mut self) -> &mut Tensor

Get the layer output mutably.

pub fn encoder_with_attention(&self) -> Option<&HiddenLayer>

Get the output of an encoder layer.

Return a Some value if the layer output is from an encoder layer, or None otherwise.

Trait Implementations

impl Debug for LayerOutput

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

Formats the value using the given formatter.