nnterp · encyclopedia gemma4_text · nnterp.families.gemma4_text

family gemma4_text

Gemma 4

A sandwich block with normed queries, keys and values, whose sum a learned per-block scalar multiplies; on E2B and E4B per-layer embeddings add a third term and the last blocks borrow keys and values from earlier ones.

Sandwich normsPer-block sizesScaled residual addsQuery/key normsSliding windowProportional rotaryBorrowed keys and valuesScaled embeddingsSoftcapped logitsMulti-query attentionVision-languageVariable 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× 35
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitssoftcapped
i = 0

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

The vision encoder

· Gemma 4 ViT · Gemma4VisionModel

The image's path into the text model, then one block of the vision encoder. Hover any part for its nnterp name; click to pin the card.

imagepixel_values
→
patch_embedpatch_embeddings
→
layers× 16 · tower_output
→
projectorinput: the pooler's soft tokens, after vision.tower_output
→
image_featuresat the scatter
→
layers[0].inputat image_token_mask
rows
One row per image: its patches, then padding up to max_soft_tokens * pooling_kernel_size**2 rows (2520 by default); the processor's image_position_ids is (-1, -1) on the padded rows.
positions
patch_embeddings is patch_embedder.input_proj's output; patch_embedder then adds a learned x and a learned y embedding per patch (zero on the padded rows), and every attention rotates queries and keys by the patch's 2D position.
masking
No causal mask: a row attends to every patch of its image; the padded rows are masked as keys.
final norm
None: vision.tower_output is the encoder's output, padded rows included, and the vision encoder's pooler follows it, outside the blocks.

identity vision.layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

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

model Gemma4ForConditionalGenerationlayoutwhat 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] Gemma4TextDecoderLayerlayoutwhat 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
per_layer_output Residual [batch seq hidden_size] What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output

