mirror of
https://github.com/ollama/ollama.git
synced 2026-11-12 03:46:14 -06:00
Jesse Gross
7343e22764
nn.go held every layer type in the package apart from attention, recurrence and rope: the Linear and Embedding interfaces with their dense and quantized types, Conv1d, RMSNorm, LayerNorm and MultiLinear, with one test file to match. Finding a layer meant scanning the file named after the package. Each layer kind gets its own file: linear.go and embedding.go hold the interface and the dense and quantized types, conv.go, norm.go and multilinear.go take the rest, and nn_test.go splits the same way. The Layer and MultiLinearLayer interfaces go; nothing implemented or accepted them. No code changes otherwise.
36 lines
721 B
Go
36 lines
721 B
Go
package nn
|
|
|
|
import "github.com/ollama/ollama/mlx"
|
|
|
|
// RMSNorm represents an RMS normalization layer.
|
|
type RMSNorm struct {
|
|
Weight *mlx.Array
|
|
Eps float32
|
|
}
|
|
|
|
func NewRMSNorm(weight *mlx.Array, eps float32) *RMSNorm {
|
|
return &RMSNorm{Weight: weight, Eps: eps}
|
|
}
|
|
|
|
func (rn *RMSNorm) Forward(x *mlx.Array, eps float32) *mlx.Array {
|
|
if eps == 0 {
|
|
eps = rn.Eps
|
|
}
|
|
return mlx.RMSNormFn(x, rn.Weight, eps)
|
|
}
|
|
|
|
// LayerNorm represents a standard layer normalization layer (with bias).
|
|
type LayerNorm struct {
|
|
Weight *mlx.Array
|
|
Bias *mlx.Array
|
|
Eps float32
|
|
}
|
|
|
|
func (ln *LayerNorm) Forward(x *mlx.Array) *mlx.Array {
|
|
eps := ln.Eps
|
|
if eps == 0 {
|
|
eps = 1e-5
|
|
}
|
|
return mlx.LayerNormFn(x, ln.Weight, ln.Bias, eps)
|
|
}
|