From 480cc5eebd4af3457b111cc6a2a5048c44938e27 Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Sat, 25 Jul 2026 11:40:01 +0530 Subject: [PATCH] Fix QueryPreAttentionNormalisation docstring to describe query scaling The enum's docstring read 'Initialization strategy.', but every member (BY_ONE_OVER_SQRT_HEAD_DIM, BY_EMBED_DIM_DIV_NUM_HEADS, ...) selects how the query is scaled before attention -- unrelated to initialization. The line appears copy-pasted; replace with an accurate description. --- gemma/gm/nn/_config.py | 2 +- gemma/gm/nn/gemma3n/_config.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/gemma/gm/nn/_config.py b/gemma/gm/nn/_config.py index 591c865b..89210042 100644 --- a/gemma/gm/nn/_config.py +++ b/gemma/gm/nn/_config.py @@ -43,7 +43,7 @@ def make_attention_layers_types( class QueryPreAttentionNormalisation(enum.Enum): - """Initialization strategy.""" + """Query pre-attention scaling strategy.""" # Whether to scale the query by 1/sqrt(head_dim) BY_ONE_OVER_SQRT_HEAD_DIM = enum.auto() diff --git a/gemma/gm/nn/gemma3n/_config.py b/gemma/gm/nn/gemma3n/_config.py index d9449dce..e56ecfa1 100644 --- a/gemma/gm/nn/gemma3n/_config.py +++ b/gemma/gm/nn/gemma3n/_config.py @@ -84,7 +84,7 @@ def create_kv_cache_sharing_patterns( # pylint: disable=invalid-name class QueryPreAttentionNormalisation(enum.Enum): - """Initialization strategy.""" + """Query pre-attention scaling strategy.""" # Apply no scaling. NONE = enum.auto()