Repository: huggingface/transformers
Pull request: Fix ShieldGemma2 non-reproducible outputs by adding _tied_weights_keys

## Summary

- `ShieldGemma2ForImageClassification` was missing `_tied_weights_keys`, so `model.lm_head.weight` was randomly re-initialized on every `from_pretrained` call instead of being tied to `embed_tokens.weight`.
- This caused non-deterministic classification outputs across model loads despite identical inputs.
- Fix: add `_tied_weights_keys = {"model.lm_head.weight": "model.model.language_model.embed_tokens.weight"}` to the class, matching the dict-based pattern used by `Gemma3ForConditionalGeneration`.

## Root Cause

The checkpoint has `text_config.tie_word_embeddings = True`, so `lm_head.weight` is not saved separately in the checkpoint — it is expected to be restored by the weight-tying mechanism. The dict-based `_tied_weights_keys` attribute tells `from_pretrained`'s loading machinery exactly which parameter is tied and what its source is, so it can restore it correctly via `tie_weights()`.

This regression was introduced when the explicit `tie_weights()` override was removed in #41580 without adding the equivalent `_tied_weights_keys` declaration.

The load report shows the symptom clearly:
```
Key                  | Status  |
---------------------+---------+
model.lm_head.weight | MISSING |
```

## Reproduction

```python
import torch
from transformers import AutoProcessor, ShieldGemma2ForImageClassification
from PIL import Image
import requests

model_id = "google/shieldgemma-2-4b-it"
processor = AutoProcessor.from_pretrained(model_id)
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"
image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
inputs = processor(images=[image], return_tensors="pt")
```

**Before fix — different `lm_head.weight` sum and wildly varying probabilities on each load:**

| Load | lm_head.weight sum | Result |
|------|-------------------|--------|
| 1 | 568.72 | bee/dangerous → Safe:0.0, Unsafe:1.0 |
| 2 | 359.00 | bee/dangerous → Safe:0.0003, Unsafe:1.0 |
| 3 | -883.37 | bee/dangerous → Safe:1.0, Unsafe:0.0 |

**After fix — all 3 loads produce identical results** with `lm_head.weight sum=-105549.42` and zero variance.

Full discussion with benchmarks: https://huggingface.co/google/shieldgemma-2-4b-it/discussions/10

## Test plan

- [x] Load `google/shieldgemma-2-4b-it` multiple times and verified `model.model.lm_head.weight` has the same sum/norm across all loads after fix.
- [x] Confirmed `model.model.lm_head.weight.data_ptr() == model.model.get_input_embeddings().weight.data_ptr()` after loading (weights share memory).

Changed files:
filename: src/transformers/models/shieldgemma2/configuration_shieldgemma2.py
patch: @@ -30,7 +30,8 @@ class ShieldGemma2Config(PreTrainedConfig):
         The begin-of-image token index to wrap the image prompt.
     eoi_token_index (`int`, *optional*, defaults to 256000):
         The end-of-image token index to wrap the image prompt.
-
+    tie_word_embeddings (`bool`, *optional*):
+            Whether to tie the word embeddings. Defaults to the value of `text_config.tie_word_embeddings` if not set.
     Example:
 
     ```python
@@ -69,6 +70,7 @@ def __init__(
         eoi_token_index: int = 256_000,
         image_token_index: int = 262_144,
         initializer_range: float = 0.02,
+        tie_word_embeddings=None,
         **kwargs,
     ):
         if isinstance(vision_config, dict):
@@ -92,7 +94,9 @@ def __init__(
         self.eoi_token_index = eoi_token_index
         self.image_token_index = image_token_index
         self.initializer_range = initializer_range
-
+        if tie_word_embeddings is None:
+            tie_word_embeddings = getattr(self.text_config, "tie_word_embeddings", True)
+        self.tie_word_embeddings = tie_word_embeddings
         super().__init__(**kwargs)