nnterp · encyclopedia gemma4_unified_text · nnterp.families.gemma4_unified_text

family gemma4_unified_text

Gemma 4 unified

A sandwich block whose sum is multiplied by a learned per-block scalar, with normed queries, keys and values, and full blocks that take their values from the key projection.

Sandwich normsPer-block sizesScaled residual addsQuery/key normsSliding windowProportional rotaryScaled embeddingsSoftcapped logitsGrouped-query attentionVision-languageEncoder-freeVariable resolutionPadded patches

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block; this family's blocks differ by layer_types.

embed_tokenstoken_embeddings
→
layers× 48
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitssoftcapped
i = 0

identity (layers[i].input + self_attn.attention_output + mlp.mlp_output) * layer_scalar == layer_output

The vision encoder

· Gemma 4 unified embedder · Gemma4UnifiedVisionEmbedder

The image's path into the text model. Hover any part for its nnterp name; click to pin the card.

imagepixel_values
→
patch_embedpatch_embeddings
tower_output
projectorinput: vision.tower_output: the embedder's states, the padded rows included
→
image_featuresat the scatter
→
layers[0].inputat image_token_mask
rows
One row per image: its merged patches, then padding up to max_soft_tokens rows (280); the processor's image_position_ids is (-1, -1) on the padded rows.
positions
patch_embeddings is patch_dense's output, before patch_ln2; a factorized 2D position embedding (a learned x and a learned y row per patch, zero on the padded rows) is added after patch_ln2, and pos_norm norms the sum.
masking
None: the embedder has no attention, so no row reads another.
final norm
pos_norm, a LayerNorm, ends the embedder's states, so vision.tower_output is its output; it is not vision.norm.

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model Gemma4UnifiedForConditionalGenerationlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] Gemma4UnifiedTextDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module pre_feedforward_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module pre_feedforward_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn Gemma4UnifiedTextAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream: the post-attention norm's output

the sibling module post_attention_layernorm's output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp Gemma4UnifiedTextMLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream: the post-feedforward norm's output

the sibling module post_feedforward_layernorm's output

the vision encoder · Gemma 4 unified embedder

model.vision Gemma4UnifiedVisionEmbedderlayoutwhat it isread at
image_token_mask ImageTokenMask [batch seq] Which positions hold image tokens: input_ids == the config's image_token_id; read-only

the model's own inputs, at the root
patch_embeddings Patches [images patches vision_hidden_size] The patch embedding's output, one row per patch

the child module patch_embed's output
tower_output Patches [images patches vision_hidden_size] The embedder's states over the patches before the projection

the child module multimodal_embedder's input
image_features ImageFeatures [image_tokens hidden_size] The image features the text model receives at the image tokens, flat over them

inside model.model's forward: the arguments of inputs_embeds_masked_scatter_0, element 1

Printout

What print(model) shows on google/gemma-4-12B: the native modules, standard names and every standard value with its layout.

