Gemma-4 small (E2B/E4B): emit vision-tower shapes in HF export map

[lokic233]

Description

GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE (used by gemma4-e2b / gemma4-e4b) returned only language-model / PLE
shapes and emitted zero model.vision_tower.* / model.embed_vision.* keys, while
GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING maps the full vision encoder. Because
checkpoint_conversion/utils/utils.py::_process requires every mapped HF target to be present in the shape
map, to_huggingface use_multimodal=true on a converted Gemma-4 E2B/E4B checkpoint died on the first vision
tensor:

ValueError: HF path 'model.vision_tower.patch_embedder.input_proj.weight' not found in hf_shape_map

This adds the vision block to GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE, gated on vision_config presence:

• Mirrors the GEMMA4_HF_WEIGHTS_TO_SHAPE (26B/31B) convention: each clipped-linear vision projection emits
  a .linear.weight; attention q/k/v/o + q_norm/k_norm, four layernorms, gate/up/down MLP, patch-embedder
  input_proj + position_embedding_table, embed_vision projection.
• Clip bounds: also emits the four scalar activation clip bounds per clipped projection
  (input_min/input_max/output_min/output_max) that GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING maps under
  use_clipped_linears_for_vit (the E2B/E4B image-parity path). Bounds are scalar [] (a rank-1 [1] shape
  makes the HF loader reinitialize the bound to a non-finite sentinel). The shape fn receives only the HF
  config (not the MaxText config), and _process iterates the param map, so emitting the bound shapes
  whenever the vision block is present is a harmless superset when clipped-linears is off and required
  coverage when it is on.
• std_scale / std_bias are emitted only under standardize=true (E2B/E4B ship standardize=false).

Scope: this PR completes the vision HF export shape map for gemma4-small. Shared-layer text K/V /
k_norm are intentionally omitted on the HF side (HF transformers Gemma4TextAttention itself omits them for
KV-shared layers), so this is HF-faithful; the serving-side shared-layer checkpoint-completeness invariant is
a separate concern and is not part of this PR.

FIXES: to_huggingface multimodal export for gemma4-e2b / gemma4-e4b.

Tests

tests/unit/hf_shape_test.py (CPU, no gated checkpoint):

• test_shape_map_covers_every_mapped_vision_target — fail-closed: every mapped vision HF target has a shape.
• test_shape_map_covers_clipped_mapping_targets — fail-closed incl. the 448 clip bounds when the param map maps them.
• test_clip_bound_keys_present_and_scalar — clip bounds present and scalar [].
• test_vision_tower_keys_present, test_std_keys_only_when_standardize, test_text_only_config_emits_no_vision_keys.

Validated composed with #4790 (which adds the clip-bound mapping): mapping<->shape coverage = 0 missing
across clipped ON/OFF x standardize ON/OFF. pyink/pylint (10.00/10)/codespell clean.

Related work

• Gemma-4 E2B: enable image (multimodal) E2E parity — vision clipped-linears + padded-patch masking + PLE substitution (E4B vision-tower validated; E4B E2E pending) #4790 — E2B multimodal enablement + GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING (incl. clip-bound mapping this PR covers).
• Gemma-4 vision: correct std_bias/std_scale classification per HF standardize (fidelity fix) #4807 — vision std_bias/std_scale fidelity.

Checklist

• Self-review. - [X] Comments in hard-to-understand areas. - [ ] E2E tests (CPU coverage invariant; full checkpoint export needs gated weights). - [X] Docs n/a.

---

Authored with Navi on behalf of Loki Chen (@lokic233).

[gemini-code-assist]

Code Review

This pull request adds support for the vision tower (multimodal E2B / E4B) in the GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE function within hf_shape.py. It maps the shapes of the vision tower components, including the patch embedder, position embedding table, standardization parameters, projection layers, and encoder layers, when a vision_config is present. Additionally, a comprehensive suite of unit tests has been introduced in hf_shape_test.py to validate this mapping behavior. There are no review comments, and I have no additional feedback to provide.