From 16e240a73395fe426d76467d09d12713e8cff28a Mon Sep 17 00:00:00 2001
From: Priyanka Bolem <priyankabolem@gmail.com>
Date: Mon, 22 Sep 2025 16:46:42 -0700
Subject: [PATCH] Fix: add num_hidden_layers property to T5GemmaConfig and add
 test for use_cache

---
 .../models/t5gemma/configuration_t5gemma.py   |  9 ++++++
 .../models/t5gemma/modular_t5gemma.py         |  9 ++++++
 .../models/t5gemma/test_generation_t5gemma.py | 32 +++++++++++++++++++
 3 files changed, 50 insertions(+)
 create mode 100644 tests/models/t5gemma/test_generation_t5gemma.py

diff --git a/src/transformers/models/t5gemma/configuration_t5gemma.py b/src/transformers/models/t5gemma/configuration_t5gemma.py
index 217a24df0417..044ddeddc9f4 100644
--- a/src/transformers/models/t5gemma/configuration_t5gemma.py
+++ b/src/transformers/models/t5gemma/configuration_t5gemma.py
@@ -327,5 +327,14 @@ def get_text_config(self, *args, **kwargs):
         # Always return self, regardless of the decoder option.
         return self
 
+    # Bridge for generation/cache utils which expect `config.num_hidden_layers`.
+    # Prefer a top-level override if present; otherwise use the decoder's count.
+    @property
+    def num_hidden_layers(self):
+        if "num_hidden_layers" in self.__dict__:
+            return self.__dict__["num_hidden_layers"]
+        dec = getattr(self, "decoder", None)
+        return getattr(dec, "num_hidden_layers", None)
+
 
 __all__ = ["T5GemmaConfig", "T5GemmaModuleConfig"]
diff --git a/src/transformers/models/t5gemma/modular_t5gemma.py b/src/transformers/models/t5gemma/modular_t5gemma.py
index d358a51d0e68..28a08fe6a382 100644
--- a/src/transformers/models/t5gemma/modular_t5gemma.py
+++ b/src/transformers/models/t5gemma/modular_t5gemma.py
@@ -210,6 +210,15 @@ def get_text_config(self, *args, **kwargs):
         # Always return self, regardless of the decoder option.
         return self
 
+    # Bridge for generation/cache utils which expect `config.num_hidden_layers`.
+    # Prefer a top-level override if present; otherwise use the decoder's count.
+    @property
+    def num_hidden_layers(self):
+        if "num_hidden_layers" in self.__dict__:
+            return self.__dict__["num_hidden_layers"]
+        dec = getattr(self, "decoder", None)
+        return getattr(dec, "num_hidden_layers", None)
+
 
 class T5GemmaRMSNorm(Gemma2RMSNorm):
     pass
diff --git a/tests/models/t5gemma/test_generation_t5gemma.py b/tests/models/t5gemma/test_generation_t5gemma.py
new file mode 100644
index 000000000000..729177de2a0d
--- /dev/null
+++ b/tests/models/t5gemma/test_generation_t5gemma.py
@@ -0,0 +1,32 @@
+import torch
+
+from transformers import T5GemmaConfig, T5GemmaForConditionalGeneration, T5GemmaModuleConfig
+
+
+def _tiny():
+    return T5GemmaModuleConfig(
+        vocab_size=33,
+        hidden_size=32,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=32,
+        max_position_embeddings=1024,
+        tie_word_embeddings=False,
+        layer_types=["full_attention"] * 2,
+        rope_theta=10000,
+        bos_token_id=0,
+        eos_token_id=1,
+        pad_token_id=2,
+    )
+
+
+def test_generate_use_cache_works_for_t5gemma():
+    cfg = T5GemmaConfig(encoder=_tiny(), decoder=_tiny(), vocab_size=33, attn_implementation="eager")
+    model = T5GemmaForConditionalGeneration(cfg)
+
+    output = model.generate(torch.randint(0, 33, (1, 10)), use_cache=True, max_new_tokens=2)
+
+    assert output.shape[0] == 1
+    assert output.shape[1] > 0