From 22afad1f64831709fbeea5295d7d4d5ead8bebf3 Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Sat, 25 Jul 2026 11:41:12 +0530 Subject: [PATCH] Fix _maybe_unbatch return annotation Float -> Array (dtype-preserving) _maybe_unbatch is a dtype-preserving passthrough (unbatch) over an Array['B *d'] input, and its only caller (the `tokens` property) declares Int and passes the integer predicted_tokens buffer. Annotating the return as Float['*d'] is a copy-paste error; mirror the generic input as Array['*d']. --- gemma/research/t5gemma/sampling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gemma/research/t5gemma/sampling.py b/gemma/research/t5gemma/sampling.py index 3ea754b2..698b325c 100644 --- a/gemma/research/t5gemma/sampling.py +++ b/gemma/research/t5gemma/sampling.py @@ -94,7 +94,7 @@ def tokens(self) -> Int['B L'] | Int['L']: """Predicted tokens.""" return self._maybe_unbatch(self.state.predicted_tokens) - def _maybe_unbatch(self, x: Array['B *d']) -> Float['*d']: + def _maybe_unbatch(self, x: Array['B *d']) -> Array['*d']: if isinstance(self.text, str): (x,) = x return x