Add Gemma3 allHiddenStates for frozen-text-encoder use

[xocialize]

Motivation

Multimodal generators increasingly use a frozen LLM as a text encoder, conditioning on the hidden states of every layer rather than on generated tokens. A concrete example: Lightricks LTX-2 conditions its video DiT on all 49 hidden states of Gemma-3-12B-it. We've been carrying this as a small fork while building a Swift port of LTX-2 — upstreaming so the next consumer of Gemma-3-as-encoder doesn't have to.

API

extension Gemma3Model {
    /// Embedding output plus each transformer layer's output —
    /// numHiddenLayers + 1 states, each (B, T, hiddenSize).
    public func allHiddenStates(
        _ inputs: MLXArray,
        mask: MLXFast.ScaledDotProductAttentionMaskMode
    ) -> [MLXArray]
}
// + forwarding convenience on Gemma3TextModel

Design notes:

• A single uniform caller-supplied mask is applied to every layer (the encoder use supplies a combined causal+padding mask); the per-layer sliding-window/global mask selection of the generation path is intentionally bypassed. This matches the reference LTX-2 encoder implementation.
• Per-layer eval keeps each Metal command buffer small so long-sequence (1024-token) encodes on 12B-class checkpoints stay under the macOS GPU watchdog — without it all 48 layers fuse into one dispatch. Measured as load-bearing in production use.
• Purely additive — no existing code paths change.

Testing

• New Gemma3AllHiddenStatesTests (tiny randomly-initialized model, no downloads, follows the existing ChatSessionTests pattern): state count/shapes, first state == scaled embedding, layers transform the state, determinism. 3/3 pass via xcodebuild test.
• Field validation: numerically parity-checked against the Python LTX-2 reference across all 49 layers of Gemma-3-12B-it-4bit at T=1024 (mean cosine 0.999866 vs reference goldens).

🤖 Generated with Claude Code