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 / Per-Layer-Embedding shapes and emitted zero
model.vision_tower.* / model.embed_vision.* keys — even though
GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING already 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 change adds the vision block to GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE:

• Gated on vision_config presence, matching the mapping's
  use_multimodal and vcfg gate (no effect on text-only exports).
• Mirrors the already-landed GEMMA4_HF_WEIGHTS_TO_SHAPE (26B/31B) vision
  convention: each clipped-linear vision projection emits a plain
  .linear.weight — no clip-bound keys, because the small param map does not
  map input_min/input_max/output_min/output_max.
• Emits attention q/k/v/o_proj.linear.weight + q_norm/k_norm, the four
  layernorms, gate/up/down_proj.linear.weight, the patch-embedder
  input_proj + position_embedding_table, and the embed_vision projection.
• std_scale / std_bias are emitted only under standardize=true; E2B/E4B
  ship standardize=false and store no std keys (consistent with the mapping's
  placeholder path).

The 26B/31B (GEMMA4_HF_WEIGHTS_TO_SHAPE) and text-only paths are untouched.

Follow-up (not in this PR to keep it surgical): the two
# TODO: gemma4-small multimodal not yet supported — ... are dead comments in
param_mapping.py are now stale and can be dropped in a separate cleanup.

Tests

Added tests/unit/hf_shape_test.py (pytest.mark.decoupled_target, CPU-only, no
checkpoint / no gated weights — builds a config dict shaped like the HF
gemma4-e2b config):

• test_shape_map_covers_every_mapped_vision_target — fail-closed: every HF
  target GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING produces must exist in
  GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE. This is the exact invariant _process
  enforces; it fails on main (missing vision keys) and passes with this change.
• test_vision_tower_keys_present — canonical vision entry / projector / per-layer keys.
• test_no_clip_bound_keys — no input_min/max/output_min/output_max shape keys.
• test_std_keys_only_when_standardize — std keys appear only under standardize=true.
• test_text_only_config_emits_no_vision_keys — no regression to the text-only path.

Reproduce:

python -m pytest tests/unit/hf_shape_test.py -v

Validation performed on this branch (grr-maxtext image):

• The fail-closed cross-check reports 0 missing vision targets for both
  standardize=false (42/42 vision targets) and standardize=true (44/44);
  the same check fails on main (vision keys absent) — confirming the fix.
• pytest tests/unit/hf_shape_test.py → 5 passed.
• pyink --pyink-indentation=2 --line-length=122 --check → clean.
• pylint --rcfile=pylintrc --disable=R0401,R0917,W0201,W0613 → 10.00/10.
• codespell → clean.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

---

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.