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.
45 lines
1.2 KiB
Go
45 lines
1.2 KiB
Go
package nn
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/ollama/ollama/mlx"
|
|
)
|
|
|
|
func TestQuantizedEmbeddingAsLinearPreservesGlobalScale(t *testing.T) {
|
|
weight := &mlx.Array{}
|
|
scales := &mlx.Array{}
|
|
qbiases := &mlx.Array{}
|
|
globalScale := &mlx.Array{}
|
|
|
|
embedding := &QuantizedEmbedding{
|
|
Weight: weight,
|
|
Scales: scales,
|
|
QBiases: qbiases,
|
|
GlobalScale: globalScale,
|
|
GroupSize: 16,
|
|
Bits: 4,
|
|
Mode: "nvfp4",
|
|
}
|
|
|
|
linear, ok := embedding.AsLinear().(*QuantizedLinear)
|
|
if !ok {
|
|
t.Fatalf("AsLinear type = %T, want *QuantizedLinear", embedding.AsLinear())
|
|
}
|
|
if linear.Weight != weight {
|
|
t.Fatalf("AsLinear Weight = %p, want %p", linear.Weight, weight)
|
|
}
|
|
if linear.Scales != scales {
|
|
t.Fatalf("AsLinear Scales = %p, want %p", linear.Scales, scales)
|
|
}
|
|
if linear.QBiases != qbiases {
|
|
t.Fatalf("AsLinear QBiases = %p, want %p", linear.QBiases, qbiases)
|
|
}
|
|
if linear.GlobalScale != globalScale {
|
|
t.Fatalf("AsLinear GlobalScale = %p, want %p", linear.GlobalScale, globalScale)
|
|
}
|
|
if linear.GroupSize != 16 || linear.Bits != 4 || linear.Mode != "nvfp4" {
|
|
t.Fatalf("AsLinear quant params = (%d, %d, %q), want (16, 4, %q)", linear.GroupSize, linear.Bits, linear.Mode, "nvfp4")
|
|
}
|
|
}
|