Gemma4UnifiedForConditionalGeneration(
  (lm_head): Linear(in_features=3840, out_features=262144, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Gemma4UnifiedTextScaledWordEmbedding(262144, 3840, padding_idx=0)
  (layers): ModuleList(
    (0-4): 5 x Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=4096, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (v_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (o_proj): Linear(in_features=4096, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (5): Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=8192, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=512, bias=False)
        (o_proj): Linear(in_features=8192, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (6-10): 5 x Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=4096, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (v_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (o_proj): Linear(in_features=4096, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (11): Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=8192, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=512, bias=False)
        (o_proj): Linear(in_features=8192, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (12-16): 5 x Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=4096, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (v_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (o_proj): Linear(in_features=4096, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (17): Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=8192, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=512, bias=False)
        (o_proj): Linear(in_features=8192, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (18-22): 5 x Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=4096, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (v_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (o_proj): Linear(in_features=4096, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (23): Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=8192, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=512, bias=False)
        (o_proj): Linear(in_features=8192, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (24-28): 5 x Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=4096, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (v_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (o_proj): Linear(in_features=4096, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (29): Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=8192, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=512, bias=False)
        (o_proj): Linear(in_features=8192, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (30-34): 5 x Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=4096, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (v_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (o_proj): Linear(in_features=4096, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (35): Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=8192, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=512, bias=False)
        (o_proj): Linear(in_features=8192, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (36-40): 5 x Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=4096, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (v_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (o_proj): Linear(in_features=4096, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (41): Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=8192, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=512, bias=False)
        (o_proj): Linear(in_features=8192, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (42-46): 5 x Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=4096, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (v_proj): Linear(in_features=3840, out_features=2048, bias=False)
        (o_proj): Linear(in_features=4096, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (47): Gemma4UnifiedTextDecoderLayer(
      (self_attn): Gemma4UnifiedTextAttention(
        (q_proj): Linear(in_features=3840, out_features=8192, bias=False)
        (q_norm): Gemma4UnifiedRMSNorm()
        (k_norm): Gemma4UnifiedRMSNorm()
        (v_norm): Gemma4UnifiedRMSNorm()
        (k_proj): Linear(in_features=3840, out_features=512, bias=False)
        (o_proj): Linear(in_features=8192, out_features=3840, bias=False)
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma4UnifiedTextMLP(
        (gate_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (up_proj): Linear(in_features=3840, out_features=15360, bias=False)
        (down_proj): Linear(in_features=15360, out_features=3840, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma4UnifiedRMSNorm()
      (post_attention_layernorm): Gemma4UnifiedRMSNorm()
      (pre_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (post_feedforward_layernorm): Gemma4UnifiedRMSNorm()
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): Gemma4UnifiedRMSNorm()
  (vision): Gemma4UnifiedVisionEmbedder(
    (patch_ln1): LayerNorm((6912,), eps=1e-05, elementwise_affine=True, bias=True)
    (patch_embed/patch_dense): Linear(in_features=6912, out_features=3840, bias=True)
    (patch_ln2): LayerNorm((3840,), eps=1e-05, elementwise_affine=True, bias=True)
    (pos_norm): LayerNorm((3840,), eps=1e-05, elementwise_affine=True, bias=True)
    (multimodal_embedder): Gemma4UnifiedMultimodalEmbedder(
      (embedding_projection): Linear(in_features=3840, out_features=3840, bias=False)
      (embedding_pre_projection_norm): Gemma4UnifiedRMSNorm()
    )
    (tower_output) -> Patches [images patches vision_hidden_size]: The embedder's states over the patches before the projection
    (image_token_mask) -> ImageTokenMask [batch seq]: Which positions hold image tokens: input_ids == the config's image_token_id; read-only
    (patch_embeddings) -> Patches [images patches vision_hidden_size]: The patch embedding's output, one row per patch
    (image_features) -> ImageFeatures [image_tokens hidden_size]: The image features the text model receives at the image tokens, flat over them
  )
  (projector): Gemma4UnifiedMultimodalEmbedder(
    (embedding_projection): Linear(in_features=3840, out_features=3840, bias=False)
    (embedding_pre_projection_norm): Gemma4UnifiedRMSNorm()
  )
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers48
model.hidden_size3840
model.num_heads16
model.num_kv_heads8
model.head_dim256
model.qk_head_dim256
model.vocab_size262144
model.intermediate_size15360
architecturesGemma4UnifiedForConditionalGeneration
num_hidden_layers48
hidden_size3840
intermediate_size15360
num_attention_heads16
num_key_value_heads8
head_dim256
vocab_size262144
max_position_embeddings262144
sliding_window1024
hidden_activationgelu_pytorch_tanh
rms_norm_eps1e-06
final_logit_softcapping30.0
tie_word_embeddingsTrue
dtypebfloat16
num_kv_shared_layers0
attention_biasFalse
enable_moe_blockFalse
hidden_size_per_layer_input0
attention_k_eq_vTrue
use_bidirectional_attentionvision
Vision Config

Vision Config

model.vision.num_layers0
model.vision.vision_hidden_size3840
model.vision.num_headsnone
model.vision.head_dimnone
model.vision.intermediate_sizenone
model.vision.patch_size48
model.vision.image_sizevaries

The Notes

03 · what to know before you run

The block, in order

a   = input_layernorm(x)
q   = rope(q_norm(q_proj(a)))         # per head
k   = rope(k_norm(k_proj(a)))         # per head
v   = v_norm(v_proj(a))               # per head, no gain; k_proj(a) on 12B's full blocks
h   = x + post_attention_layernorm(o_proj(attend(q, k, v)))      # no 1/√d
h   = h + post_feedforward_layernorm(mlp(pre_feedforward_layernorm(h)))
out = h * layer_scalar

Four RMSNorms around the sublayers and three inside the attention. post_attention_layernorm follows the attention; the MLP's input norm is pre_feedforward_layernorm. The block has no per-layer branch and no mixture of experts: per_layer_output and the mixture's values are not values of this family. layer_scalar is a buffer, layers[i]._module.layer_scalar, applied in place.

The block's sum is scaled; its terms are served unscaled

attention_output and mlp_output are the post-norms' outputs, before layer_scalar. The identity is exact in float32 on the pinned checkpoint:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

scalar = model.layers[1]._module.layer_scalar
torch.testing.assert_close((x + attn + mlp) * scalar, out)

A vector added to mlp_output reaches layer_output multiplied by the block's scalar; one added to layer_output does not. Block i's terms reach the last stream multiplied by its own scalar and every later one, so direct logit attribution weights each term by that product first. The scalars are far from one on 12B: block 11's is 0.005.

Queries, keys and values are normed, and the scores are not scaled

q_norm and k_norm norm each head before the rotary; v_norm norms each value head with no gain, so every head of attention_values has RMS 1. The attention's scaling is 1.0: the scores are q · k with no 1/√head_dim, and the norms' gains set the temperature. attention_queries and attention_keys are read after their norms and the rotary.

Full blocks take their values from the key projection

12B sets attention_k_eq_v: its full blocks have no v_proj, and their values are v_norm(k_proj(a)), the keys' projection before k_norm and the rotary. An edit to a full block's k_proj output moves its keys and its values at once; attention_keys and attention_values are served apart and edit apart. Each full block has one key/value head for its 16 query heads, so an edit to a full block's attention_keys or attention_values reaches all 16; a sliding block has 8, each serving two.

Sliding and full blocks differ in width and rotary

config.layer_types makes every sixth block full on 12B (5, 11, ..., 47), the last block among them; the window is 1024 tokens. Full blocks have 512-wide heads and one key/value head; sliding blocks 256-wide heads and 8. The root's sizes are the sliding blocks'; each block's own are on its attention:

model.head_dim, model.layers[5].self_attn.head_dim        # 256, 512 on 12B
model.num_kv_heads, model.layers[5].self_attn.num_kv_heads  # 8, 1 on 12B

Sliding blocks rotate all 256 dimensions with base 10,000. Full blocks use proportional rotary (partial_rotary_factor 0.25, base 1,000,000): only a quarter of each head turns, at frequencies spaced over the whole head; the other dimensions carry no position.

The family serves borrowed keys and values; 12B borrows none

The attention is Gemma 4's: on a config with num_kv_shared_layers above 0 the last blocks attend with an earlier block's keys and values, served as a copy private to the block. 12B sets num_kv_shared_layers to 0, so every block computes its own. The pinned checkpoint sets 2: its blocks 2 and 3 attend with blocks 0's and 1's keys.

The readout is softcapped, and the norm gain is the weight

logits is 30 · tanh(lm_head.output / 30) on 12B; project_on_vocab applies the final norm, lm_head and the cap, so a lens at the last block reproduces logits. The vocabulary is 262,144 tokens and lm_head is embed_tokens' weight. Gemma4UnifiedRMSNorm multiplies by weight, not 1 + weight: folding the final norm into lm_head uses model.norm._module.weight as it is.

Every checkpoint is a multimodal wrapper

The released repos are gemma4_unified checkpoints: the text-generation task builds Gemma4UnifiedForConditionalGeneration, nnterp picks this family from config.text_config, and the text stack sits under model.model.language_model with lm_head at the root. Gemma4UnifiedForCausalLM is built only from a gemma4_unified_text config, as the suite's tiny checkpoint is.

Vision notes

· Gemma 4 unified embedder · Gemma 4 unified

The Gemma 4 unified embedder vision encoder

The embedder is five modules and no blocks

h = patch_dense(patch_ln1(pixels))       # patch_embeddings
h = pos_norm(patch_ln2(h) + pos_x + pos_y)
out = multimodal_embedder(h)             # projector: RMS norm, linear

model.vision is model.model.embed_vision, a Vision with no layers: vision.num_layers is 0 and vision.num_heads, head_dim and intermediate_size raise Unavailable. Each row of pixel_values is one merged patch of 48 × 48 pixels, 6912 values, and patch_dense maps it straight to vision.vision_hidden_size (mm_embed_dim, 3840 on 12B). patch_embeddings and tower_output are the embedder's only values.

The padded rows run through the embedder

The processor fits each image into at most max_soft_tokens merged patches (280) and pads the row to that length. The padded rows are rows of patch_embeddings, tower_output and projector.output; the wrapper drops them before the scatter. Select the image's rows with the processor's positions:

encoding = model.processor(text=prompt, images=[image], return_tensors="pt")
valid = (encoding["image_position_ids"] != -1).all(-1)    # [images, 280]

with model.trace(dict(encoding)):
    states = model.vision.tower_output.save()

states[valid]                           # the image's merged patches

Gemma 4 unified

image_features is the projector's output without the padded rows

model.projector is embed_vision.multimodal_embedder, and it receives vision.tower_output itself: projector.input is [images, 280, mm_embed_dim], padded rows included. The wrapper keeps the rows where image_position_ids is not (-1, -1) and scatters them, so vision.image_features == projector.output[valid], and layers[0].input[vision.image_token_mask] == vision.image_features holds exactly. On the pinned tiny wrapper:

valid = (encoding["image_position_ids"] != -1).all(-1)

with model.trace(dict(encoding)):
    mask = model.vision.image_token_mask.save()
    projected = model.projector.output.save()
    features = model.vision.image_features.save()
    first = model.layers[0].input.save()

torch.equal(projected[valid], features)   # True
torch.equal(first[mask], features)        # True: the scatter

A write to tower_output at an image's rows reaches the text model; at the padded rows it reaches nothing.

One image token per merged patch

An image is as many image tokens as it has merged patches, at most 280. With google/gemma-4-12B's processor a square image, 224 × 224 or 2000 × 2000, is 256 image tokens, a 1024 × 768 image 266 and a 96 × 48 image 253. Assign one image's features into another's run only when the two give the same count.

Image tokens attend to each other

12B's text config sets use_bidirectional_attention to "vision": an image's tokens attend to every token of the same image, before and after them, within the window on a sliding block. The wrapper builds that mask from the processor's mm_token_type_ids, so pass the whole encoding (model.trace(prompt, images=[image]) does). Unlike Gemma 4's 26B-A4B and 31B, the full blocks get the same bidirectional image blocks as the sliding ones. The first step of generate builds the same masks as a trace (with the default dynamic cache), and the steps after it attend causally.

The base checkpoint has no chat template

google/gemma-4-12B ships no chat template: put the processor's image token in the prompt yourself, f"{model.processor.image_token} What is this?". Audio enters through model.model.embed_audio, another embedder without blocks, which keeps its native name.

From the family module

04 · nnterp/families/gemma4_unified_text.py

Gemma 4 unified, text (Gemma4UnifiedForCausalLM, model_type gemma4_unified_text): Gemma-4-12B.

Gemma-4's tree and block without per-layer embeddings or a mixture of experts: the sandwich, then out = (x + post_attention_layernorm(...) + post_feedforward_layernorm(...)) * layer_scalar, in place. The contributions are the post-norms' outputs, served unscaled, so the identity is (input + attention_output + mlp_output) * layer_scalar == layer_output. KV sharing, attention_k_eq_v and the per-layer head_dim are Gemma-4's (gemma4_text), and so is the attention envoy (keys and values served as copies private to the block), and so are the root's head_dim and num_kv_heads: the config's top-level values. A gemma4_unified checkpoint (Gemma4UnifiedForConditionalGeneration) keeps the text stack at model.language_model, so RENAME carries both spellings.

The wrapper is encoder-free: no vision tower, no blocks. model.embed_vision embeds the processor's merged patches (model_patch_size pixels square) directly: patch_ln1, patch_dense, patch_ln2, a factorized 2D position embedding, pos_norm, then multimodal_embedder (an RMS norm and a linear onto the text width, the same projection Gemma-4's embed_vision is). So vision is the embedder (a Vision with no layers), vision.patch_embed its patch_dense, and projector its multimodal_embedder, the module mapping onto the text width as on every other wrapper. The processor pads each image's patches to max_soft_tokens rows at position (-1, -1), and the embedder runs on every row; the wrapper strips the padded rows of the projector's output before scattering, so vision.image_features is read at the scatter (the first inputs_embeds.masked_scatter of Gemma4UnifiedModel.forward, an ImageScatter). Audio is embedded the same way, without blocks: model.embed_audio keeps its native name, and there is no audio_tower.