Add validation for attention type in the Attention module.

Use enum.auto() for AttentionType members. Add a ValueError check to ensure attn_type is either GLOBAL or LOCAL_SLIDING when processing attention masks.

PiperOrigin-RevId: 871311809

Author: MaxGuigon

gemma/gm/nn/_modules.py

@@ -53,8 +53,8 @@ def create_sliding_mask(
 
 
 class AttentionType(enum.Enum):
-  GLOBAL = 1
-  LOCAL_SLIDING = 2
+  GLOBAL = enum.auto()
+  LOCAL_SLIDING = enum.auto()
 
 
 class Embedder(nn.Module):
@@ -267,6 +267,11 @@ def __call__(
       )
       # [batch_size, seq_len, cache_size]
       attn_mask *= sliding_mask
+    elif self.attn_type != AttentionType.GLOBAL:
+      raise ValueError(
+          'Attn_type must be either AttentionType.GLOBAL or'
+          f' AttentionType.GLOBAL not {self.attn_type}'
+      )
 
     # [batch_size, seq_len, num_heads, cache_size]
     padded_logits = jnp.where((jnp.expand_dims(attn_mask, -2)), logits, K_MASK)