layer

Documentation

Index

• type KeysValuesPairs
• type Layer
  • func NewLayer(config config.Config) *Layer
  • func (m *Layer) Forward(xs []ag.Node, encoderHiddenStates []ag.Node, ...) ([]ag.Node, KeysValuesPairs)

Types

type KeysValuesPairs

type KeysValuesPairs struct {
	// SelfAttKeyValues contains the keys and values used by self-attention.
	SelfAttKeyValues multiheadattention.KeysValuesPairs
	// CrossAttKeyValues contains the keys and values used by cross-attention.
	CrossAttKeyValues multiheadattention.KeysValuesPairs
}

KeysValuesPairs contains the keys and values used by the self-attention and cross-attention blocks.

type Layer

type Layer struct {
	nn.BaseModel
	Config                    config.Config
	SelfAttention             *multiheadattention.Model
	SelfAttentionLayerNorm    *layernorm.Model
	EncoderAttention          *multiheadattention.Model
	EncoderAttentionLayerNorm *layernorm.Model
	FFN                       *stack.Model
	LayerNorm                 *layernorm.Model
}

Layer implements a BART decoder layer.

func NewLayer

func NewLayer(config config.Config) *Layer

NewLayer returns a new BART decoder Layer.

func (*Layer) Forward

func (m *Layer) Forward(
	xs []ag.Node,
	encoderHiddenStates []ag.Node,
	pastProjKeysValues KeysValuesPairs,
) ([]ag.Node, KeysValuesPairs)

Forward performs the forward step for each input and returns the result.