Code update

PiperOrigin-RevId: 948315473

Author: frt03

gemma/gm/text/_sampler_loop.py

@@ -90,7 +90,10 @@ def attention_mask_for_step(self) -> Bool['B cache_length']:  # pyrefly: ignore[
     cache_length = self.full_attention_mask.shape[-1]
 
     # +1 because the current step can be self-attended too.
-    step_mask = jnp.arange(cache_length) < self.used_cache_length + 1
+    used_cache_len = self.used_cache_length
+    if used_cache_len.ndim > 0:
+      used_cache_len = used_cache_len[..., None]
+    step_mask = jnp.arange(cache_length) < used_cache_len + 1
     attention_mask = self.full_attention_mask * step_mask
     return attention_mask
 

gemma/gm/utils/_cache_helper.py

@@ -71,10 +71,10 @@ def at(self) -> _CacheProxyAt:
     return _CacheProxyAt(self)
 
   @property
-  def end_index(self) -> Int['']:  # pyrefly: ignore[not-a-type]
+  def end_index(self) -> Int['*B']:  # pyrefly: ignore[not-a-type]
     """End index of the cache."""
     layer_data = next(iter(self.cache.values()))
-    return layer_data['end_index'][0]
+    return layer_data['end_index']
 
   def set_end_index(self, value: int) -> Cache:
     """Set the end index of the cache."""
@@ -86,7 +86,7 @@ def set_end_index(self, value: int) -> Cache:
   def is_full(self) -> Bool['']:  # pyrefly: ignore[not-a-type]
     """Returns whether the cache is full."""
     # Maybe will lose the last token.
-    return self.end_index >= self.total_cache_length - 1
+    return jnp.any(self.end_index >= self.total_cache_length - 1)
 
 
 @dataclasses.dataclass(frozen=True)