nn

package

v1.9.12 Latest Latest Go to latest Published: Nov 9, 2025 License: MIT Imports: 3 Imported by: 0

Details

This section is empty.

This section is empty.

This section is empty.

type Attention struct {
	// contains filtered or unexported fields
}

func NewAttention(embedDim, numHeads int64, dropout float64) *Attention

func (a *Attention) Forward(q, k, v, mask *tensor.Tensor, isCausal bool) (*tensor.Tensor, *tensor.Tensor)

func (m *Attention) Parameters() []*tensor.Tensor

func (m *Attention) ToDevice(device consts.DeviceType)

func (m *Attention) ToScalarType(t consts.ScalarType)

type Embedding struct {
	// contains filtered or unexported fields
}

func NewEmbedding(numEmbeddings, embeddingDim, paddingIdx int64) *Embedding

func (m *Embedding) Forward(x *tensor.Tensor) *tensor.Tensor

func (m *Embedding) Parameters() []*tensor.Tensor

func (m *Embedding) ToDevice(device consts.DeviceType)

func (m *Embedding) ToScalarType(t consts.ScalarType)

type LayerNorm struct {
	// contains filtered or unexported fields
}

func NewLayerNorm(shapes ...int64) *LayerNorm

func (m *LayerNorm) Forward(x *tensor.Tensor) *tensor.Tensor

func (m *LayerNorm) Parameters() []*tensor.Tensor

func (m *LayerNorm) ToDevice(device consts.DeviceType)

func (m *LayerNorm) ToScalarType(t consts.ScalarType)

type Linear struct {
	// contains filtered or unexported fields
}

func NewLinear(inFeatures, outFeatures int64) *Linear

func (m *Linear) Forward(x *tensor.Tensor) *tensor.Tensor

func (m *Linear) Parameters() []*tensor.Tensor

func (m *Linear) ToDevice(device consts.DeviceType)

func (m *Linear) ToScalarType(t consts.ScalarType)