the child module post_per_layer_input_norm's output
model.layers[i].self_attn Gemma4TextAttentionlayoutwhat 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 Gemma4TextMLPlayoutwhat 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 (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output

the sibling module post_feedforward_layernorm's output
⚠︎router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

the child module router.proj's output
⚠︎expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2
⚠︎expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1
⚠︎expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1
⚠︎routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output
⚠︎shared_expert_output Residual [batch seq hidden_size] The dense MLP's output, before its post-norm: the shared expert beside the mixture

the module's own output

the vision encoder · Gemma 4 ViT

model.vision Gemma4VisionModellayoutwhat 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 last block's stream after the tower's final norm, before any pooling or adapter

the child module encoder's output
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
model.vision.layers[i] Gemma4VisionEncoderLayerlayoutwhat it isread at
layer_input Patches [images patches vision_hidden_size] The tower's stream entering the block

the module's own input
attention_norm_output Patches [images patches vision_hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Patches [images patches vision_hidden_size] The tower's stream after the attention add

the child module pre_feedforward_layernorm's input
mlp_norm_output Patches [images patches vision_hidden_size] The normed stream entering the MLP

the child module pre_feedforward_layernorm's output
mlp_input Patches [images patches vision_hidden_size] What the MLP is called with

the child module mlp's input
layer_output Patches [images patches vision_hidden_size] The tower's stream leaving the block

the module's own output
model.vision.layers[i].self_attn Gemma4VisionAttentionlayoutwhat 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 Patches [images patches vision_hidden_size] What the attention adds to the tower's 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.vision.layers[i].mlp Gemma4VisionMLPlayoutwhat 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

the child module down_proj's input
mlp_output Patches [images patches vision_hidden_size] What the MLP adds to the tower's stream: the post-feedforward norm's output

the sibling module post_feedforward_layernorm's output

Printout

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

Gemma4ForConditionalGeneration(
  (lm_head): Linear(in_features=1536, out_features=262144, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Gemma4TextScaledWordEmbedding(262144, 1536, padding_idx=0)
  (layers): ModuleList(
    (0-3): 4 x Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=2048, bias=False)
        (q_norm): Gemma4RMSNorm()
        (k_norm): Gemma4RMSNorm()
        (v_norm): Gemma4RMSNorm()
        (k_proj): Linear(in_features=1536, out_features=256, bias=False)
        (v_proj): Linear(in_features=1536, out_features=256, bias=False)
        (o_proj): Linear(in_features=2048, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (up_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (down_proj): Linear(in_features=6144, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (4): Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=4096, bias=False)
        (q_norm): Gemma4RMSNorm()
        (k_norm): Gemma4RMSNorm()
        (v_norm): Gemma4RMSNorm()
        (k_proj): Linear(in_features=1536, out_features=512, bias=False)
        (v_proj): Linear(in_features=1536, out_features=512, bias=False)
        (o_proj): Linear(in_features=4096, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (up_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (down_proj): Linear(in_features=6144, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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-8): 4 x Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=2048, bias=False)
        (q_norm): Gemma4RMSNorm()
        (k_norm): Gemma4RMSNorm()
        (v_norm): Gemma4RMSNorm()
        (k_proj): Linear(in_features=1536, out_features=256, bias=False)
        (v_proj): Linear(in_features=1536, out_features=256, bias=False)
        (o_proj): Linear(in_features=2048, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (up_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (down_proj): Linear(in_features=6144, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (9): Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=4096, bias=False)
        (q_norm): Gemma4RMSNorm()
        (k_norm): Gemma4RMSNorm()
        (v_norm): Gemma4RMSNorm()
        (k_proj): Linear(in_features=1536, out_features=512, bias=False)
        (v_proj): Linear(in_features=1536, out_features=512, bias=False)
        (o_proj): Linear(in_features=4096, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (up_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (down_proj): Linear(in_features=6144, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (10-13): 4 x Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=2048, bias=False)
        (q_norm): Gemma4RMSNorm()
        (k_norm): Gemma4RMSNorm()
        (v_norm): Gemma4RMSNorm()
        (k_proj): Linear(in_features=1536, out_features=256, bias=False)
        (v_proj): Linear(in_features=1536, out_features=256, bias=False)
        (o_proj): Linear(in_features=2048, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (up_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (down_proj): Linear(in_features=6144, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (14): Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=4096, bias=False)
        (q_norm): Gemma4RMSNorm()
        (k_norm): Gemma4RMSNorm()
        (v_norm): Gemma4RMSNorm()
        (k_proj): Linear(in_features=1536, out_features=512, bias=False)
        (v_proj): Linear(in_features=1536, out_features=512, bias=False)
        (o_proj): Linear(in_features=4096, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (up_proj): Linear(in_features=1536, out_features=6144, bias=False)
        (down_proj): Linear(in_features=6144, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (15-18): 4 x Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=2048, bias=False)
        (q_norm): Gemma4RMSNorm()
        (o_proj): Linear(in_features=2048, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (up_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (19): Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=4096, bias=False)
        (q_norm): Gemma4RMSNorm()
        (o_proj): Linear(in_features=4096, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (up_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (20-23): 4 x Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=2048, bias=False)
        (q_norm): Gemma4RMSNorm()
        (o_proj): Linear(in_features=2048, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (up_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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): Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=4096, bias=False)
        (q_norm): Gemma4RMSNorm()
        (o_proj): Linear(in_features=4096, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (up_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (25-28): 4 x Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=2048, bias=False)
        (q_norm): Gemma4RMSNorm()
        (o_proj): Linear(in_features=2048, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (up_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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): Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=4096, bias=False)
        (q_norm): Gemma4RMSNorm()
        (o_proj): Linear(in_features=4096, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (up_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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-33): 4 x Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=2048, bias=False)
        (q_norm): Gemma4RMSNorm()
        (o_proj): Linear(in_features=2048, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (up_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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
    )
    (34): Gemma4TextDecoderLayer(
      (self_attn): Gemma4TextAttention(
        (q_proj): Linear(in_features=1536, out_features=4096, bias=False)
        (q_norm): Gemma4RMSNorm()
        (o_proj): Linear(in_features=4096, out_features=1536, 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): Gemma4TextMLP(
        (gate_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (up_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (down_proj): Linear(in_features=12288, out_features=1536, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP (and, on a mixture block, the experts) adds to the residual stream: the post-feedforward norm's output
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (shared_expert_output) -> Residual [batch seq hidden_size]: The dense MLP's output, before its post-norm: the shared expert beside the mixture
        (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
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (input_layernorm): Gemma4RMSNorm()
      (post_attention_layernorm): Gemma4RMSNorm()
      (pre_feedforward_layernorm): Gemma4RMSNorm()
      (post_feedforward_layernorm): Gemma4RMSNorm()
      (act_fn): GELUTanh()
      (per_layer_input_gate): Linear(in_features=1536, out_features=256, bias=False)
      (per_layer_projection): Linear(in_features=256, out_features=1536, bias=False)
      (post_per_layer_input_norm): Gemma4RMSNorm()
      (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
      (per_layer_output) -> Residual [batch seq hidden_size]: What the per-layer-embedding branch adds to the residual stream: the post-per-layer-input norm's output
      (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): Gemma4RMSNorm()
  (vision): Gemma4VisionModel(
    (encoder): Gemma4VisionEncoder(
      (rotary_emb): Gemma4VisionRotaryEmbedding()
      (layers): ModuleList(
        (0-15): 16 x Gemma4VisionEncoderLayer(
          (self_attn): Gemma4VisionAttention(
            (q_proj): Gemma4ClippableLinear(
              (linear): Linear(in_features=768, out_features=768, bias=False)
            )
            (k_proj): Gemma4ClippableLinear(
              (linear): Linear(in_features=768, out_features=768, bias=False)
            )
            (v_proj): Gemma4ClippableLinear(
              (linear): Linear(in_features=768, out_features=768, bias=False)
            )
            (o_proj): Gemma4ClippableLinear(
              (linear): Linear(in_features=768, out_features=768, bias=False)
            )
            (q_norm): Gemma4RMSNorm()
            (k_norm): Gemma4RMSNorm()
            (v_norm): Gemma4RMSNorm()
            (attention_output) -> Patches [images patches vision_hidden_size]: What the attention adds to the tower's stream: the post-attention norm's output
            (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
            (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
            (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_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
            (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): Gemma4VisionMLP(
            (gate_proj): Gemma4ClippableLinear(
              (linear): Linear(in_features=768, out_features=3072, bias=False)
            )
            (up_proj): Gemma4ClippableLinear(
              (linear): Linear(in_features=768, out_features=3072, bias=False)
            )
            (down_proj): Gemma4ClippableLinear(
              (linear): Linear(in_features=3072, out_features=768, bias=False)
            )
            (act_fn): GELUTanh()
            (mlp_output) -> Patches [images patches vision_hidden_size]: What the MLP adds to the tower's 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
          )
          (input_layernorm): Gemma4RMSNorm()
          (post_attention_layernorm): Gemma4RMSNorm()
          (pre_feedforward_layernorm): Gemma4RMSNorm()
          (post_feedforward_layernorm): Gemma4RMSNorm()
          (layer_mid) -> Patches [images patches vision_hidden_size]: The tower's stream after the attention add
          (mlp_norm_output) -> Patches [images patches vision_hidden_size]: The normed stream entering the MLP
          (layer_input) -> Patches [images patches vision_hidden_size]: The tower's stream entering the block
          (attention_norm_output) -> Patches [images patches vision_hidden_size]: The normed stream entering the attention
          (mlp_input) -> Patches [images patches vision_hidden_size]: What the MLP is called with
          (layer_output) -> Patches [images patches vision_hidden_size]: The tower's stream leaving the block
        )
      )
    )
    (pooler): Gemma4VisionPooler()
    (patch_embed): Linear(in_features=768, out_features=768, bias=False)
    (layers): ModuleList(
      (0-15): 16 x Gemma4VisionEncoderLayer(
        (self_attn): Gemma4VisionAttention(
          (q_proj): Gemma4ClippableLinear(
            (linear): Linear(in_features=768, out_features=768, bias=False)
          )
          (k_proj): Gemma4ClippableLinear(
            (linear): Linear(in_features=768, out_features=768, bias=False)
          )
          (v_proj): Gemma4ClippableLinear(
            (linear): Linear(in_features=768, out_features=768, bias=False)
          )
          (o_proj): Gemma4ClippableLinear(
            (linear): Linear(in_features=768, out_features=768, bias=False)
          )
          (q_norm): Gemma4RMSNorm()
          (k_norm): Gemma4RMSNorm()
          (v_norm): Gemma4RMSNorm()
          (attention_output) -> Patches [images patches vision_hidden_size]: What the attention adds to the tower's stream: the post-attention norm's output
          (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
          (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
          (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_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
          (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): Gemma4VisionMLP(
          (gate_proj): Gemma4ClippableLinear(
            (linear): Linear(in_features=768, out_features=3072, bias=False)
          )
          (up_proj): Gemma4ClippableLinear(
            (linear): Linear(in_features=768, out_features=3072, bias=False)
          )
          (down_proj): Gemma4ClippableLinear(
            (linear): Linear(in_features=3072, out_features=768, bias=False)
          )
          (act_fn): GELUTanh()
          (mlp_output) -> Patches [images patches vision_hidden_size]: What the MLP adds to the tower's 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
        )
        (input_layernorm): Gemma4RMSNorm()
        (post_attention_layernorm): Gemma4RMSNorm()
        (pre_feedforward_layernorm): Gemma4RMSNorm()
        (post_feedforward_layernorm): Gemma4RMSNorm()
        (layer_mid) -> Patches [images patches vision_hidden_size]: The tower's stream after the attention add
        (mlp_norm_output) -> Patches [images patches vision_hidden_size]: The normed stream entering the MLP
        (layer_input) -> Patches [images patches vision_hidden_size]: The tower's stream entering the block
        (attention_norm_output) -> Patches [images patches vision_hidden_size]: The normed stream entering the attention
        (mlp_input) -> Patches [images patches vision_hidden_size]: What the MLP is called with
        (layer_output) -> Patches [images patches vision_hidden_size]: The tower's stream leaving the block
      )
    )
    (tower_output) -> Patches [images patches vision_hidden_size]: The last block's stream after the tower's final norm, before any pooling or adapter
    (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): Gemma4MultimodalEmbedder(
    (embedding_projection): Linear(in_features=768, out_features=1536, bias=False)
    (embedding_pre_projection_norm): Gemma4RMSNorm()
  )
  (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_layers35
model.hidden_size1536
model.num_heads8
model.num_kv_heads1
model.head_dim256
model.qk_head_dim256
model.vocab_size262144
model.intermediate_size6144
architecturesGemma4ForConditionalGeneration
num_hidden_layers35
hidden_size1536
intermediate_size6144
num_attention_heads8
num_key_value_heads1
head_dim256
vocab_size262144
max_position_embeddings131072
sliding_window512
hidden_activationgelu_pytorch_tanh
rms_norm_eps1e-06
final_logit_softcapping30.0
tie_word_embeddingsTrue
dtypebfloat16
num_kv_shared_layers20
attention_biasFalse
enable_moe_blockFalse
hidden_size_per_layer_input256
attention_k_eq_vFalse
Vision Config

Vision Config

model.vision.num_layers16
model.vision.vision_hidden_size768
model.vision.num_heads12
model.vision.head_dim64
model.vision.intermediate_size3072
model.vision.patch_size16
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; borrowed on the last blocks
v   = v_norm(v_proj(a))               # per head, no gain; k_proj(a) on 26B-A4B's and 31B'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)))
p   = per_layer_projection(gelu(per_layer_input_gate(h)) * per_layer_input)
h   = h + post_per_layer_input_norm(p)
out = h * layer_scalar

Four RMSNorms around the sublayers as on Gemma 2 and 3, a fifth after the per-layer branch on E2B and E4B, and three inside the attention (only q_norm on a block that borrows keys and values). A mixture block on 26B-A4B adds three more around the experts. post_attention_layernorm follows the attention; the MLP's input norm is pre_feedforward_layernorm. The per-layer branch has no module of its own: its three modules are children of the block, and per_layer_output is a value of layers[i]. 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, mlp_output and per_layer_output are the post-norms' outputs, before layer_scalar. The identity is exact in float32 on every E2B block:

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()
    ple = model.layers[1].per_layer_output.save()
    out = model.layers[1].layer_output.save()

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

No scalar is 1 on E2B: they run from 0.019 (block 0) to 0.88 (block 30), with 0.074 and 0.049 at blocks 13 and 14 and 0.16 at the last block. A vector added to mlp_output reaches layer_output multiplied by the block's scalar; one added to layer_output does not.

A term's weight in the last stream is the product of the scalars from its block on

Block i's terms are multiplied by its own scalar and by every later one; the embeddings by all of them. In float32 the last block's output is that weighted sum:

scalars = torch.cat([layer._module.layer_scalar for layer in model.layers])
weight = scalars.flip(0).cumprod(0).flip(0)

added = []
with model.trace(prompt):
    emb = model.token_embeddings.save()
    for layer in model.layers:
        attn = layer.self_attn.attention_output
        added.append((attn + layer.mlp.mlp_output + layer.per_layer_output).save())
    final = model.layers[-1].layer_output.save()

total = weight[0] * emb + sum(w * a for w, a in zip(weight, added))
torch.testing.assert_close(total, final)

On E2B the weights are 1.9e-13 for block 0 (and the embeddings), 1.8e-6 for block 14, 1.2e-3 for block 20, 0.06 for block 30 and 0.16 for block 34. The final norm divides out the common scale, so only the ratios count: block 34's terms outweigh block 0's by about 10¹². At the last position of a 9-token prompt, block 0's attention_output has norm 1476 and block 34's mlp_output 122; weighted, they are 2.8e-10 and 20 in a last stream of norm 28. Direct logit attribution multiplies each term by its weight first; the early blocks act on the logits only through the blocks that read the stream after them.

Per-layer embeddings add a third term

E2B and E4B (hidden_size_per_layer_input 256) give every block a 256-wide input of its own. The model looks each token up in model.model.language_model.embed_tokens_per_layer ([262144, 35 × 256] on E2B, scaled by 16: 2.35B of the language model's 4.63B parameters), adds a normed projection of token_embeddings, and scales the sum by 2^-1/2; block i receives slice i as its second argument, model.layers[i].inputs[0][1] ([batch, seq, 256]). The block gates it by gelu(per_layer_input_gate(h)), projects it to the hidden size and norms it: per_layer_output. The model depends on it: zeroing it on every block takes p(" Paris") after "The Eiffel Tower is in the city of" from 0.91 to 6e-10 on E2B.

with model.trace(prompt):
    for layer in model.layers:
        layer.per_layer_output[:] = 0
    logits = model.logits.save()

Patching token_embeddings does not swap the token

The per-layer table is looked up by token id, not from token_embeddings, so writing another token's embedding into token_embeddings changes only the projected half of every block's per-layer input. On E2B, putting " Rome"'s embedding in place of " Paris"'s at the end of "The capital of France is Paris" leaves a KL of 0.51 to the " Rome" run. Patching the table's output as well reproduces that run exactly:

with model.trace(rome):
    emb = model.token_embeddings.save()
    table = model.model.language_model.embed_tokens_per_layer.output.save()

with model.trace(paris):
    model.token_embeddings[:, -1] = emb[:, -1]
    model.model.language_model.embed_tokens_per_layer.output[:, -1] = table[:, -1]
    logits = model.logits.save()

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; q_norm.output is the normed queries before it. E2B has one key/value head for its 8 query heads, so an edit to attention_keys or attention_values reaches all 8; E4B has two, each serving 4.

Sliding and full blocks differ in width and rotary

config.layer_types makes every fifth block full on E2B (4, 9, ..., 34) and every sixth on E4B, 26B-A4B and 31B; the last block is full on all four. The window is 512 tokens on E2B and E4B (a query attends to itself and the 511 before it) and 1024 on 26B-A4B and 31B. Full blocks have 512-wide heads against the sliding blocks' 256, on every size; 26B-A4B and 31B also give them fewer key/value heads (2 against 8, 4 against 16). The root's sizes are the sliding blocks'; each block's own are on its attention and MLP:

model.head_dim, model.layers[4].self_attn.head_dim        # 256, 512 on E2B
model.layers[4].self_attn.num_heads                       # 8: the same on every block
model.layers[20].mlp.intermediate_size                    # 12288 on E2B; 6144 at the root

Sliding blocks rotate all 256 dimensions with base 10,000. Full blocks use proportional rotary: only dimensions 0 to 63 and 256 to 319 of the 512 turn, at frequencies spaced over the whole head with base 1,000,000 (the slowest has a wavelength of 188 tokens); the other 384 carry no position.

26B-A4B and 31B take the full blocks' values from the key projection

26B-A4B and 31B set attention_k_eq_v: their 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. The sliding blocks, and every block on E2B and E4B, have their own v_proj.

Blocks 15 to 34 attend with borrowed keys and values on E2B

The last num_kv_shared_layers blocks (20 of 35 on E2B, 18 of 42 on E4B, none on 26B-A4B and 31B) have no k_proj, v_proj, k_norm or v_norm. On E2B the sliding ones attend with block 13's keys and values and the full ones with block 14's. Their attention_keys are block 13's or 14's, as a copy private to the block: zeroing block 13's leaves block 20's as they were, while an edit of block 13's k_proj output reaches every sliding block that borrows from it.

with model.trace(prompt):
    source = model.layers[13].self_attn.attention_keys.save()
    borrowed = model.layers[20].self_attn.attention_keys.save()

assert torch.equal(source, borrowed)

skip_layers over block 13 or 14 leaves the borrowers nothing to read and fails with a KeyError inside transformers.

Eager and sdpa compute the same attention

Gemma 4 has no attention softcap, and the scale is 1.0 under every implementation. On E2B in float32 a default (sdpa) load's logits are within 2.1e-5 of an eager load's on a 9-token prompt and 3.5e-5 on an 881-token one. attn_implementation="eager" is needed only for the values read inside the attention function (the ones marked ⚠).

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

logits is 30 · tanh(lm_head.output / 30) on every released size; 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 on every size, and lm_head is embed_tokens' weight. Unlike Gemma 2 and 3, Gemma4RMSNorm multiplies by weight, not 1 + weight: folding the final norm into lm_head uses model.norm._module.weight as it is (mean 14.2 on E2B, maximum 118.5).

Every checkpoint is a multimodal wrapper

The released repos are gemma4 checkpoints: the text-generation task builds Gemma4ForConditionalGeneration (vision and audio towers on E2B and E4B, vision only on 26B-A4B and 31B), nnterp picks this family from config.text_config, and the text stack sits under model.model.language_model with lm_head at the root. Gemma4ForCausalLM is built only from a gemma4_text config, as the suite's tiny checkpoint is. google/gemma-4-12B is a gemma4_unified checkpoint and loads as gemma4_unified_text: the same block without per-layer embeddings or a mixture of experts.

26B-A4B runs a mixture of experts beside the dense MLP

Under enable_moe_block every block runs 128 experts, 8 per token and each 704 wide, beside the 2112-wide dense MLP. The router reads the stream after the attention add, through its own norm without gain, scales it by router.scale and hidden_size ** -0.5, takes a softmax over the 128 experts and keeps the top 8, renormalised to sum to one and then multiplied by per_expert_scale: expert_weights are those products. The experts read pre_feedforward_layernorm_2 of the same stream. shared_expert_output is the dense MLP's output, mlp.output, and the two meet in three norms:

mlp_output = post_feedforward_layernorm(
    post_feedforward_layernorm_1(shared_expert_output)
    + post_feedforward_layernorm_2(routed_output))

On E2B, E4B and 31B every mixture value is unavailable.

Vision notes

· Gemma 4 ViT · Gemma 4

The Gemma 4 ViT vision encoder

The block is a sandwich

h   = x + post_attention_layernorm(self_attn(input_layernorm(x)))
out = h + post_feedforward_layernorm(mlp(pre_feedforward_layernorm(h)))

Four RMSNorms per block, no per-layer branch and no scalar. attention_output and mlp_output are the post-norms' outputs, so vision.layers[i].input + attention_output + mlp_output == layer_output holds exactly in float32, and an ablation or a steering vector on a sublayer goes on attention_output or mlp_output.

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

The attention norms each head's queries and keys (q_norm, k_norm) and its values (v_norm, without gain), rotates queries and keys by the patch's x and y position, and computes the scores with scaling 1.0. model.vision.head_dim is the config's head_dim: 64 on E2B and E4B, with 12 heads over a 768-wide stream.

The padded rows run through every block

The processor resizes an image within a budget of max_soft_tokens * pooling_kernel_size**2 patches of 16 pixels and pads the row to that length: a square image is 48 × 48 = 2304 patches and 216 padded rows of 2520. The padded rows are masked as keys, so no patch reads them, but they are rows of every block's layer_output and of tower_output, with values. Select the image's patches with the processor's positions:

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

with model.trace(dict(encoding)):
    stream = model.vision.layers[0].layer_output.save()

stream[~padded]                         # the 2304 patches

The pooler makes the soft tokens

After the last block the vision encoder's pooler zeroes the padded rows, averages each 3 × 3 block of patches by position (pooling_kernel_size), multiplies by √vision_hidden_size and drops the padding: 2304 patches become 256 soft tokens, [soft_tokens, vision_hidden_size] flat over the images. These soft tokens are model.projector.input. A write to tower_output at the padded rows changes nothing downstream; at the patches it reaches the text model. On 26B-A4B and 31B (standardize) the pooled tokens are then shifted by std_bias and scaled by std_scale, inside the vision encoder.

Gemma 4

embed_vision projects the pooled soft tokens

model.embed_vision (the projector) receives the pooler's soft tokens, not tower_output: model.projector.input is [soft_tokens, vision_hidden_size], flat over the images. It norms them without a gain (embedding_pre_projection_norm) and projects them to the text model's width (embedding_projection); vision.image_features is model.projector.output, and layers[0].input[vision.image_token_mask] == vision.image_features holds exactly. On the pinned tiny checkpoint:

with model.trace(prompt, images=[image]):
    mask = model.vision.image_token_mask.save()
    out = model.vision.tower_output.save()
    fed = model.projector.input.save()
    features = model.vision.image_features.save()
    first = model.layers[0].input.save()

out.shape[1], fed.shape[0]               # 2520 rows, 256 soft tokens
torch.equal(first[mask], features)       # True: the scatter

256 image tokens for a square image

On google/gemma-4-E2B a 224 × 224 image is 2304 patches and 216 padded rows: under eager every vision encoder block's pattern is (1, 12, 2520, 2520), and the text model receives 256 image tokens, vision.image_features (256, 1536). The trace ran under torch.no_grad() in bfloat16 and peaked at 11 GB. An image of another shape gets another grid in the same 2520 rows (a 96 × 48 image is 2277 patches and 253 image tokens), so the count of image tokens depends on the image: assign one image's features into another's run only when the two give the same count.

On 26B-A4B and 31B an image's tokens attend to each other on the sliding blocks

26B-A4B and 31B set use_bidirectional_attention to "vision": on a sliding block an image's tokens attend to every token of the same image, before and after them, within the window; the full blocks stay causal. The wrapper builds that mask from the processor's mm_token_type_ids, so pass the whole encoding (model.trace(prompt, images=[image]) does). E2B and E4B leave the setting unset and run causal masks throughout.

The base checkpoints have no chat template

google/gemma-4-E2B ships no chat template: put the processor's image token in the prompt yourself, f"{model.processor.image_token} What is this?".

The audio side keeps its native names

E2B and E4B also carry an audio encoder. model.model.audio_tower and model.model.embed_audio keep their native names, and an audio prompt, model.trace(prompt, audio=[waveform]), runs under the same load.

From the family module

04 · nnterp/families/gemma4_text.py

Gemma 4, text (Gemma4ForCausalLM, model_type gemma4_text): E2B, E4B, 26B-A4B and 31B.

model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, pre_feedforward_layernorm, mlp, post_feedforward_layernorm}, norm} and lm_head: Gemma-3's sandwich tree. A gemma4 checkpoint (Gemma4ForConditionalGeneration, the published repos) nests its config's text_config, of this type; the text-generation task builds the wrapper, whose text stack sits at model.language_model.{embed_tokens, layers, norm} with lm_head at the root, so RENAME carries both spellings and whichever the tree has binds.

The block, in order::

x1 = x  + post_attention_layernorm(self_attn(input_layernorm(x)))
x2 = x1 + post_feedforward_layernorm(mlp(pre_feedforward_layernorm(x1)) [+ experts])
p  = gelu(per_layer_input_gate(x2)) * per_layer_input
x3 = x2 + post_per_layer_input_norm(per_layer_projection(p))
out = x3 * layer_scalar          # in place, on x3
  • The contributions are the post-norms' outputs, as on Gemma-2/3. On a mixture-of-experts checkpoint (26B-A4B, enable_moe_block) the block also runs router and experts, direct children of the block, not of mlp: post_feedforward_layernorm_1(mlp(...)) + post_feedforward_layernorm_2(experts(pre_feedforward_layernorm_2(x1))), and post_feedforward_layernorm norms that sum, so mlp_output is the dense MLP and the experts together, what the block adds. The mixture has no module of its own, so mlp (a Moe) hosts its values, reading the block's router and experts through its parent. The router runs on x1, the stream after the attention's add (the block's residual at that point, not its input), with its own norm, and returns probabilities; router_logits are its projection's output (router.proj). shared_expert_output is the dense MLP's output, before its post-norm: mlp_output == post_feedforward_layernorm( post_feedforward_layernorm_1(shared_expert_output) + post_feedforward_layernorm_2(routed_output)). On a dense checkpoint every mixture value is unavailable.
  • Per-layer embeddings (E2B/E4B, hidden_size_per_layer_input): a third add, served as layers[i].per_layer_output (the post-per-layer-input norm's output). per_layer_input is the block's second argument, one slice of the model's per-layer embedding table mixed with a projection of the token embeddings. On a checkpoint without them the value is unavailable.
  • layer_scalar, a per-block buffer, multiplies the sum in place. The contributions are served unscaled, so the identity is (input + attention_output + mlp_output [+ per_layer_output]) * layer_scalar == layer_output.
  • Keys and values that are not the block's own. The last num_kv_shared_layers blocks have no k_proj/v_proj: each borrows the keys and values of the last earlier block of the same kind (sliding or full) before the sharing starts. On a checkpoint with attention_k_eq_v (26B-A4B, 31B) the full-attention blocks have no v_proj: their values are v_norm(k_proj(x)), the keys' projection before k_norm and the rotary embedding. attention_keys and attention_values are what the attention receives on every block, borrowed or not, served as copies private to the block: an edit, in place or by assignment, changes what that block attends with and nothing else, and the blocks that borrow a source block's tensors keep the originals.
  • Per-layer sizes. Sliding and full blocks differ in head_dim (and on 26B-A4B/31B in num_key_value_heads), and on E2B (use_double_wide_mlp) the KV-sharing blocks' MLP is twice intermediate_size wide; the config marks the attention sizes per-layer and refuses a plain config.head_dim. The root's head_dim and num_kv_heads report the config's top-level values (the sliding blocks'); each block's own are on layers[i].self_attn and layers[i].mlp, read off the module.

final_logit_softcapping is in the text config; the root's project_on_vocab reads it there.

On the wrapper the vision tower model.vision_tower is vision (a Vision) and model.embed_vision, the last module before the scatter, is projector. The tower is a ViT over each image's patches: the patch embedder's linear over the processor's flattened patches (patch_embedder.input_proj) is vision.patch_embed, so vision.patch_embeddings is before the 2D position embedding patch_embedder adds, and the blocks (encoder.layers) are vision.layers. Each block is the text block's sandwich, so its contributions are the post-norms' outputs too (VisionAttention, VisionMlp). The processor pads every image's patches to max_soft_tokens * pooling_kernel_size**2 rows with zero pixels at position (-1, -1); the padded rows run through every block (masked as keys only), so they are present in vision.layers[i].layer_output. There is no final norm: the tower's pooler zeroes the padded rows, average-pools the patches pooling_kernel_size by pooling_kernel_size and multiplies by sqrt(vision_hidden_size) (in float32); the tower then strips the padding and, where the vision config sets standardize (26B-A4B, 31B), subtracts std_bias and multiplies by std_scale. So vision.tower_output is the encoder's output (the last block's stream, padded) and projector.input is the pooled, scaled (and standardized) soft tokens, flat over the images. The audio tower (model.audio_tower, a Conformer) and its embedder model.embed_audio keep their native names.