multiheadattention

Documentation

Index

• type Cache
  • func (r Cache) At(i int) selfattention.Cache
• type Model
  • func New[T float.DType](size, numOfHeads int, useCausalMask, isCrossAttention bool) *Model
  • func (m *Model) Forward(cache Cache, q, x []mat.Tensor) ([]mat.Tensor, [][]mat.Tensor, Cache)
  • func (m *Model) Init(rng *rand.LockedRand)

Types

type Cache

type Cache []selfattention.Cache

Cache contains the self-attention cache for each head.

func (Cache) At

func (r Cache) At(i int) selfattention.Cache

type Model

type Model struct {
	nn.Module
	Heads       []*selfattention.Model
	OutputMerge *linear.Model
}

Model contains the serializable parameters.

func New

func New[T float.DType](size, numOfHeads int, useCausalMask, isCrossAttention bool) *Model

New returns a new model with parameters initialized to zeros.

func (*Model) Forward

func (m *Model) Forward(cache Cache, q, x []mat.Tensor) ([]mat.Tensor, [][]mat.Tensor, Cache)

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

func (*Model) Init

func (m *Model) Init(rng *rand.LockedRand)

Init initializes the self-attention heads and the merge layer with uniform Xavier random distribution.