feat(model/gemma4): neutral vision weight-name canonicalisation (vision loader 1/n)

First slice of the gemma4 vision-tower port (#8): the neutral front of metal's vision_load.go — canonicalise
vision weight names off a checkpoint's tensors (strip the multimodal wrapper prefixes via the existing
profile.TrimWeightWrapperPrefix, then the vision_tower./vision_model. tower prefix; multi_modal_projector./
embed_vision. keep their prefix) + detect tower-vs-projector-only packs (via the existing model.WeightAny).
Reuses what's there rather than re-porting it. The device build of the SigLIP tower from these weights is
the next piece, backend-side in pkg/native — same split as the text path (neutral assemble -> native build).

pkg/model green.

Co-Authored-By: Virgil <virgil@lethean.io>

Author: Snider

go/pkg/model/gemma4/vision_weights.go

@@ -0,0 +1,80 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package gemma4
+
+import (
+	core "dappco.re/go"
+	"dappco.re/go/mlx/pkg/model"
+	"dappco.re/go/mlx/pkg/safetensors"
+	"dappco.re/go/mlx/profile"
+)
+
+// vision_weights.go is the neutral front of the gemma4 vision-tower loader: it canonicalises the vision
+// weight names off a checkpoint's tensors (strip the multimodal wrapper prefixes, then the
+// vision_tower./vision_model. tower prefix) and detects whether a pack carries a full SigLIP tower or
+// only the projector. The DEVICE build of the tower from these weights lives backend-side (pkg/native),
+// the same split as the text path (neutral assemble → native device-build). Lifted from
+// pkg/metal/model/gemma4/vision_load.go, reusing profile.TrimWeightWrapperPrefix + model.WeightAny
+// rather than re-porting them.
+
+// canonicalGemma4VisionWeightName strips the wrapper prefixes then the vision-tower prefix, returning the
+// canonical vision-weight name and whether the tensor is a vision weight at all. multi_modal_projector.
+// and embed_vision. weights keep their prefix — the projector reads them by full name.
+func canonicalGemma4VisionWeightName(name string) (string, bool) {
+	trimmed := name
+	for {
+		next, changed := profile.TrimWeightWrapperPrefix("gemma4", trimmed)
+		if !changed {
+			break
+		}
+		trimmed = next
+	}
+	for _, prefix := range []string{"vision_tower.", "vision_model."} {
+		if core.HasPrefix(trimmed, prefix) {
+			return core.TrimPrefix(trimmed, prefix), true
+		}
+	}
+	for _, prefix := range []string{"multi_modal_projector.", "embed_vision."} {
+		if core.HasPrefix(trimmed, prefix) {
+			return trimmed, true
+		}
+	}
+	return "", false
+}
+
+// SanitizeVisionWeights returns the vision weights from a checkpoint's tensors, keyed by their canonical
+// names — the input the device-side vision build consumes.
+func SanitizeVisionWeights(raw map[string]safetensors.Tensor) map[string]safetensors.Tensor {
+	vision := make(map[string]safetensors.Tensor)
+	for name, t := range raw {
+		if canonical, ok := canonicalGemma4VisionWeightName(name); ok {
+			vision[canonical] = t
+		}
+	}
+	return vision
+}
+
+// HasVisionTowerWeights reports whether the pack carries a full SigLIP vision tower (a patch embedder),
+// vs only the multimodal projector.
+func HasVisionTowerWeights(weights map[string]safetensors.Tensor) bool {
+	_, ok := model.WeightAny(weights,
+		"patch_embedder.input_proj.weight",
+		"patch_embedder.input_proj.linear.weight",
+		"embeddings.patch_embedding.weight",
+		"patch_embedding.weight",
+	)
+	return ok
+}
+
+// HasVisionProjectionWeights reports whether the pack carries the vision-to-text multimodal projector.
+func HasVisionProjectionWeights(weights map[string]safetensors.Tensor) bool {
+	_, ok := model.WeightAny(weights,
+		"embed_vision.embedding_projection.weight",
+		"embed_vision.embedding_projection.linear.weight",
+		"multi_modal_projector.embedding_projection.weight",
+		"multi_modal_projector.embedding_projection.linear.weight",
+		"multi_modal_projector.proj.weight",
+		"multi_modal_projector.weight",
+	)
+	return ok
+}

go/pkg/model/gemma4/vision_weights_test.go

@@ -0,0 +1,55 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package gemma4
+
+import (
+	"testing"
+
+	"dappco.re/go/mlx/pkg/safetensors"
+)
+
+// TestCanonicalGemma4VisionWeightName pins the canonicalisation: the vision_tower./vision_model. prefix
+// is stripped, the projector prefixes are kept, and a text weight is rejected.
+func TestCanonicalGemma4VisionWeightName(t *testing.T) {
+	cases := []struct {
+		in, want string
+		ok       bool
+	}{
+		{"vision_tower.encoder.layers.0.self_attn.q_proj.weight", "encoder.layers.0.self_attn.q_proj.weight", true},
+		{"vision_model.embeddings.patch_embedding.weight", "embeddings.patch_embedding.weight", true},
+		{"multi_modal_projector.proj.weight", "multi_modal_projector.proj.weight", true},
+		{"embed_vision.embedding_projection.weight", "embed_vision.embedding_projection.weight", true},
+		{"model.layers.0.self_attn.q_proj.weight", "", false},
+	}
+	for _, c := range cases {
+		got, ok := canonicalGemma4VisionWeightName(c.in)
+		if got != c.want || ok != c.ok {
+			t.Fatalf("canonicalGemma4VisionWeightName(%q) = (%q,%v), want (%q,%v)", c.in, got, ok, c.want, c.ok)
+		}
+	}
+}
+
+// TestSanitizeAndDetectVisionWeights pins the gather + tower/projector detection over a tiny tensor set.
+func TestSanitizeAndDetectVisionWeights(t *testing.T) {
+	raw := map[string]safetensors.Tensor{
+		"vision_tower.embeddings.patch_embedding.weight": {Dtype: "BF16", Shape: []int{8, 4}},
+		"multi_modal_projector.proj.weight":              {Dtype: "BF16", Shape: []int{4, 8}},
+		"model.layers.0.self_attn.q_proj.weight":         {Dtype: "BF16", Shape: []int{4, 4}}, // text — dropped
+	}
+	vision := SanitizeVisionWeights(raw)
+	if _, ok := vision["embeddings.patch_embedding.weight"]; !ok {
+		t.Fatal("tower weight should be present under its canonical name")
+	}
+	if _, ok := vision["multi_modal_projector.proj.weight"]; !ok {
+		t.Fatal("projector weight should keep its prefix")
+	}
+	if len(vision) != 2 {
+		t.Fatalf("the text weight should be dropped: got %d vision weights", len(vision))
+	}
+	if !HasVisionTowerWeights(vision) {
+		t.Fatal("should detect a full vision tower")
+	}
+	if !HasVisionProjectionWeights(vision) {
+		t.Fatal("should detect the multimodal projector")
+	}
+}