Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.3#4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16import copy17import math18from collections.abc import Callable, Sequence19from typing import Any, Optional, Union20 21import torch22import torch.nn as nn23import torch.nn.functional as F24 25from ...activations import ACT2FN26from ...cache_utils import Cache, DynamicCache27from ...configuration_utils import PretrainedConfig, layer_type_validation28from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask29from ...modeling_flash_attention_utils import FlashAttentionKwargs30from ...modeling_outputs import BaseModelOutputWithPast31from ...modeling_rope_utils import rope_config_validation32from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel33from ...processing_utils import Unpack34from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging35from ...utils.deprecation import deprecate_kwarg36from ..auto import AutoModel37from ..gemma2.configuration_gemma2 import Gemma2Config38from ..gemma2.modeling_gemma2 import (39 Gemma2MLP,40 Gemma2PreTrainedModel,41 Gemma2RotaryEmbedding,42 eager_attention_forward,43 rotate_half,44)45from ..gemma3.modeling_gemma3 import (46 Gemma3Attention,47 Gemma3DecoderLayer,48 Gemma3ForCausalLM,49 Gemma3RMSNorm,50 Gemma3TextModel,51 Gemma3TextScaledWordEmbedding,52)53from ..paligemma.modeling_paligemma import (54 PaliGemmaCausalLMOutputWithPast,55 PaliGemmaForConditionalGeneration,56 PaliGemmaModel,57 PaligemmaModelOutputWithPast,58)59from ..timm_wrapper.configuration_timm_wrapper import TimmWrapperConfig60 61 62logger = logging.get_logger(__name__)63 64 65class Gemma3nTextConfig(Gemma2Config, PretrainedConfig):66 r"""67 This is the configuration class to store the configuration of a [`Gemma3nTextModel`]. It is used to instantiate an68 Gemma3nTextModel model according to the specified arguments, defining the model architecture. Instantiating a69 configuration with the defaults will yield a similar configuration to that of the Gemma 3n E4B, e.g.70 [google/gemma-3n-E4B](https://huggingface.co/google/gemma-3n-E4B).71 72 Configuration objects that inherit from [`Gemma3nTextConfig`] and can be used to control the model outputs. Read73 the documentation from [`Gemma3nTextConfig`] for more information.74 75 Args:76 vocab_size (`int`, *optional*, defaults to 262400):77 Vocabulary size of the Gemma3nText model. Defines the number of different tokens that can be represented by78 the `inputs_ids` passed when calling [`Gemma3nTextModel`]79 vocab_size_per_layer_input (`int`, *optional*, defaults to 262144):80 Vocabulary size of the per-layer text embeddings that augment the standard embeddings.81 hidden_size (`int`, *optional*, defaults to 2048):82 Dimension of the hidden representations.83 hidden_size_per_layer_input (`int`, *optional*, defaults to 256):84 Dimension of the hidden representations for per-layer emebeddings.85 intermediate_size (`int` or `Sequence[int]`, *optional*, defaults to 16384):86 Dimension of the MLP representations. MatFormer configurations may wish to provide a sequence of integers87 to account for variable intermediate_size values across layers. In such cases,88 `len(intermediate_size) == num_hidden_layers`.89 num_hidden_layers (`int`, *optional*, defaults to 35):90 Number of hidden layers in the Transformer decoder.91 num_attention_heads (`int`, *optional*, defaults to 8):92 Number of attention heads for each attention layer in the Transformer decoder.93 num_key_value_heads (`int`, *optional*, defaults to 2):94 This is the number of key_value heads that should be used to implement Grouped Query Attention. If95 `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if96 `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When97 converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed98 by meanpooling all the original heads within that group. For more details checkout this99 [paper](https://huggingface.co/papers/2305.13245). If not specified, will default to `num_attention_heads`.100 head_dim (`int`, *optional*, defaults to 256):101 The attention head dimension.102 hidden_activation (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):103 The non-linear activation function (function or string) in the decoder. Will default to104 `"gelu_pytorch_tanh"` if not specified. `"gelu_pytorch_tanh"` uses an approximation of the `"gelu"`105 activation function.106 max_position_embeddings (`int`, *optional*, defaults to 32768):107 The maximum sequence length that this model might ever be used with.108 initializer_range (`float`, *optional*, defaults to 0.02):109 The standard deviation of the truncated_normal_initializer for initializing all weight matrices.110 rms_norm_eps (`float`, *optional*, defaults to 1e-06):111 The epsilon used by the rms normalization layers.112 use_cache (`bool`, *optional*, defaults to `True`):113 Whether or not the model should return the last key/values attentions (not used by all models). Only114 relevant if `config.is_decoder=True`.115 pad_token_id (`int`, *optional*, defaults to 0):116 Padding token id.117 eos_token_id (`int`, *optional*, defaults to 1):118 End of stream token id.119 bos_token_id (`int`, *optional*, defaults to 2):120 Beginning of stream token id.121 rope_theta (`float`, *optional*, defaults to 1000000.0):122 The base period of the RoPE embeddings.123 rope_scaling (`Dict`, *optional*):124 Dictionary containing the scaling configuration for the RoPE embeddings used in global attention.125 NOTE: if you apply new rope type and you expect the model to work on longer `max_position_embeddings`, we126 recommend you to update this value accordingly.127 Expected contents:128 `rope_type` (`str`):129 The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',130 'llama3'], with 'default' being the original RoPE implementation.131 `factor` (`float`, *optional*):132 Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In133 most scaling types, a `factor` of x will enable the model to handle sequences of length x *134 original maximum pre-trained length.135 `original_max_position_embeddings` (`int`, *optional*):136 Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during137 pretraining.138 `attention_factor` (`float`, *optional*):139 Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention140 computation. If unspecified, it defaults to value recommended by the implementation, using the141 `factor` field to infer the suggested value.142 `beta_fast` (`float`, *optional*):143 Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear144 ramp function. If unspecified, it defaults to 32.145 `beta_slow` (`float`, *optional*):146 Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear147 ramp function. If unspecified, it defaults to 1.148 `short_factor` (`List[float]`, *optional*):149 Only used with 'longrope'. The scaling factor to be applied to short contexts (<150 `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden151 size divided by the number of attention heads divided by 2152 `long_factor` (`List[float]`, *optional*):153 Only used with 'longrope'. The scaling factor to be applied to long contexts (<154 `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden155 size divided by the number of attention heads divided by 2156 `low_freq_factor` (`float`, *optional*):157 Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE158 `high_freq_factor` (`float`, *optional*):159 Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE160 rope_local_base_freq (float, *optional*, defaults to 10000.0):161 The base period of the RoPE embeddings for local attention.162 attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):163 Whether to use a bias in the query, key, value and output projection layers during self-attention.164 attention_dropout (`float`, *optional*, defaults to 0.0):165 The dropout ratio for the attention probabilities.166 sliding_window (`int`, *optional*, defaults to 512):167 This is the size of the sliding window used by local attention layers.168 layer_types (`Optional`, *optional*):169 A sequence of strings defining the attention type for that layer as either "sliding_attention" or170 "full_attention". If not provided, `layer_types` will de inferred from `num_hidden_layers` using a pattern171 of four "sliding_attention" layers followed one "full_attention". The last layer in the model should always172 be a "full_attention" layer.173 final_logit_softcapping (`float`, *optional*, defaults to 30.0):174 Scaling factor when applying tanh softcapping on the logits.175 altup_active_idx (`int`, *optional*, defaults to 0):176 The index of the prediction from which AltUp will compute additional predictions or correct177 altup_coef_clip (`float`, *optional*, defaults to 120.0):178 The maximum amplitude of an AltUp prediction or correction coefficient weight.179 altup_correct_scale (`bool`, *optional*, defaults to `True`):180 If True, apply the `AltUp.correct_output_scale` to the corrected prediction at `altup_active_idx`.181 altup_num_inputs (`int`, *optional*, defaults to 4):182 The number of predictions that AltUp should be make given the input sequence.183 num_kv_shared_layers (`int`, *optional*, defaults to 15):184 The number of layer that share KV cache values. During the forward pass, the last `num_kv_shared_layers`185 layers in the model "share" the KV values in that each local and global layer in this range uses the KV186 cache values computed for the last local or global layer, respectively, before entering this range. The187 value should be a multiple of the attention pattern size (see `layer_types` parameter).188 laurel_rank (int, *optional*, defaults to 64):189 The intermediate size for the linear projections in the Learned Augmented Residual Layer.190 activation_sparsity_pattern (Sequence[float], *optional*):191 The sparsity factor used to extract the top-k activations for a given layer. The provided Sequence must192 explicitly provide a sparsity value for each layer in the model. By default, the first 10 layers are193 sparse with a sparsity factor of 0.95 and the rest are dense.194 195 ```python196 >>> from transformers import Gemma3nTextModel, Gemma3nTextConfig197 198 >>> # Initializing a Gemma3nText gemma3n_text-E4B style configuration199 >>> configuration = Gemma3nTextConfig()200 201 >>> # Initializing a model from the gemma3n_text-E4B style configuration202 >>> model = Gemma3nTextModel(configuration)203 204 >>> # Accessing the model configuration205 >>> configuration = model.config206 ```207 """208 209 model_type = "gemma3n_text"210 211 def __init__(212 self,213 vocab_size: int = 262_400,214 vocab_size_per_layer_input: int = 262_144,215 hidden_size: int = 2048,216 hidden_size_per_layer_input: int = 256,217 intermediate_size: Union[int, Sequence[int]] = 16_384,218 num_hidden_layers: int = 35,219 num_attention_heads: int = 8,220 num_key_value_heads: int = 2,221 head_dim: int = 256,222 hidden_activation: str = "gelu_pytorch_tanh",223 max_position_embeddings: int = 32_768,224 initializer_range: float = 0.02,225 rms_norm_eps: float = 1e-6,226 use_cache: bool = True,227 pad_token_id: int = 0,228 eos_token_id: int = 1,229 bos_token_id: int = 2,230 rope_theta: float = 1_000_000.0,231 rope_scaling: Optional[dict[str, Any]] = None,232 rope_local_base_freq: float = 10_000.0,233 attention_bias: bool = False,234 attention_dropout: float = 0.0,235 sliding_window: int = 512,236 layer_types: Optional[Sequence[str]] = None,237 final_logit_softcapping: float = 30.0,238 altup_active_idx: int = 0,239 altup_coef_clip: float = 120.0,240 altup_correct_scale: bool = True,241 altup_num_inputs: int = 4,242 num_kv_shared_layers: int = 15,243 laurel_rank: int = 64,244 activation_sparsity_pattern: Optional[Union[float, Sequence[float]]] = None,245 **kwargs,246 ):247 PretrainedConfig.__init__(248 pad_token_id=pad_token_id,249 bos_token_id=bos_token_id,250 eos_token_id=eos_token_id,251 **kwargs,252 )253 254 if isinstance(intermediate_size, Sequence) and (intsize_len := len(intermediate_size)) != num_hidden_layers:255 raise ValueError(256 "intermediate_size must have an explicit intermediate size for every layer or one for all layers. "257 f"Expected {num_hidden_layers} values but got {intsize_len}."258 )259 elif not isinstance(intermediate_size, Sequence):260 intermediate_size = [intermediate_size] * num_hidden_layers261 262 self.vocab_size = vocab_size263 self.vocab_size_per_layer_input = vocab_size_per_layer_input264 self.max_position_embeddings = max_position_embeddings265 self.hidden_size = hidden_size266 self.intermediate_size = intermediate_size267 self.num_hidden_layers = num_hidden_layers268 self.num_attention_heads = num_attention_heads269 self.head_dim = head_dim270 self.num_key_value_heads = num_key_value_heads271 self.initializer_range = initializer_range272 self.rms_norm_eps = rms_norm_eps273 self.use_cache = use_cache274 self.rope_theta = rope_theta275 self.attention_bias = attention_bias276 self.attention_dropout = attention_dropout277 self.hidden_activation = hidden_activation278 self.sliding_window = sliding_window279 self.final_logit_softcapping = final_logit_softcapping280 self.layer_types = layer_types281 282 self.rope_local_base_freq = rope_local_base_freq283 self.rope_scaling = rope_scaling284 rope_config_validation(self)285 286 if layer_types is None:287 self.layer_types = [288 "full_attention" if (i + 1) % 5 == 0 else "sliding_attention" for i in range(self.num_hidden_layers)289 ]290 else:291 self.layer_types = layer_types292 293 layer_type_validation(self.layer_types, self.num_hidden_layers)294 295 self.hidden_size_per_layer_input = hidden_size_per_layer_input296 self.num_kv_shared_layers = num_kv_shared_layers297 298 self.altup_active_idx = altup_active_idx299 self.altup_coef_clip = altup_coef_clip300 self.altup_correct_scale = altup_correct_scale301 self.altup_num_inputs = altup_num_inputs302 303 self.laurel_rank = laurel_rank304 305 if activation_sparsity_pattern is None:306 num_sparse_layers = 10 if num_hidden_layers > 10 else 0307 activation_sparsity_pattern = [0.95] * num_sparse_layers + [0.0] * (num_hidden_layers - num_sparse_layers)308 309 if (len_asp := len(activation_sparsity_pattern)) != num_hidden_layers:310 raise ValueError(311 "activation_sparsity_pattern must have an explicit activation sparsity value for every layer."312 f"Expected {num_hidden_layers} values but got {len_asp}."313 )314 self.activation_sparsity_pattern = activation_sparsity_pattern315 316 317class Gemma3nAudioConfig(PretrainedConfig):318 r"""319 This is the configuration class to store the configuration of a [`Gemma3nAudioEncoder`]. It is used to instantiate320 an `Gemma3nAudioEncoder` model according to the specified arguments, defining the model architecture. Instantiating321 a configuration with the defaults will yield a similar configuration to that of the Gemma 3n E4B, e.g.,322 [google/gemma-3n-E4B](https://huggingface.co/google/gemma-3n-E4B).323 324 Configuration objects that inherit from [`Gemma3nAudioConfig`] and can be used to control the model outputs. Read325 the documentation from [`Gemma3nAudioConfig`] for more information.326 327 Args:328 vocab_size (`int`, *optional*, defaults to 128):329 Vocabulary size of the additional hard-token embeddings for audio model. These augment the embeddings330 included in the `Gemma3nTextModel` to provide, e.g., the end of audio and audio soft token placeholder331 tokens when converting `input_ids` to embeddings in the `Gemma3nForConditionalGeneration` model.332 vocab_offset (`int`, *optional*, defaults to 262272):333 Offset between the tokenizer vocab index for the token ids embedded by `Gemma3nMultimodalEmbedder` and the334 0-indexed `Gemma3nMultimodalEmbedder.embedding` table.335 input_feat_size (`int`, *optional*, defaults to 128):336 The number of channels in each mel-spectrogram frame.337 hidden_size (`int`, *optional*, defaults to 1536):338 Dimension of the hidden representations.339 rms_norm_eps (`float`, *optional*, defaults to 1e-06):340 The epsilon used by the rms normalization layers.341 gradient_clipping (`float`, *optional*, defaults to 10000000000.0):342 Clipping value used to stabilize extremely large gradient values.343 conf_attention_chunk_size (`int`, *optional*, defaults to 12):344 The sub-sequence size for local attention processing inside the Conformer ("conf") section of the345 Universal Speech Model.346 conf_attention_context_left (`int`, *optional*, defaults to 13):347 The left context size of the local attention inside the Conformer ("conf") section of the348 Universal Speech Model.349 conf_attention_context_right (`int`, *optional*, defaults to 0):350 The right context size of the local attention inside the Conformer ("conf") section of the351 Universal Speech Model.352 conf_attention_logit_cap (`float`, *optional*, defaults to 50.0):353 Logit cap applied during local attention inside the Conformer ("conf") section of the354 Universal Speech Model.355 conf_num_attention_heads (`int`, *optional*, defaults to 8):356 The number of attention heads in local attention inside the Conformer ("conf") section of the357 Universal Speech Model.358 conf_num_hidden_layers (`int`, *optional*, defaults to 12):359 The number of layers that use local attention inside the Conformer ("conf") section of the360 Universal Speech Model.361 conf_conv_kernel_size (`int`, *optional*, defaults to 5):362 Convolution kernel size for the conformer block inside the Conformer ("conf") section of the363 Universal Speech Model.364 conf_reduction_factor (`int`, *optional*, defaults to 4):365 Reduction factor used in the conformer block inside the Conformer ("conf") section of the366 Universal Speech Model.367 conf_residual_weight (`float`, *optional*, defaults to 0.5):368 Residual connection weight inside the Conformer ("conf") section of the369 Universal Speech Model.370 sscp_conv_channel_size (`tuple(int, int)`, *optional*, defaults to `(128, 32)`):371 The channel sizes for the first and second convolutional layers in the Sub-sample Convolution Projection372 ("sscp") section of the Universal Speech Model.373 sscp_conv_group_norm_eps (`float`, *optional*, defaults to 0.001):374 Epsilon used in group normalization in the subsample convolution projection in the Sub-sample Convolution375 Projection ("sscp") section of the Universal Speech Model.376 sscp_conv_kernel_size (`tuple(tuple(int, int), tuple(int, int))`, *optional*, defaults to `((3, 3), (3, 3))`):377 Kernel sizes of the two convolutional layers in the subsample convolution projection in the Sub-sample378 Convolution Projection ("sscp") section of the Universal Speech Model. The kernel sizes are specified as a379 tuple of height and width for each layer, where the height corresponds to the time dimension and the width380 corresponds to the frequency dimension.381 sscp_conv_stride_size (`tuple(tuple(int, int), tuple(int, int))`, *optional*, defaults to `((2, 2), (2, 2))`):382 Stride sizes of the two convolutional layers in the subsample convolution projection in the Sub-sample383 Convolution Projection ("sscp") section of the Universal Speech Model. The stride sizes are specified as a384 tuple of height and width for each layer, where the height corresponds to the time dimension and the width385 corresponds to the frequency dimension.386 387 Example:388 389 ```python390 >>> from transformers import Gemma3nAudioConfig, Gemma3nAudioEncoder391 392 >>> # Initializing a Gemma3nAudioEncoder gemma3n_audio-E4B-style configuration393 >>> configuration = Gemma3nAudioConfig()394 395 >>> # Initializing a model from the gemma3n_audio-E4B style configuration396 >>> model = Gemma3nAudioEncoder(configuration)397 398 >>> # Accessing the model configuration399 >>> configuration = model.config400 ```401 """402 403 model_type = "gemma3n_audio"404 405 def __init__(406 self,407 vocab_size: int = 128,408 vocab_offset: int = 262_144 + 128, # text vocab size + vision vocab size409 input_feat_size: int = 128,410 hidden_size: int = 1536,411 rms_norm_eps: float = 1e-6,412 gradient_clipping: float = 10_000_000_000.0,413 conf_attention_chunk_size: int = 12,414 conf_attention_context_left: int = 13,415 conf_attention_context_right: int = 0,416 conf_attention_logit_cap: float = 50.0,417 conf_num_attention_heads: int = 8,418 conf_num_hidden_layers: int = 12,419 conf_conv_kernel_size: int = 5,420 conf_reduction_factor: int = 4,421 conf_residual_weight: float = 0.5,422 sscp_conv_channel_size: tuple[int, int] = (128, 32),423 sscp_conv_group_norm_eps: float = 1e-3,424 sscp_conv_kernel_size: tuple[tuple[int, int], tuple[int, int]] = (425 (3, 3),426 (3, 3),427 ),428 sscp_conv_stride_size: tuple[tuple[int, int], tuple[int, int]] = (429 (2, 2),430 (2, 2),431 ),432 **kwargs,433 ):434 super().__init__(**kwargs)435 self.input_feat_size = input_feat_size436 self.hidden_size = hidden_size437 self.rms_norm_eps = rms_norm_eps438 self.vocab_size = vocab_size439 self.vocab_offset = vocab_offset440 self.gradient_clipping = gradient_clipping441 self.conf_attention_chunk_size = conf_attention_chunk_size442 self.conf_attention_context_left = conf_attention_context_left443 self.conf_attention_context_right = conf_attention_context_right444 self.conf_attention_logit_cap = conf_attention_logit_cap445 self.conf_num_attention_heads = conf_num_attention_heads446 self.conf_num_hidden_layers = conf_num_hidden_layers447 self.conf_conv_kernel_size = conf_conv_kernel_size448 self.conf_reduction_factor = conf_reduction_factor449 self.conf_residual_weight = conf_residual_weight450 self.sscp_conv_channel_size = sscp_conv_channel_size451 self.sscp_conv_group_norm_eps = sscp_conv_group_norm_eps452 self.sscp_conv_kernel_size = sscp_conv_kernel_size453 self.sscp_conv_stride_size = sscp_conv_stride_size454 455 456class Gemma3nVisionConfig(TimmWrapperConfig):457 r"""458 This is the configuration class to store the configuration for a timm backbone [`TimmWrapper`]. It is used to459 instantiate an timm model model according to the specified arguments, defining the model architecture.460 Instantiating a configuration with the defaults will yield a similar configuration to that of the Gemma 3n E4B461 vision tower, e.g. [google/gemma-3n-E4B](https://huggingface.co/google/gemma-3n-E4B).462 463 Configuration objects inherit from [`Gemma3nVisionConfig`] and can be used to control the model outputs. Read the464 documentation from [`Gemma3nVisionConfig`] for more information.465 466 Config loads imagenet label descriptions and stores them in `id2label` attribute, `label2id` attribute for default467 imagenet models is set to `None` due to occlusions in the label descriptions.468 469 Args:470 initializer_range (`float`, *optional*, defaults to 0.02):471 The standard deviation of the truncated_normal_initializer for initializing all weight matrices.472 do_pooling (`bool`, *optional*, defaults to `False`):473 Whether to do pooling for the last_hidden_state in `TimmWrapper` or not.474 architecture (`str`, *optional*, defaults to `"mobilenetv5_300m_enc"`):475 Determines vision architecture for TimmWrapper.476 hidden_size (`int`, *optional*, defaults to 2048):477 Dimension of the hidden representations.478 vocab_size (`int`, *optional*, defaults to 128):479 Vocabulary size of the additional hard-token embeddings for vision model.480 vocab_offset (`int`, *optional*, defaults to 262144):481 Offset between the tokenizer vocab index for the token ids embedded by `Gemma3nMultimodalEmbedder` and the482 0-indexed `Gemma3nMultimodalEmbedder.embedding` table.483 rms_norm_eps (`float`, *optional*, defaults to 1e-06):484 The epsilon used by the rms normalization layers.485 486 Example:487 ```python488 >>> from transformers import Gemma3nVisionConfig, TimmWrapper489 490 >>> # Initializing a TimmWrapper gemma3n_vision-E4B-style configuration491 >>> configuration = Gemma3nVisionConfig()492 493 >>> # Initializing a gemma3n_vision-E4B-style TimmWrapper from the configuration494 >>> model = TimmWrapper(configuration)495 496 >>> # Accessing the model configuration497 >>> configuration = model.config498 ```499 """500 501 model_type = "gemma3n_vision"502 503 def __init__(504 self,505 initializer_range: float = 0.02,506 do_pooling: bool = False,507 architecture: str = "mobilenetv5_300m_enc",508 hidden_size: int = 2048,509 vocab_size: int = 128,510 vocab_offset: int = 262_144,511 rms_norm_eps: float = 1e-06,512 model_args: Optional[dict] = None,513 **kwargs,514 ):515 super().__init__(**kwargs)516 self.architecture = architecture517 self.initializer_range = initializer_range518 self.do_pooling = do_pooling519 self.hidden_size = hidden_size520 self.vocab_size = vocab_size521 self.vocab_offset = vocab_offset522 self.rms_norm_eps = rms_norm_eps523 524 525class Gemma3nConfig(PretrainedConfig):526 r"""527 This is the configuration class to store the configuration of a [`Gemma3nForConditionalGeneration`]. It is used to528 instantiate a Gemma3nForConditionalGeneration according to the specified arguments, defining the model529 architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of530 Gemma3n-E4B.531 532 e.g. [google/gemma-3n-E4B](https://huggingface.co/google/gemma-3n-E4B)533 534 Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the535 documentation from [`PretrainedConfig`] for more information.536 537 Args:538 text_config (`Union[Gemma3nTextConfig, dict]`, *optional*):539 The config object of the text backbone.540 vision_config (`Union[AutoConfig, dict]`, *optional*):541 Custom vision config or dict.542 audio_config (`Union[AutoConfig, dict]`, *optional*):543 Custom audio config or dict.544 audio_soft_tokens_per_image (`int`, *optional*, defaults to 188):545 The number of soft tokens per audio clip.546 vision_soft_tokens_per_image (`int`, *optional*, defaults to 256):547 The number of soft tokens per image.548 boi_token_id (`int`, *optional*, defaults to 255999):549 The begin-of-image token index to wrap the image prompt.550 eoi_token_id (`int`, *optional*, defaults to 262144):551 The end-of-image token index to wrap the image prompt.552 image_token_id (`int`, *optional*, defaults to 262145):553 The image token index to encode the image prompt.554 boa_token_id (`int`, *optional*, defaults to 256000):555 The begin-of-audio token index to wrap the audio prompt.556 eoa_token_id (`int`, *optional*, defaults to 262272):557 The end-of-audio token index to wrap the audio prompt.558 audio_token_id (`int`, *optional*, defaults to 262273):559 The audio token index to encode the audio prompt.560 initializer_range (`float`, *optional*, defaults to 0.02):561 The standard deviation of the truncated_normal_initializer for initializing all weight matrices.562 563 564 Example:565 566 ```python567 >>> from transformers import Gemma3nForConditionalGeneration, Gemma3nConfig, Gemma3nTextConfig568 569 >>> # Initializing a MobileNet vision config, which is loaded from TIMM570 >>> vision_config = Gemma3nVisionConfig()571 572 >>> # Initializing a Gemma3n Audio config573 >>> audio_config = Gemma3nAudioConfig()574 575 >>> # Initializing a Gemma3n Text config576 >>> text_config = Gemma3nTextConfig()577 578 >>> # Initializing a Gemma3n gemma-3-4b style configuration579 >>> configuration = Gemma3nConfig(text_config, vision_config, audio_config)580 581 >>> # Initializing a model from the gemma-3-4b style configuration582 >>> model = Gemma3nTextConfig(configuration)583 584 >>> # Accessing the model configuration585 >>> configuration = model.config586 ```"""587 588 model_type = "gemma3n"589 sub_configs = {590 "text_config": Gemma3nTextConfig,591 "vision_config": Gemma3nVisionConfig,592 "audio_config": Gemma3nAudioConfig,593 }594 595 def __init__(596 self,597 text_config: Optional[Union[Gemma3nTextConfig, dict[str, Any]]] = None,598 vision_config: Optional[Union[Gemma3nVisionConfig, dict[str, Any]]] = None,599 audio_config: Optional[Union[Gemma3nAudioConfig, dict[str, Any]]] = None,600 audio_soft_tokens_per_image: int = 188,601 vision_soft_tokens_per_image: int = 256,602 boi_token_id: int = 255_999,603 eoi_token_id: int = 262_144,604 image_token_id: int = 262_145,605 boa_token_id: int = 256_000,606 eoa_token_id: int = 262_272,607 audio_token_id: int = 262_273,608 initializer_range: float = 0.02,609 **kwargs,610 ):611 super().__init__(**kwargs)612 613 if isinstance(text_config, dict):614 text_config = Gemma3nTextConfig(**text_config)615 elif text_config is None:616 text_config = Gemma3nTextConfig()617 logger.info("text_config is None. Using default Gemma3nTextConfig.")618 619 if isinstance(vision_config, dict):620 vision_config = Gemma3nVisionConfig(**vision_config)621 elif vision_config is None:622 vision_config = Gemma3nVisionConfig()623 logger.info("vision_config is None. Using default Gemma3nVisionConfig.")624 625 if isinstance(audio_config, dict):626 audio_config = Gemma3nAudioConfig(**audio_config)627 elif audio_config is None:628 audio_config = Gemma3nAudioConfig()629 logger.info("audio_config is None. Using default Gemma3nAudioConfig.")630 631 self.text_config = text_config632 self.vision_config = vision_config633 self.audio_config = audio_config634 635 self.audio_soft_tokens_per_image = audio_soft_tokens_per_image636 self.vision_soft_tokens_per_image = vision_soft_tokens_per_image637 self.boi_token_id = boi_token_id638 self.eoi_token_id = eoi_token_id639 self.image_token_id = image_token_id640 self.boa_token_id = boa_token_id641 self.eoa_token_id = eoa_token_id642 self.audio_token_id = audio_token_id643 self.initializer_range = initializer_range644 645 646class Gemma3nModelOutputWithPast(PaligemmaModelOutputWithPast):647 r"""648 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):649 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).650 651 Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see652 `past_key_values` input) to speed up sequential decoding.653 image_hidden_states (`torch.FloatTensor`, *optional*):654 A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.655 image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.656 audio_hidden_states (`torch.FloatTensor`, *optional*):657 A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.658 audio_hidden_states of the model produced by the audio encoder and after projecting the last hidden state.659 """660 661 audio_hidden_states: Optional[torch.FloatTensor] = None662 663 664class Gemma3nCausalLMOutputWithPast(PaliGemmaCausalLMOutputWithPast):665 r"""666 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):667 Language modeling loss (for next-token prediction).668 logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.text_config.vocab_size)`):669 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).670 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):671 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).672 673 Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see674 `past_key_values` input) to speed up sequential decoding.675 image_hidden_states (`torch.FloatTensor`, *optional*):676 A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.677 image_hidden_states of the model produced by the vision encoder after projecting last hidden state.678 audio_hidden_states (`torch.FloatTensor`, *optional*):679 A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.680 audio_hidden_states of the model produced by the audio encoder and after projecting the last hidden state.681 """682 683 audio_hidden_states: Optional[torch.FloatTensor] = None684 685 686class Gemma3nRMSNorm(Gemma3RMSNorm):687 def __init__(self, dim: int, eps: float = 1e-6, with_scale: bool = True):688 super().__init__(dim, eps=eps)689 del self.weight690 self.with_scale = with_scale691 692 if self.with_scale:693 self.weight = nn.Parameter(torch.ones(dim))694 else:695 self.register_buffer("weight", torch.tensor(1.0), persistent=False)696 697 def _norm(self, x):698 return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)699 700 def forward(self, x: torch.Tensor) -> torch.Tensor:701 # Llama does x.to(float16) * w whilst Gemma2 is (x * w).to(float16)702 # See https://github.com/huggingface/transformers/pull/29402703 output = self._norm(x.float()) * self.weight.float()704 return output.type_as(x)705 706 707# ==== Audio Encoder ====708 709 710class Gemma3nAudioRelativePositionEmbedding(nn.Module):711 def __init__(self, config: Gemma3nAudioConfig):712 super().__init__()713 self.config = config714 715 self.num_heads = self.config.conf_num_attention_heads716 self.channels = self.config.hidden_size717 self.head_dim = self.channels // self.num_heads718 self.max_backward = max(0, self.config.conf_attention_context_left - 1)719 self.max_forward = self.config.conf_attention_context_right720 721 self.pos_proj = nn.Linear(self.channels, self.num_heads * self.head_dim, bias=False)722 723 min_timescale = 1.0724 max_timescale = 1.0e4725 num_timescales = self.channels // 2726 log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / max(num_timescales - 1, 1)727 inv_timescales = min_timescale * torch.exp(torch.arange(num_timescales) * -log_timescale_increment)728 self.register_buffer(729 "inv_timescales",730 inv_timescales.float().unsqueeze(0).unsqueeze(0),731 persistent=False,732 )733 734 def _get_timing_signal_1d_pos(self, position: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:735 position = position.float().unsqueeze(-1)736 scaled_time = position * self.inv_timescales.to(device=position.device, dtype=torch.float32)737 timing_signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=-1)738 return timing_signal.type(dtype)739 740 def _relative_shift(741 self,742 term_bd_before_shift: torch.Tensor,743 batch_size: int,744 num_heads: int,745 num_query_blocks: int,746 query_block_size: int,747 key_context_size: int,748 max_span_plus_1: int,749 ) -> torch.Tensor:750 """Performs the relative shift.751 752 Args:753 term_bd_before_shift: Tensor of shape [B, N, U, W, F_span]. batch_size754 (B), num_heads (N), num_query_blocks (U), query_block_size (W),755 key_context_size (C = W+L+R), max_span_plus_1 (F_span = L+R+1).756 757 Returns:758 Tensor of shape [B, N, U, W, C].759 """760 # term_bd_before_shift shape: [B, N, U, W, F_span]761 # Target shape after shift: [B, N, U, W, C]762 763 # Padding amount for the last dimension (F_span) to become (C + 1)764 # C = key_context_size765 # F_span = max_span_plus_1766 pad_amount_last_dim = (key_context_size + 1) - max_span_plus_1767 768 # PyTorch F.pad expects (pad_left, pad_right, pad_top, pad_bottom ...)769 # We only pad the last dimension on the right.770 padding_tuple = (0, pad_amount_last_dim)771 772 term_bd_padded = nn.functional.pad(term_bd_before_shift, padding_tuple)773 # Shape after pad: [B, N, U, W, C+1]774 775 # Reshape for slicing (emulating JAX's behavior)776 # [B, N, U, W * (C+1)]777 term_bd_reshaped = term_bd_padded.reshape(778 (779 batch_size,780 num_heads,781 num_query_blocks,782 query_block_size * (key_context_size + 1),783 )784 )785 786 # Slice to effective [B, N, U, W * C]787 term_bd_sliced = term_bd_reshaped[:, :, :, : query_block_size * key_context_size]788 789 # Reshape back to [B, N, U, W, C]790 term_bd_shifted = term_bd_sliced.reshape(791 (792 batch_size,793 num_heads,794 num_query_blocks,795 query_block_size,796 key_context_size,797 )798 )799 return term_bd_shifted800 801 def forward(self, queries: torch.Tensor, keys: torch.Tensor) -> torch.Tensor:802 # queries: [B, U, W, N, H] (batch, num_query_blocks, query_block_size, num_heads, head_dim)803 # keys: [B, U, C, N, H] (batch, num_query_blocks, key_context_size, num_heads, head_dim)804 # C = W + L + R (key_context_size)805 # F_span = L + R + 1 (max_span + 1)806 807 batch_size, num_query_blocks, query_block_size, num_heads, head_dim = queries.shape808 _, _, key_context_size, _, _ = keys.shape809 810 # Relative positions for sinusoidal embeddings: [L, L-1, ..., -R]811 # Length is L+R+1 = self.max_span + 1812 pos_indices = torch.arange(self.max_backward, -self.max_forward - 1, -1, device=queries.device).unsqueeze(813 0814 ) # Shape [1, F_span]815 816 max_span_plus_1 = pos_indices.shape[1] # F_span817 818 sin_emb_timing_signal = self._get_timing_signal_1d_pos(819 pos_indices, dtype=queries.dtype820 ) # Shape [1, F_span, self.channels]821 822 # Project sinusoidal embeddings: [1, F_span, self.channels] -> [1, F_span, N*H]823 projected_sin_emb = self.pos_proj(sin_emb_timing_signal)824 # Reshape to [1, F_span, N, H] then squeeze to [F_span, N, H]825 sin_emb = projected_sin_emb.reshape(1, max_span_plus_1, self.num_heads, self.head_dim).squeeze(826 0827 ) # Shape [F, N, H]828 829 # term_ac: Query-Key content interaction830 # queries: [B, U, W, N, H] -> permute to [B, N, U, W, H] for matmul831 # keys: [B, U, C, N, H] -> permute to [B, N, U, H, C] for matmul832 queries_p = queries.permute(0, 3, 1, 2, 4) # [B, N, U, W, H]833 keys_p_t = keys.permute(0, 3, 1, 4, 2) # [B, N, U, H, C]834 term_ac = torch.matmul(queries_p, keys_p_t) # [B, N, U, W, C]835 836 # term_bd: Query-Position interaction837 # Original einsum: term_bd_unshifed = torch.einsum('buwnh,fnh->bnuwf', queries, sin_emb)838 # queries shape: [B, U, W, N, H]839 # sin_emb shape: [F, N, H]840 # Target output shape: [B, N, U, W, F]841 842 # Permute queries to [B, N, U, W, H] for easier broadcasting with sin_emb843 q_permuted = queries.permute(0, 3, 1, 2, 4)844 845 # Permute sin_emb to [N, H, F] to prepare for matmul846 # sin_emb original is [F, N, H]847 s_permuted = sin_emb.permute(1, 2, 0) # Shape: [N, H, F]848 849 # Reshape queries for matmul: [B, N, U*W, H]850 q_reshaped = q_permuted.reshape(batch_size, num_heads, num_query_blocks * query_block_size, head_dim)851 852 # Perform matmul: [B, N, U*W, H] @ [N, H, F]853 # s_permuted ([N, H, F]) will be broadcast to [B, N, H, F]854 # Result: [B, N, U*W, F]855 term_bd_unshifed_matmul = torch.matmul(q_reshaped, s_permuted)856 857 # Reshape to target [B, N, U, W, F]858 term_bd_unshifed = term_bd_unshifed_matmul.reshape(859 batch_size,860 num_heads,861 num_query_blocks,862 query_block_size,863 max_span_plus_1,864 )865 866 # Apply relative shift to term_bd_unshifed867 term_bd_shifted = self._relative_shift(868 term_bd_unshifed,869 batch_size,870 num_heads,871 num_query_blocks,872 query_block_size,873 key_context_size,874 max_span_plus_1,875 ) # Shape [B, N, U, W, C]876 877 return term_ac + term_bd_shifted878 879 880class Gemma3nAudioAttention(nn.Module):881 def __init__(self, config: Gemma3nAudioConfig):882 super().__init__()883 self.config = config884 885 self.num_heads = self.config.conf_num_attention_heads886 self.hidden_size = self.config.hidden_size887 self.head_dim = self.hidden_size // self.num_heads888 889 self.chunk_size = self.config.conf_attention_chunk_size890 self.max_future_horizon = self.config.conf_attention_context_right891 self.max_past_horizon = max(0, self.config.conf_attention_context_left - 1)892 self.attention_logits_soft_cap = self.config.conf_attention_logit_cap893 self.context_size = self.chunk_size + self.max_past_horizon + self.max_future_horizon894 895 self.relative_position_embedding = Gemma3nAudioRelativePositionEmbedding(config)896 self.per_dim_scale = nn.Parameter(torch.zeros((self.head_dim,)))897 898 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)899 self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)900 self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)901 902 q_scale = self.head_dim**-0.5903 r_softplus_0 = 1.0 / torch.nn.functional.softplus(torch.tensor(0.0))904 self.register_buffer("q_scale", (q_scale * r_softplus_0).clone().detach(), persistent=False)905 906 lower_causal_mask = torch.tril(907 torch.ones((self.context_size, self.chunk_size), dtype=torch.bool),908 diagonal=0,909 ).T910 upper_causal_mask = torch.tril(911 torch.ones((self.chunk_size, self.context_size), dtype=torch.bool),912 diagonal=self.max_past_horizon + self.max_future_horizon,913 )914 local_causal_valid_mask = torch.ones((self.chunk_size, self.context_size), dtype=torch.bool)915 local_causal_valid_mask = local_causal_valid_mask * lower_causal_mask * upper_causal_mask916 self.register_buffer("local_causal_valid_mask", local_causal_valid_mask, persistent=False)917 918 self.register_buffer(919 "softcap",920 torch.tensor(self.attention_logits_soft_cap).float(),921 persistent=False,922 )923 924 def _pad_dim1(self, x: torch.Tensor, pad_left: int, pad_right: int) -> torch.Tensor:925 batch, _, *tail_shape = x.shape926 left = x.new_zeros((batch, pad_left, *tail_shape))927 right = x.new_zeros((batch, pad_right, *tail_shape))928 x = torch.cat([left, x, right], dim=1)929 return x930 931 def _convert_to_block(self, hidden_states: torch.Tensor) -> torch.Tensor:932 """Turns a sequence to non overlapping blocks.933 934 Args:935 hidden_states: a tensor of [batch, time, ...].936 937 Returns:938 A tensor of [batch, num_blocks, block_size, ...], with necessary939 paddings,940 where output[:, i, ...] are x[:, i*block_size:(i+1)*block_size, ...].941 """942 shape = hidden_states.shape943 b, t = shape[:2]944 num_blocks = (t + self.chunk_size - 1) // self.chunk_size945 946 if (padding_len := num_blocks * self.chunk_size - t) > 0:947 hidden_states = self._pad_dim1(hidden_states, 0, padding_len)948 949 permute_dims = (b, num_blocks, self.chunk_size) + shape[2:]950 hidden_states = hidden_states.reshape(permute_dims).contiguous()951 return hidden_states952 953 def _extract_block_context(self, hidden_states: torch.Tensor) -> torch.Tensor:954 """Extracts temporal context for every block.955 956 Args:957 hidden_states: a tensor of [batch, time, ...].958 959 Returns:960 A tensor of [batch, num_blocks, context_size, ...], with necessary961 paddings,962 where context_size = block_size + left_context + right_context,963 and output[:, i, ...] are x[:, start-left_context:end+right_context,964 ...],965 start = i * block_size, end = (i + 1) * block_size.966 """967 pad_left = self.max_past_horizon968 # The JAX equivalent padding for signal.frame with pad_mode='valid' is969 # (left_context, right_context + block_size - 1) on the time dimension.970 # PyTorch's _pad_dim1 applies padding symmetrically if only one value is given,971 # or (pad_dim_start, pad_dim_end) if two are given.972 # Our _pad_dim1(x, pad_left, pad_right) pads dim -2 (time for [B,T,N,H])973 # or dim 1 (time for [B,T]).974 # The current pad_right calculation matches the JAX effective padding.975 pad_right = self.max_future_horizon + self.chunk_size - 1976 hidden_states = self._pad_dim1(hidden_states, pad_left, pad_right)977 978 frame_len = self.context_size979 frame_step = self.chunk_size980 981 # Directly use unfold without the subframe_factor logic982 # x.unfold(dimension, size, step)983 # dimension=1 (time dimension, assuming x is [B, T_padded, ...])984 # size=frame_len (context_size)985 # step=frame_step (chunk_size)986 x_unfolded = hidden_states.unfold(dimension=1, size=frame_len, step=frame_step)987 988 # If x was [B, T_padded], x_unfolded is [B, num_blocks, frame_len]989 # If x was [B, T_padded, N, H], x_unfolded is [B, num_blocks, N, H, frame_len]990 # We want to match JAX's typical output for such operations which might be991 # [B, num_blocks, frame_len, N, H] if N, H are present.992 # The relative_position_embedding expects keys as [B, U, C, N, H].993 # If x_unfolded is [B, U, N, H, C(frame_len)], we need to move C.994 if hidden_states.ndim > 2 and x_unfolded.ndim > 3: # Check if inner dimensions (like N, H) exist995 # Current shape after unfold for [B, T_pad, N, H] is [B, U, N, H, C]996 # Target shape for keys in RPE: [B, U, C, N, H]997 x_unfolded = torch.movedim(x_unfolded, source=-1, destination=2)998 999 return x_unfolded.contiguous()1000 1001 def forward(self, hidden_states: torch.Tensor, mask: torch.BoolTensor) -> torch.Tensor:1002 # sl.Dense uses jax.numpy.einsum("...a,abcd->...bcd") and jax.numpy.select()1003 qkv_shape = (*hidden_states.shape[:-1], self.num_heads, self.head_dim)1004 query_states = self.q_proj(hidden_states).reshape(qkv_shape).contiguous()1005 key_states = self.k_proj(hidden_states).reshape(qkv_shape).contiguous()1006 value_states = self.v_proj(hidden_states).reshape(qkv_shape).contiguous()1007 1008 per_dim_scale_sp = torch.nn.functional.softplus(self.per_dim_scale)1009 1010 broadcast_shape = (1, 1, 1, self.head_dim)1011 per_dim_scale_sp_broadcast = per_dim_scale_sp.view(broadcast_shape)1012 query_states = query_states * self.q_scale * per_dim_scale_sp_broadcast1013 1014 batch_size, q_time = query_states.shape[:2]1015 1016 query_blocks = self._convert_to_block(query_states)1017 key_blocks = self._extract_block_context(key_states)1018 value_blocks = self._extract_block_context(value_states)1019 num_query_blocks = query_blocks.shape[1]1020 1021 # 1. Create a mask indicating originally valid positions.1022 original_valid_mask = ~mask # True for valid, False for padded1023 1024 # 2. Extract blocks from this validity mask.1025 extracted_valid_mask_blocks = self._extract_block_context(original_valid_mask)1026 1027 # If subframe_factor was used in _extract_block_context for a [B, T] input mask,1028 # the shape might be [B, U, C/SF, SF]. Reshape to [B, U, C].1029 # batch_size and num_query_blocks are known from query_blocks.1030 # self.context_size is C.1031 if (1032 extracted_valid_mask_blocks.ndim == 41033 and extracted_valid_mask_blocks.shape[2] * extracted_valid_mask_blocks.shape[3] == self.context_size1034 ):1035 extracted_valid_mask_blocks = extracted_valid_mask_blocks.reshape(1036 batch_size, num_query_blocks, self.context_size1037 )1038 # After potential reshape, ensure it's [B, U, C] if it was from a [B,T] mask.1039 # This assertion might be too strict if _extract_block_context handles higher-rank inputs differently,1040 # but for the mask case, this should hold.1041 if extracted_valid_mask_blocks.shape != (1042 batch_size,1043 num_query_blocks,1044 self.context_size,1045 ):1046 raise ValueError(1047 "Shape of extracted_valid_mask_blocks"1048 f" {extracted_valid_mask_blocks.shape} is not ({batch_size},"1049 f" {num_query_blocks}, {self.context_size}) after potential reshape."1050 )1051 1052 # 3. Expand dimensions for broadcasting with logits and causal mask.1053 # Target shape for broadcasting with logits [B,N,U,W,C]1054 # extracted_valid_mask_blocks to [B, 1, U, 1, C]1055 condition_from_input_validity = extracted_valid_mask_blocks.unsqueeze(1).unsqueeze(-2)1056 1057 # self.local_causal_valid_mask is [W, C], True where allowed by local window.1058 # Expand to [1, 1, 1, W, C]1059 condition_from_causality = self.local_causal_valid_mask.unsqueeze(0).unsqueeze(0).unsqueeze(0)1060 1061 # 4. Combine the two conditions.1062 # final_condition will be True where a key is *both* originally valid *and* causally accessible.1063 # Broadcasts to [B, 1, U, W, C]1064 final_condition_for_where = torch.logical_and(1065 condition_from_input_validity,1066 condition_from_causality.to(condition_from_input_validity.device), # Ensure same device1067 )1068 1069 # Embed queries and keys1070 logits = self.relative_position_embedding(query_blocks, key_blocks)1071 1072 # Apply attention logit softcap1073 # Ensure softcap is on the same device as logits1074 softcap_val = self.softcap.to(logits.device)1075 logits = logits / softcap_val1076 logits = torch.tanh(logits)1077 logits = logits * softcap_val1078 1079 # Apply the combined mask.1080 # final_condition_for_where will broadcast with logits [B,N,U,W,C]1081 logits = torch.where(final_condition_for_where, logits, torch.finfo(logits.dtype).min)1082 probabilities = torch.nn.functional.softmax(logits, dim=-1, dtype=torch.float32).to(dtype=value_blocks.dtype)1083 1084 # context_vectors is adapted from jax.numpy.einsum("BNuwc,BucNH->BuwNH", ...)1085 b_dim, n_dim, u_dim, w_dim, c_dim = probabilities.shape1086 h_dim = value_blocks.shape[-1]1087 prob_bun = probabilities.permute(0, 2, 1, 3, 4).reshape(-1, w_dim, c_dim)1088 v_bun = value_blocks.permute(0, 1, 3, 2, 4).reshape(-1, c_dim, h_dim)1089 result_bmm = torch.bmm(prob_bun, v_bun)1090 context_vectors = result_bmm.reshape(b_dim, u_dim, n_dim, w_dim, h_dim).permute(0, 1, 3, 2, 4)1091 context_vectors = context_vectors.reshape(1092 (1093 batch_size,1094 num_query_blocks * self.chunk_size,1095 self.num_heads,1096 self.head_dim,1097 )1098 )1099 context_vectors = context_vectors[:, :q_time]1100 1101 return context_vectors1102 1103 1104class Gemma3nAudioCumulativeGroupNorm(nn.Module):1105 """Applies Group Normalization cumulatively over the time dimension.1106 1107 This layer normalizes the input by calculating the mean and variance1108 cumulatively over the time dimension (dim 1). The statistics are computed1109 over all feature dimensions (specified by `feature_dims` and `num_channels`)1110 for elements marked as valid by the optional `mask`.1111 1112 If a `mask` is provided (True for valid, False for invalid/padded),1113 invalid time steps do not contribute to the statistics calculation, and1114 their corresponding output values are zeroed out.1115 1116 Scale and bias, if enabled, are applied per-channel (last dimension).1117 This behavior is similar to JAX's `GroupNormalization` with `num_groups=1`1118 and `cumulative=True`.1119 """1120 1121 def __init__(1122 self,1123 num_channels: int, # Number of channels (size of the last dimension)1124 feature_dims: Sequence[int], # Sizes of non-channel feature dimensions, e.g., (H, W) for input [B,T,H,W,C]1125 eps: float = 1e-3,1126 ):1127 super().__init__()1128 self.num_channels = num_channels1129 self.feature_dims = tuple(feature_dims)1130 self.eps = eps1131 1132 # Scale parameter depends only on the channel dimension1133 self.weight = nn.Parameter(torch.ones(num_channels))1134 1135 # Axes for normalization: all dimensions except Batch (0) and Time (1).1136 # For input [B, T, *feature_dims, C], these are dims from 2 onwards.1137 self.reduction_axes = tuple(range(2, 2 + len(self.feature_dims) + 1))1138 1139 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:1140 """Applies cumulative group norm, optionally using a mask.1141 1142 Args:1143 hidden_states: Input tensor, shape [B, T, *feature_dims, C].1144 1145 Returns:1146 Normalized tensor with the same shape as x.1147 """1148 expected_input_suffix = self.feature_dims + (self.num_channels,)1149 if hidden_states.shape[2:] != expected_input_suffix:1150 raise ValueError(1151 f"Input tensor shape suffix {hidden_states.shape[2:]} does not match expected"1152 f" suffix (feature_dims + num_channels) {expected_input_suffix}"1153 )1154 1155 input_dtype = hidden_states.dtype1156 # Calculations are performed in float32 for numerical stability.1157 calc_dtype = torch.float321158 x_calc = hidden_states.to(calc_dtype)1159 1160 # Prepare a broadcastable mask (`mask_calc`).1161 # If no mask is provided, treat all elements as valid1162 # (mask_calc is all ones).1163 # Otherwise, expand the [B, T] mask to [B, T, 1, ..., 1] for broadcasting.1164 mask_calc = torch.ones_like(x_calc, dtype=calc_dtype)1165 1166 # Cumulative Statistics Calculation1167 # 1. Sum of values over reduction axes at each time step.1168 sum_values_at_t = torch.sum(x_calc, dim=self.reduction_axes, keepdim=True)1169 # 2. Cumulative sum of values over time.1170 cum_sum_values = torch.cumsum(sum_values_at_t, dim=1)1171 1172 # 3. Count of valid elements in the normalization group at each time step.1173 # (A "group" here consists of all features at a given Batch, Time).1174 elements_in_group_at_t = torch.sum(mask_calc, dim=self.reduction_axes, keepdim=True)1175 # 4. Cumulative count of valid elements over time.1176 cum_count_elements = torch.cumsum(elements_in_group_at_t, dim=1)1177 # Avoid division by zero if all preceding elements were masked.1178 safe_cum_count_elements = torch.clamp(cum_count_elements, min=1.0)1179 1180 # 5. Cumulative mean.1181 cum_mean = cum_sum_values / safe_cum_count_elements1182 1183 # 6. Sum of squared differences from the cumulative mean.1184 # Only sum for valid elements: (x_calc - cum_mean)^2 * mask_calc.1185 # Using x_calc here for the difference, as cum_mean already accounts for masking.1186 squared_diff_from_mean = (x_calc - cum_mean).pow(2)1187 sum_sq_diff_at_t = torch.sum(squared_diff_from_mean, dim=self.reduction_axes, keepdim=True)1188 1189 # 7. Cumulative sum of squared differences over time.1190 cum_sum_sq_diff = torch.cumsum(sum_sq_diff_at_t, dim=1)1191 1192 # 8. Cumulative variance.1193 cum_variance = cum_sum_sq_diff / safe_cum_count_elements1194 1195 # Normalize the input using the calculated cumulative statistics:1196 # (x - E[x]) / sqrt(Var[x] + eps)1197 normalized_x = (x_calc - cum_mean) * torch.rsqrt(cum_variance + self.eps)1198 1199 # Apply affine transformation (scale and bias) if enabled.1200 # Scale and bias are applied per-channel (last dimension).