§
    ‚ŠtjÐ ã            	       óˆ  — d dl Z d dlmZ d dlmZmZ d dlmZ d dlm	Z	 d dl
Z
d dlmZ d dlmc mZ d dlmZ ddlmZ dd	lmZ dd
lmZmZ ddlmZ ddlmZmZ ddlm Z m!Z! ddl"m#Z# ddl$m%Z%m&Z& ddl'm(Z( ddl)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/ ddl0m1Z1 ddl2m3Z3 ddl4m5Z5 ddl6m7Z7m8Z8m9Z9m:Z: ddl;m<Z< ddl=m>Z>m?Z?m@Z@mAZAmBZB ddlCmDZDmEZEmFZFmGZG ddlHmIZI  e-¦   «         rd dlJmKZK  e.jL        eM¦  «        ZN e+d¬¦  «        e G d„ de<¦  «        ¦   «         ¦   «         ZO e+d¬¦  «        e G d „ d!e¦  «        ¦   «         ¦   «         ZP e+d¬¦  «        e G d"„ d#eI¦  «        ¦   «         ¦   «         ZQ e+d¬¦  «        e G d$„ d%e¦  «        ¦   «         ¦   «         ZRe+e G d&„ d'e!¦  «        ¦   «         ¦   «         ZS G d(„ d)eG¦  «        ZT G d*„ d+eD¦  «        ZU G d,„ d-ejV        ¦  «        ZW G d.„ d/ejV        ¦  «        ZX G d0„ d1ejV        ¦  «        ZY G d2„ d3ejV        ¦  «        ZZ G d4„ d5ejV        ¦  «        Z[ G d6„ d7ejV        ¦  «        Z\ G d8„ d9ejV        ¦  «        Z] G d:„ d;ejV        ¦  «        Z^ G d<„ d=ejV        ¦  «        Z_ G d>„ d?ejV        ¦  «        Z` G d@„ dAeB¦  «        Za G dB„ dCejV        ¦  «        Zb G dD„ dEe7¦  «        Zc G dF„ dGejV        ¦  «        ZddhdIe
je        dJe
je        dKe
je        dLeffdM„Zg G dN„ dOejV        ¦  «        Zh G dP„ dQe>¦  «        Zi G dR„ dSe8¦  «        Zj G dT„ dUej¦  «        Zk G dV„ dWe@¦  «        Zl e+dX¬Y¦  «         G dZ„ d[eA¦  «        ¦   «         Zm e+d\¬Y¦  «         G d]„ d^e?¦  «        ¦   «         Zn G d_„ d`ejV        ¦  «        Zo e+da¬Y¦  «         G db„ dceF¦  «        ¦   «         Zp e+dd¬Y¦  «         G de„ dfeE¦  «        ¦   «         Zqg dg¢ZrdS )ié    N)ÚUserDict)ÚCallableÚSequence)Ú	dataclass)ÚAny)Ústricté   )Úinitialization)ÚACT2FN)ÚCacheÚDynamicCache)ÚPreTrainedConfig)Úcreate_causal_maskÚ!create_sliding_window_causal_mask)ÚBaseModelOutputWithPastÚBaseModelOutputWithPooling)ÚROPE_INIT_FUNCTIONS)ÚALL_ATTENTION_FUNCTIONSÚPreTrainedModel)ÚUnpack)ÚTransformersKwargsÚauto_docstringÚcan_return_tupleÚis_accelerate_availableÚloggingÚtorch_compilable_check)Úmerge_with_config_defaults)Úcapture_outputsé   )Ú	AutoModel)Ú	Gemma2MLPÚGemma2PreTrainedModelÚeager_attention_forwardÚrotate_half)ÚGemma3TextConfig)ÚGemma3DecoderLayerÚGemma3ForCausalLMÚGemma3RotaryEmbeddingÚGemma3TextModelÚGemma3TextScaledWordEmbedding)ÚPaliGemmaCausalLMOutputWithPastÚ!PaliGemmaForConditionalGenerationÚPaliGemmaModelÚPaligemmaModelOutputWithPast)ÚTimmWrapperConfig)Úadd_hook_to_modulezgoogle/gemma-3n-E4B)Ú
checkpointc                   ó¼  — e Zd ZU dZdZdddddddddddœ
Zddd	œZd
Zee	d<   dZ
ee	d<   dZee	d<   dZee	d<   dZeee         z  e	d<   dZee	d<   dZee	d<   dZee	d<   dZee	d<   dZee         dz  e	d<   dZee	d<   d Zee	d!<   d"Zee	d#<   d$Zee	d%<   d&Zee	d'<   d(Zee	d)<   d*Zee	d+<   dZeee         z  dz  e	d,<    e¦   «         Z  e¦   «         Z! e¦   «         Z"d-„ Z#d.„ Z$dS )/ÚGemma3nTextConfigaL	  
    vocab_size_per_layer_input (`int`, *optional*, defaults to 262144):
        Vocabulary size of the per-layer text embeddings that augment the standard embeddings.
    hidden_size_per_layer_input (`int`, *optional*, defaults to 256):
        Dimension of the hidden representations for per-layer embeddings.
    altup_active_idx (`int`, *optional*, defaults to 0):
        The index of the prediction from which AltUp will compute additional predictions or correct the active prediction.
    altup_coef_clip (`float`, *optional*, defaults to 120.0):
        The maximum amplitude of an AltUp prediction or correction coefficient weight.
    altup_correct_scale (`bool`, *optional*, defaults to `True`):
        If True, apply the `AltUp.correct_output_scale` to the corrected prediction at `altup_active_idx`.
    altup_num_inputs (`int`, *optional*, defaults to 4):
        The number of predictions that AltUp should make given the input sequence.
    num_kv_shared_layers (`int`, *optional*, defaults to 15):
        The number of layers that share KV cache values. During the forward pass, the last `num_kv_shared_layers`
        layers in the model "share" the KV values in that each local and global layer in this range uses the KV
        cache values computed for the last local or global layer, respectively, before entering this range. The
        value should be a multiple of the attention pattern size (see `layer_types` parameter).
    laurel_rank (`int`, *optional*, defaults to 64):
        The intermediate size for the linear projections in the Learned Augmented Residual Layer.
    activation_sparsity_pattern (`Sequence[float]`, *optional*):
        The sparsity factor used to extract the top-k activations for a given layer. The provided Sequence must
        explicitly provide a sparsity value for each layer in the model. By default, the first 10 layers are
        sparse with a sparsity factor of 0.95 and the rest are dense.

    ```python
    >>> from transformers import Gemma3nTextModel, Gemma3nTextConfig

    >>> # Initializing a Gemma3nText gemma3n_text-E4B style configuration
    >>> configuration = Gemma3nTextConfig()

    >>> # Initializing a model from the gemma3n_text-E4B style configuration
    >>> model = Gemma3nTextModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    Úgemma3n_textÚcolwiseÚreplicated_with_grad_allreduceÚrowwise)
zlayers.*.self_attn.q_projzlayers.*.self_attn.k_projzlayers.*.self_attn.v_projzlayers.*.self_attn.q_normzlayers.*.self_attn.k_normzlayers.*.self_attn.v_normzlayers.*.self_attn.o_projzlayers.*.mlp.gate_projzlayers.*.mlp.up_projzlayers.*.mlp.down_projg    €„.Aç     ˆÃ@)ÚglobalÚlocali  Ú
vocab_sizeé   Úvocab_size_per_layer_inputé   Úhidden_sizeé   Úhidden_size_per_layer_inputi @  Úintermediate_sizeé#   Únum_hidden_layersr   Únum_key_value_headsi €  Úmax_position_embeddingsi   Úsliding_windowNÚlayer_typesg      >@Úfinal_logit_softcappingr   Úaltup_active_idxg      ^@Úaltup_coef_clipTÚaltup_correct_scaleé   Úaltup_num_inputsé   Únum_kv_shared_layersé@   Úlaurel_rankÚactivation_sparsity_patternc                 óV  — t          | j        t          ¦  «        r:t          | j        ¦  «        x}| j        k    rt          d| j        › d|› d�¦  «        ‚t          | j        t          ¦  «        s| j        g| j        z  | _        | j        €#d„ t          | j        ¦  «        D ¦   «         | _        | j        €)| j        dk    rdnd}dg|z  dg| j        |z
  z  z   | _        t          | j        ¦  «        x}| j        k    rt          d	| j        › d|› d�¦  «        ‚t          j
        d
i |¤Ž d S )Nzjintermediate_size must have an explicit intermediate size for every layer or one for all layers. Expected z values but got ú.c                 ó.   — g | ]}|d z   dz  dk    rdnd‘ŒS )é   é   r   Úfull_attentionÚsliding_attention© )Ú.0Úis     úi/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/transformers/models/gemma3n/modular_gemma3n.pyú
<listcomp>z3Gemma3nTextConfig.__post_init__.<locals>.<listcomp>§   s>   € ð  ð  ð  ØRS Q¨¡U¨a¡K°1Ò$4Ð$4Ð Ð Ð:Mð ð  ð  ó    é
   r   gffffffî?ç        zeactivation_sparsity_pattern must have an explicit activation sparsity value for every layer.Expected r[   )Ú
isinstancerB   r   ÚlenrD   Ú
ValueErrorrH   ÚrangerS   r   Ú__post_init__)ÚselfÚkwargsÚintsize_lenÚnum_sparse_layersÚlen_asps        r^   rg   zGemma3nTextConfig.__post_init__š   s¤  € å�tÔ-­xÑ8Ô8ð		Wå # DÔ$:Ñ ;Ô ;Ð;�ÀÔ@VÒVÐVåðSØ Ô2ðSð SØDOðSð Sð Sñô ð õ ˜DÔ2µHÑ=Ô=ð 	WØ&*Ô&<Ð%=ÀÔ@VÑ%VˆDÔ"àÔÐ#ð ð  ÝW\Ð]aÔ]sÑWtÔWtð ñ  ô  ˆDÔð Ô+Ð3Ø&*Ô&<¸rÒ&AÐ&A  ÀqÐØ04¨vÐ8IÑ/IÈSÈEØÔ&Ð):Ñ:ñMñ 0ˆDÔ,õ ˜4Ô;Ñ<Ô<Ð<ˆGÀÔAWÒWÐWÝðOØ Ô2ðOð OØDKðOð Oð Oñô ð õ
 	Ô&Ð0Ð0¨Ð0Ð0Ð0Ð0Ð0r`   c                 ór  — |                      dd ¦  «        }ddiddidœ}| j        �| j        n|| _        |� | j        d                              |¦  «         | j                             d¦  «        €ddi| j        d<   | j        d                              d|                      d| j        d         ¦  «        ¦  «         | j                             d¦  «        €ddi| j        d<   | j        d                              d|                      d	| j        d
         ¦  «        ¦  «         |                      ¦   «          |S )NÚrope_scalingÚ	rope_typeÚdefault)rZ   rY   rY   Ú
rope_thetar9   rZ   Úrope_local_base_freqr:   )ÚpopÚrope_parametersÚupdateÚgetÚ
setdefaultÚdefault_thetaÚstandardize_rope_params)rh   ri   rn   Údefault_rope_paramss       r^   Úconvert_rope_params_to_dictz-Gemma3nTextConfig.convert_rope_params_to_dict¹   s_  € Ø—z’z .°$Ñ7Ô7ˆð
 #.¨yÐ!9Ø*¨IÐ6ð
ð 
Ðð 8<Ô7KÐ7W˜tÔ3Ð3Ð]pˆÔØÐ#ØÔ Ð!1Ô2×9Ò9¸,ÑGÔGÐGð Ô×#Ò#Ð$4Ñ5Ô5Ð=Ø6AÀ9Ð5MˆDÔ Ð!1Ñ2ØÔÐ-Ô.×9Ò9Ø˜&Ÿ*š* \°4Ô3EÀhÔ3OÑPÔPñ	
ô 	
ð 	
ð Ô×#Ò#Ð$7Ñ8Ô8Ð@Ø9DÀiÐ8PˆDÔ Ð!4Ñ5ØÔÐ0Ô1×<Ò<Ø˜&Ÿ*š*Ð%;¸TÔ=OÐPWÔ=XÑYÔYñ	
ô 	
ð 	
ð
 	×$Ò$Ñ&Ô&Ð&Øˆr`   )%Ú__name__Ú
__module__Ú__qualname__Ú__doc__Ú
model_typeÚbase_model_tp_planrx   r;   ÚintÚ__annotations__r=   r?   rA   rB   ÚlistrD   rE   rF   rG   rH   ÚstrrI   ÚfloatrJ   rK   rL   ÚboolrN   rP   rR   rS   ÚAttributeErrorÚattn_logit_softcappingÚuse_bidirectional_attentionÚquery_pre_attn_scalarrg   r{   r[   r`   r^   r3   r3   K   sõ  € € € € € € ð%ð %ðN  €Jà%.Ø%.Ø%.Ø%EØ%EØ%EØ%.Ø"+Ø )Ø"+ðð Ðð  +°XÐ>Ð>€Mà€J�ÐÐÑØ&-Ð Ð-Ð-Ñ-Ø€K�ÐÐÑØ'*Ð Ð*Ð*Ñ*Ø)/Ð�s˜T #œY‘Ð/Ð/Ñ/ØÐ�sÐÐÑØ Ð˜Ð Ð Ñ Ø#)Ð˜SÐ)Ð)Ñ)Ø€N�CÐÐÑØ$(€K��c”˜TÑ!Ð(Ð(Ñ(Ø%)Ð˜UÐ)Ð)Ñ)ØÐ�cÐÐÑØ"€O�UÐ"Ð"Ñ"Ø $Ð˜Ð$Ð$Ñ$ØÐ�cÐÐÑØ "Ð˜#Ð"Ð"Ñ"Ø€K�ÐÐÑØ>BÐ ¨¨e¬Ñ!4°tÑ!;ÐBÐBÑBØ+˜^Ñ-Ô-ÐØ"0 .Ñ"2Ô"2ÐØ*˜NÑ,Ô,Ðð1ð 1ð 1ð>ð ð ð ð r`   r3   c                   ó°  — e Zd ZU dZdZdZeed<   dZeed<   dZ	eed<   dZ
eed	<   d
Zeed<   dZeed<   dZeed<   dZeed<   dZeed<   dZeed<   dZeed<   dZeed<   dZeed<   dZeed<   dZeed<   dZee         eeef         z  ed <   d!Zeed"<   d#Zeeeeef         eeef         f         z  ed$<   d%Zeeeeef         eeef         f         z  ed&<   d'S )(ÚGemma3nAudioConfiga‡  
    vocab_offset (`int`, *optional*, defaults to 262272):
        Offset between the tokenizer vocab index for the token ids embedded by `Gemma3nMultimodalEmbedder` and the
        0-indexed `Gemma3nMultimodalEmbedder.embedding` table.
    input_feat_size (`int`, *optional*, defaults to 128):
        The number of channels in each mel-spectrogram frame.
    gradient_clipping (`float`, *optional*, defaults to 10000000000.0):
        Clipping value used to stabilize extremely large gradient values.
    conf_attention_chunk_size (`int`, *optional*, defaults to 12):
        The sub-sequence size for local attention processing inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_attention_context_left (`int`, *optional*, defaults to 13):
        The left context size of the local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_attention_context_right (`int`, *optional*, defaults to 0):
        The right context size of the local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_attention_logit_cap (`float`, *optional*, defaults to 50.0):
        Logit cap applied during local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_num_attention_heads (`int`, *optional*, defaults to 8):
        The number of attention heads in local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_num_hidden_layers (`int`, *optional*, defaults to 12):
        The number of layers that use local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_conv_kernel_size (`int`, *optional*, defaults to 5):
        Convolution kernel size for the conformer block inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_reduction_factor (`int`, *optional*, defaults to 4):
        Reduction factor used in the conformer block inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_residual_weight (`float`, *optional*, defaults to 0.5):
        Residual connection weight inside the Conformer ("conf") section of the
        Universal Speech Model.
    sscp_conv_channel_size (`tuple(int, int)`, *optional*, defaults to `(128, 32)`):
        The channel sizes for the first and second convolutional layers in the Sub-sample Convolution Projection
        ("sscp") section of the Universal Speech Model.
    sscp_conv_group_norm_eps (`float`, *optional*, defaults to 0.001):
        Epsilon used in group normalization in the subsample convolution projection in the Sub-sample Convolution
        Projection ("sscp") section of the Universal Speech Model.
    sscp_conv_kernel_size (`tuple(tuple(int, int), tuple(int, int))`, *optional*, defaults to `((3, 3), (3, 3))`):
        Kernel sizes of the two convolutional layers in the subsample convolution projection  in the Sub-sample
        Convolution Projection ("sscp") section of the Universal Speech Model. The kernel sizes are specified as a
        tuple of height and width for each layer, where the height corresponds to the time dimension and the width
        corresponds to the frequency dimension.
    sscp_conv_stride_size (`tuple(tuple(int, int), tuple(int, int))`, *optional*, defaults to `((2, 2), (2, 2))`):
        Stride sizes of the two convolutional layers in the subsample convolution projection in the Sub-sample
        Convolution Projection ("sscp") section of the Universal Speech Model. The stride sizes are specified as a
        tuple of height and width for each layer, where the height corresponds to the time dimension and the width
        corresponds to the frequency dimension.

    Example:

    ```python
    >>> from transformers import Gemma3nAudioConfig, Gemma3nAudioEncoder

    >>> # Initializing a Gemma3nAudioEncoder gemma3n_audio-E4B-style configuration
    >>> configuration = Gemma3nAudioConfig()

    >>> # Initializing a model from the gemma3n_audio-E4B style configuration
    >>> model = Gemma3nAudioEncoder(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    Úgemma3n_audioé€   r;   é€  Úvocab_offsetÚinput_feat_sizei   r?   ç�íµ ÷Æ°>Úrms_norm_epsg    _ BÚgradient_clippingé   Úconf_attention_chunk_sizeé   Úconf_attention_context_leftr   Úconf_attention_context_rightg      I@Úconf_attention_logit_capé   Úconf_num_attention_headsÚconf_num_hidden_layersrX   Úconf_conv_kernel_sizerM   Úconf_reduction_factorç      à?Úconf_residual_weight)r�   é    Ússcp_conv_channel_sizeçü©ñÒMbP?Ússcp_conv_group_norm_eps)©r	   r	   r§   Ússcp_conv_kernel_size)©r   r   r©   Ússcp_conv_stride_sizeN)r|   r}   r~   r   r€   r;   r‚   rƒ   r‘   r’   r?   r”   r†   r•   r—   r™   rš   r›   r�   rž   rŸ   r    r¢   r¤   r„   Útupler¦   r¨   rª   r[   r`   r^   r�   r�   ×   sÉ  € € € € € € ðBð BðH !€Jà€J�ÐÐÑØ%€L�#Ð%Ð%Ñ%Ø€O�SÐÐÑØ€K�ÐÐÑØ€L�%ÐÐÑØ/Ð�uÐ/Ð/Ñ/Ø%'Ð˜sÐ'Ð'Ñ'Ø')Ð Ð)Ð)Ñ)Ø()Ð  #Ð)Ð)Ñ)Ø&*Ð˜eÐ*Ð*Ñ*Ø$%Ð˜cÐ%Ð%Ñ%Ø"$Ð˜CÐ$Ð$Ñ$Ø!"Ð˜3Ð"Ð"Ñ"Ø!"Ð˜3Ð"Ð"Ñ"Ø"%Ð˜%Ð%Ð%Ñ%Ø:CÐ˜D œI¨¨c°3¨h¬Ñ7ÐCÐCÑCØ&*Ð˜eÐ*Ð*Ñ*ðMÐ˜4 %¨¨c°3¨h¬¸¸sÀC¸x¼Ð(HÔ"IÑIð ð ñ ðMÐ˜4 %¨¨c°3¨h¬¸¸sÀC¸x¼Ð(HÔ"IÑIð ð ñ ð ð r`   r�   c                   óŽ   — e Zd ZU dZdZdZeed<   dZe	ed<   dZ
eed<   d	Zeed
<   dZeed<   dZeed<   dZeed<   dZedz  ed<   dS )ÚGemma3nVisionConfiga³  
    architecture (`str`, *optional*, defaults to `"resnet50"`):
        The timm architecture to load.
    do_pooling (`bool`, *optional*, defaults to `True`):
        Whether to do pooling for the last_hidden_state in `TimmWrapperModel` or not.
    model_args (`dict[str, Any]`, *optional*):
        Additional keyword arguments to pass to the `timm.create_model` function. e.g. `model_args={"depth": 3}`
        for `timm/vit_base_patch32_clip_448.laion2b_ft_in12k_in1k` to create a model with 3 blocks. Defaults to `None`.
    vocab_offset (`int`, *optional*, defaults to 262144):
        Offset between the tokenizer vocab index for the token ids embedded by `Gemma3nMultimodalEmbedder` and the
        0-indexed `Gemma3nMultimodalEmbedder.embedding` table.

    Example:
    ```python
    >>> from transformers import Gemma3nVisionConfig, TimmWrapper

    >>> # Initializing a TimmWrapper gemma3n_vision-E4B-style configuration
    >>> configuration = Gemma3nVisionConfig()

    >>> # Initializing a gemma3n_vision-E4B-style TimmWrapper from the configuration
    >>> model = TimmWrapper(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    Úgemma3n_visionç{®Gáz”?Úinitializer_rangeFÚ
do_poolingÚmobilenetv5_300m_encÚarchitecturer>   r?   r�   r;   r<   r‘   r“   r”   NÚ
model_args)r|   r}   r~   r   r€   r°   r†   rƒ   r±   r‡   r³   r…   r?   r‚   r;   r‘   r”   r´   Údictr[   r`   r^   r­   r­   ;  s¨   € € € € € € ðð ð6 "€Jà#Ð�uÐ#Ð#Ñ#Ø€J�ÐÐÑØ.€L�#Ð.Ð.Ñ.Ø€K�ÐÐÑØ€J�ÐÐÑØ€L�#ÐÐÑØ€L�%ÐÐÑØ"€J��t‘Ð"Ð"Ñ"Ð"Ð"r`   r­   c                   óˆ  ‡ — e Zd ZU dZdZeeedœZdZ	ee
eef         z  dz  ed<   dZee
eef         z  dz  ed<   dZee
eef         z  dz  ed<   dZedz  ed	<   d
Zedz  ed<   dZedz  ed<   dZedz  ed<   dZedz  ed<   dZedz  ed<   dZedz  ed<   dZedz  ed<   dZedz  ed<   dZedz  ed<   dZeed<   ˆ fd„Zˆ xZS )ÚGemma3nConfigaæ  
    audio_soft_tokens_per_image (`int`, *optional*, defaults to 188):
        The number of soft tokens per audio clip.
    vision_soft_tokens_per_image (`int`, *optional*, defaults to 256):
        The number of soft tokens per image.
    boi_token_id (`int`, *optional*, defaults to 255999):
        The begin-of-image token index to wrap the image prompt.
    eoi_token_id (`int`, *optional*, defaults to 262144):
        The end-of-image token index to wrap the image prompt.
    boa_token_id (`int`, *optional*, defaults to 256000):
        The begin-of-audio token index to wrap the audio prompt.
    eoa_token_id (`int`, *optional*, defaults to 262272):
        The end-of-audio token index to wrap the audio prompt.

    Example:

    ```python
    >>> from transformers import Gemma3nForConditionalGeneration, Gemma3nConfig, Gemma3nTextConfig

    >>> # Initializing a MobileNet vision config, which is loaded from TIMM
    >>> vision_config = Gemma3nVisionConfig()

    >>> # Initializing a Gemma3n Audio config
    >>> audio_config = Gemma3nAudioConfig()

    >>> # Initializing a Gemma3n Text config
    >>> text_config = Gemma3nTextConfig()

    >>> # Initializing a Gemma3n gemma-3-4b style configuration
    >>> configuration = Gemma3nConfig(text_config, vision_config, audio_config)

    >>> # Initializing a model from the gemma-3-4b style configuration
    >>> model = Gemma3nTextConfig(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```Úgemma3n)Útext_configÚvision_configÚaudio_configNr¹   rº   r»   é¼   Úaudio_soft_tokens_per_imager@   Úvision_soft_tokens_per_imageiÿç Úboi_token_idr<   Úeoi_token_idi  Úimage_token_idi è Úboa_token_idr�   Úeoa_token_idi�  Úaudio_token_idr¯   r°   TÚtie_word_embeddingsÚ	use_cachec                 ó˜  •— | j         €.t          ¦   «         | _         t                               d¦  «         n0t	          | j         t
          ¦  «        rt          di | j         ¤Ž| _         t	          | j        t
          ¦  «        rt          di | j        ¤Ž| _        n4| j        €-t          ¦   «         | _        t                               d¦  «         t	          | j        t
          ¦  «        rt          di | j        ¤Ž| _        n4| j        €-t          ¦   «         | _        t                               d¦  «          t          ¦   «         j        di |¤Ž d S )NzAtext_config is None, using default Gemma3nTextConfig text config.zGvision_config is None, using default Gemma3nVisionConfig vision config.z7audio_config is None. Using default Gemma3nAudioConfig.r[   )r¹   r3   ÚloggerÚinforc   rµ   rº   r­   r»   r�   Úsuperrg   )rh   ri   Ú	__class__s     €r^   rg   zGemma3nConfig.__post_init__¤  s:  ø€ ØÔÐ#Ý0Ñ2Ô2ˆDÔÝ�KŠKÐ[Ñ\Ô\Ð\Ð\Ý˜Ô(­$Ñ/Ô/ð 	EÝ0ÐDÐD°4Ô3CÐDÐDˆDÔå�dÔ(­$Ñ/Ô/ð 	cÝ!4Ð!JÐ!J°tÔ7IÐ!JÐ!JˆDÔÐØÔÐ'Ý!4Ñ!6Ô!6ˆDÔÝ�KŠKÐaÑbÔbÐbå�dÔ'­Ñ.Ô.ð 	SÝ 2Ð GÐ G°TÔ5FÐ GÐ GˆDÔÐØÔÐ&Ý 2Ñ 4Ô 4ˆDÔÝ�KŠKÐQÑRÔRÐRà�‰ŒÔÐ'Ð' Ð'Ð'Ð'Ð'Ð'r`   ) r|   r}   r~   r   r€   r3   r­   r�   Úsub_configsr¹   rµ   r…   r   rƒ   rº   r»   r½   r‚   r¾   r¿   rÀ   rÁ   rÂ   rÃ   rÄ   r°   r†   rÅ   r‡   rÆ   rg   Ú__classcell__©rË   s   @r^   r·   r·   e  s©  ø€ € € € € € ð$ð $ðL €Jà(Ø,Ø*ðð €Kð >B€KÐ" T¨#¨s¨(¤^Ñ3°dÑ:ÐAÐAÑAØAE€MÐ&¨¨c°3¨h¬Ñ7¸$Ñ>ÐEÐEÑEØ?C€LÐ$ t¨C°¨H¤~Ñ5¸Ñ<ÐCÐCÑCØ.1Ð  t¡Ð1Ð1Ñ1Ø/2Ð  #¨¡*Ð2Ð2Ñ2Ø&€L�#˜‘*Ð&Ð&Ñ&Ø&€L�#˜‘*Ð&Ð&Ñ&Ø!(€N�C˜$‘JÐ(Ð(Ñ(Ø&€L�#˜‘*Ð&Ð&Ñ&Ø&€L�#˜‘*Ð&Ð&Ñ&Ø!(€N�C˜$‘JÐ(Ð(Ñ(Ø&*Ð�u˜t‘|Ð*Ð*Ñ*Ø'+Ð˜ ™Ð+Ð+Ñ+Ø€IˆtÐÐÑð(ð (ð (ð (ð (ð (ð (ð (ð (r`   r·   c                   ó2   — e Zd ZU dZdZej        dz  ed<   dS )ÚGemma3nAudioEncoderModelOutputzy
    audio_mel_mask (`torch.BoolTensor`, *optional*):
        A torch.BoolTensor of shape `(batch_size, num_frames)`
    NÚaudio_mel_mask)r|   r}   r~   r   rÑ   ÚtorchÚ
BoolTensorrƒ   r[   r`   r^   rÐ   rÐ   º  s6   € € € € € € ðð ð
 /3€N�EÔ$ tÑ+Ð2Ð2Ñ2Ð2Ð2r`   rÐ   c                   ó2   — e Zd ZU dZdZej        dz  ed<   dS )ÚGemma3nModelOutputWithPastaÝ  
    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    image_hidden_states (`torch.FloatTensor`, *optional*):
        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
    audio_hidden_states (`torch.FloatTensor`, *optional*):
        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
        audio_hidden_states of the model produced by the audio encoder and after projecting the last hidden state.
    NÚaudio_hidden_states©r|   r}   r~   r   rÖ   rÒ   ÚFloatTensorrƒ   r[   r`   r^   rÕ   rÕ   Å  s7   € € € € € € ðð ð 59Ð˜Ô*¨TÑ1Ð8Ð8Ñ8Ð8Ð8r`   rÕ   c                   ó2   — e Zd ZU dZdZej        dz  ed<   dS )ÚGemma3nCausalLMOutputWithPastaF  
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Language modeling loss (for next-token prediction).
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.text_config.vocab_size)`):
        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    image_hidden_states (`torch.FloatTensor`, *optional*):
        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
        image_hidden_states of the model produced by the vision encoder after projecting last hidden state.
    audio_hidden_states (`torch.FloatTensor`, *optional*):
        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
        audio_hidden_states of the model produced by the audio encoder and after projecting the last hidden state.
    NrÖ   r×   r[   r`   r^   rÚ   rÚ   ×  s7   € € € € € € ðð ð$ 59Ð˜Ô*¨TÑ1Ð8Ð8Ñ8Ð8Ð8r`   rÚ   c                   óh   ‡ — e Zd Zddededefˆ fd„Zdej        fd„Z	dej        d	ej        fd
„Z
ˆ xZS )ÚGemma3nRMSNormr“   TÚdimÚepsÚ
with_scalec                 óÐ   •— t          ¦   «                              ¦   «          || _        || _        | j        r/t	          j        t          j        |¦  «        d¬¦  «        | _        d S d S )NT)Úrequires_grad)	rÊ   Ú__init__rÞ   rß   ÚnnÚ	ParameterrÒ   ÚonesÚweight)rh   rÝ   rÞ   rß   rË   s       €r^   râ   zGemma3nRMSNorm.__init__î  s`   ø€ Ý‰Œ×ÒÑÔÐØˆŒØ$ˆŒàŒ?ð 	LÝœ,¥u¤z°#¡¤ÀdÐKÑKÔKˆDŒKˆKˆKð	Lð 	Lr`   Úhidden_statesc                 ó–   — |                      d¦  «                             dd¬¦  «        | j        z   }|t          j         |d¦  «        z  S )Nr   éÿÿÿÿT)Úkeepdimç      à¿)ÚpowÚmeanrÞ   rÒ   )rh   rç   Úmean_squareds      r^   Ú_normzGemma3nRMSNorm._normö  sF   € Ø$×(Ò(¨Ñ+Ô+×0Ò0°¸TÐ0ÑBÔBÀTÄXÑMˆà�uœy¨°tÑ<Ô<Ñ<Ð<r`   Úreturnc                 óÀ   — |                       |                     ¦   «         ¦  «        }| j        r|| j                             ¦   «         z  }|                     |¦  «        S ©N)rï   r†   rß   ræ   Útype_as)rh   rç   Únormed_outputs      r^   ÚforwardzGemma3nRMSNorm.forwardû  sV   € ØŸ
š
 =×#6Ò#6Ñ#8Ô#8Ñ9Ô9ˆØŒ?ð 	@Ø)¨D¬K×,=Ò,=Ñ,?Ô,?Ñ?ˆMØ×$Ò$ ]Ñ3Ô3Ð3r`   )r“   T)r|   r}   r~   r‚   r†   r‡   râ   rÒ   ÚTensorrï   rõ   rÍ   rÎ   s   @r^   rÜ   rÜ   í  s¤   ø€ € € € € ðLð L˜Cð L eð LÀð Lð Lð Lð Lð Lð Lð= 5¤<ð =ð =ð =ð =ð
4 U¤\ð 4°e´lð 4ð 4ð 4ð 4ð 4ð 4ð 4ð 4r`   rÜ   c                   óÄ   ‡ — e Zd Zdefˆ fd„Zdej        dej        dej        fd„Zdej        de	d	e	d
e	de	de	de	dej        fd„Z
dej        dej        dej        fd„Zˆ xZS )Ú%Gemma3nAudioRelativePositionEmbeddingÚconfigc                 ó$  •— t          ¦   «                              ¦   «          || _        | j        j        | _        | j        j        | _        | j        | j        z  | _        t          d| j        j	        dz
  ¦  «        | _
        | j        j        | _        t          j        | j        | j        | j        z  d¬¦  «        | _        d}d}| j        dz  }t!          j        t%          |¦  «        t%          |¦  «        z  ¦  «        t          |dz
  d¦  «        z  }|t'          j        t'          j        |¦  «        | z  ¦  «        z  }|                      d|                     ¦   «                              d¦  «                             d¦  «        d¬	¦  «         d S )
Nr   rW   F©Úbiasç      ð?r8   r   Úinv_timescales©Ú
persistent)rÊ   râ   rù   r�   Ú	num_headsr?   ÚchannelsÚhead_dimÚmaxr™   Úmax_backwardrš   Úmax_forwardrã   ÚLinearÚpos_projÚmathÚlogr†   rÒ   ÚexpÚarangeÚregister_bufferÚ	unsqueeze)rh   rù   Úmin_timescaleÚmax_timescaleÚnum_timescalesÚlog_timescale_incrementrþ   rË   s          €r^   râ   z.Gemma3nAudioRelativePositionEmbedding.__init__  sg  ø€ Ý‰Œ×ÒÑÔÐØˆŒàœÔ=ˆŒØœÔ/ˆŒØœ¨¬Ñ7ˆŒÝ  4¤;Ô#JÈQÑ#NÑOÔOˆÔØœ;ÔCˆÔåœ	 $¤-°´À$Ä-Ñ1OÐV[Ð\Ñ\Ô\ˆŒàˆØˆØœ¨!Ñ+ˆÝ"&¤(­5°Ñ+?Ô+?Å%ÈÑBVÔBVÑ+VÑ"WÔ"WÕZ]Ð^lÐopÑ^pÐrsÑZtÔZtÑ"tÐØ&­¬µ5´<ÀÑ3OÔ3OÐSjÐRjÑ3jÑ)kÔ)kÑkˆØ×ÒØØ× Ò Ñ"Ô"×,Ò,¨QÑ/Ô/×9Ò9¸!Ñ<Ô<Øð 	ñ 	
ô 	
ð 	
ð 	
ð 	
r`   ÚpositionÚdtyperð   c                 óN  — |                      ¦   «                              d¦  «        }|| j                             |j        t
          j        ¬¦  «        z  }t          j        t          j        |¦  «        t          j	        |¦  «        gd¬¦  «        }| 
                    |¦  «        S )Nré   )Údevicer  ©rÝ   )r†   r  rþ   Útor  rÒ   Úfloat32ÚcatÚsinÚcosÚtype)rh   r  r  Úscaled_timeÚtiming_signals        r^   Ú_get_timing_signal_1d_posz?Gemma3nAudioRelativePositionEmbedding._get_timing_signal_1d_pos  s‡   € Ø—>’>Ñ#Ô#×-Ò-¨bÑ1Ô1ˆØ Ô!4×!7Ò!7¸x¼ÕV[ÔVcÐ!7Ñ!dÔ!dÑdˆÝœ	¥5¤9¨[Ñ#9Ô#9½5¼9À[Ñ;QÔ;QÐ"RÐXZÐ[Ñ[Ô[ˆØ×!Ò! %Ñ(Ô(Ð(r`   Úterm_bd_before_shiftÚ
batch_sizer  Únum_query_blocksÚquery_block_sizeÚkey_context_sizeÚmax_span_plus_1c                 óþ   — |dz   |z
  }d|f}	t           j                             ||	¦  «        }
|
                     |||||dz   z  f¦  «        }|dd…dd…dd…d||z  …f         }|                     |||||f¦  «        }|S )aZ  Performs the relative shift.

        Args:
          term_bd_before_shift: Tensor of shape [B, N, U, W, F_span]. batch_size
            (B), num_heads (N), num_query_blocks (U), query_block_size (W),
            key_context_size (C = W+L+R), max_span_plus_1 (F_span = L+R+1).

        Returns:
          Tensor of shape [B, N, U, W, C].
        rW   r   N)rã   Ú
functionalÚpadÚreshape)rh   r!  r"  r  r#  r$  r%  r&  Úpad_amount_last_dimÚpadding_tupleÚterm_bd_paddedÚterm_bd_reshapedÚterm_bd_slicedÚterm_bd_shifteds                 r^   Ú_relative_shiftz5Gemma3nAudioRelativePositionEmbedding._relative_shift#  sÎ   € ð4  0°!Ñ3°ÑFÐð Ð/Ð0ˆåœ×*Ò*Ð+?ÀÑOÔOˆð
 *×1Ò1àØØ Ø Ð$4°qÑ$8Ñ9ð	ñ
ô 
Ðð *¨!¨!¨!¨Q¨Q¨Q°°°Ð3XÐ5EÐHXÑ5XÐ3XÐ*XÔYˆð )×0Ò0àØØ Ø Ø ðñ
ô 
ˆð Ðr`   ÚqueriesÚkeysc           	      óX  — |j         \  }}}}}|j         \  }}}	}}t          j        | j        | j         dz
  d|j        ¬¦  «                             d¦  «        }
|
j         d         }|                      |
|j        ¬¦  «        }|  	                    |¦  «        }| 
                    d|| j        | j        ¦  «                             d¦  «        }|                     ddddd¦  «        }|                     ddddd¦  «        }t          j        ||¦  «        }|                     ddddd¦  «        }|                     ddd¦  «        }| 
                    ||||z  |¦  «        }t          j        ||¦  «        }| 
                    |||||¦  «        }|                      ||||||	|¦  «        }||z   S )	NrW   ré   ©r  r   ©r  r	   r   rM   )ÚshaperÒ   r  r  r  r  r  r   r  r  r*  r  r  ÚsqueezeÚpermuteÚmatmulr1  )rh   r2  r3  r"  r#  r$  r  r  Ú_r%  Úpos_indicesr&  Úsin_emb_timing_signalÚprojected_sin_embÚsin_embÚ	queries_pÚkeys_p_tÚterm_acÚ
q_permutedÚ
s_permutedÚ
q_reshapedÚterm_bd_unshifed_matmulÚterm_bd_unshifedr0  s                           r^   rõ   z-Gemma3nAudioRelativePositionEmbedding.forward`  sò  € ð OVÌmÑKˆ
Ð$Ð&6¸	À8Ø'+¤zÑ$ˆˆ1Ð  1õ ”l 4Ô#4°tÔ7GÐ6GÈ!Ñ6KÈRÐX_ÔXfÐgÑgÔg×qÒqØñ
ô 
ˆð &Ô+¨AÔ.ˆà $× >Ò >Ø˜wœ}ð !?ñ !
ô !
Ðð
 !ŸMšMÐ*?Ñ@Ô@Ðà#×+Ò+¨A¨ÀÄÐPTÔP]Ñ^Ô^×fÒfØñ
ô 
ˆð —O’O A q¨!¨Q°Ñ2Ô2ˆ	Ø—<’<  1 a¨¨AÑ.Ô.ˆÝ”,˜y¨(Ñ3Ô3ˆð —_’_ Q¨¨1¨a°Ñ3Ô3ˆ
ð —_’_ Q¨¨1Ñ-Ô-ˆ
ð  ×'Ò'¨
°IÐ?OÐRbÑ?bÐdlÑmÔmˆ
õ
 #(¤,¨z¸:Ñ"FÔ"FÐð 3×:Ò:ØØØØØñ
ô 
Ðð ×.Ò.ØØØØØØØñ
ô 
ˆð ˜Ñ(Ð(r`   )r|   r}   r~   r�   râ   rÒ   rö   r  r   r‚   r1  rõ   rÍ   rÎ   s   @r^   rø   rø     s  ø€ € € € € ð
Ð1ð 
ð 
ð 
ð 
ð 
ð 
ð.)°%´,ð )ÀuÄ{ð )ÐW\ÔWcð )ð )ð )ð )ð;à#œlð;ð ð;ð ð	;ð
 ð;ð ð;ð ð;ð ð;ð 
Œð;ð ;ð ;ð ;ðzL)˜uœ|ð L)°5´<ð L)ÀEÄLð L)ð L)ð L)ð L)ð L)ð L)ð L)ð L)r`   rø   c                   óÐ   ‡ — e Zd Zdefˆ fd„Zd„ Zdej        dededej        fd„Z	d	ej        dej        fd
„Z
d	ej        dej        fd„Zd	ej        dej        dej        fd„Zˆ xZS )ÚGemma3nAudioAttentionrù   c                 ó   •— t          ¦   «                              ¦   «          || _        | j        j        | _        | j        j        | _        | j        | j        z  | _        | j        j        | _        | j        j	        | _
        t          d| j        j        dz
  ¦  «        | _        | j        j        | _        | j        | j        z   | j
        z   | _        t#          |¦  «        | _        t'          j        t+          j        | j        f¦  «        ¦  «        | _        t'          j        | j        | j        | j        z  d¬¦  «        | _        t'          j        | j        | j        | j        z  d¬¦  «        | _        t'          j        | j        | j        | j        z  d¬¦  «        | _        | j        dz  }dt*          j        j                             t+          j        d¦  «        ¦  «        z  }|                      d||z                        ¦   «          !                    ¦   «         d¬	¦  «         |  "                    ¦   «         }|                      d
|d¬	¦  «         |                      dt+          j        | j        ¦  «         #                    ¦   «         d¬	¦  «         d S )Nr   rW   Frû   rë   rý   rb   Úq_scalerÿ   Úlocal_causal_valid_maskÚsoftcap)$rÊ   râ   rù   r�   r  r?   r  r—   Ú
chunk_sizerš   Úmax_future_horizonr  r™   Úmax_past_horizonr›   Úattention_logits_soft_capÚcontext_sizerø   Úrelative_position_embeddingrã   rä   rÒ   ÚzerosÚper_dim_scaler  Úq_projÚk_projÚv_projr(  ÚsoftplusÚtensorr  ÚcloneÚdetachÚcreate_local_causal_valid_maskr†   )rh   rù   rK  Úr_softplus_0rL  rË   s        €r^   râ   zGemma3nAudioAttention.__init__°  s-  ø€ Ý‰Œ×ÒÑÔÐØˆŒàœÔ=ˆŒØœ;Ô2ˆÔØÔ(¨D¬NÑ:ˆŒàœ+Ô?ˆŒØ"&¤+Ô"JˆÔÝ # A t¤{Ô'NÐQRÑ'RÑ SÔ SˆÔØ)-¬Ô)MˆÔ&Ø œO¨dÔ.CÑCÀdÔF]Ñ]ˆÔå+PÐQWÑ+XÔ+XˆÔ(Ýœ\­%¬+°t´}Ð6FÑ*GÔ*GÑHÔHˆÔå”i Ô 0°$´.À4Ä=Ñ2PÐW\Ð]Ñ]Ô]ˆŒÝ”i Ô 0°$´.À4Ä=Ñ2PÐW\Ð]Ñ]Ô]ˆŒÝ”i Ô 0°$´.À4Ä=Ñ2PÐW\Ð]Ñ]Ô]ˆŒà”- Ñ%ˆØ�UœXÔ0×9Ò9½%¼,ÀsÑ:KÔ:KÑLÔLÑLˆØ×Ò˜Y¨°<Ñ)?×(FÒ(FÑ(HÔ(H×(OÒ(OÑ(QÔ(QÐ^cÐÑdÔdÐdà"&×"EÒ"EÑ"GÔ"GÐØ×ÒÐ6Ð8OÐ\aÐÑbÔbÐbà×ÒØÝŒL˜Ô7Ñ8Ô8×>Ò>Ñ@Ô@Øð 	ñ 	
ô 	
ð 	
ð 	
ð 	
r`   c                 ó’  — t          j        t          j        | j        | j        ft           j        ¬¦  «        d¬¦  «        j        }t          j        t          j        | j        | j        ft           j        ¬¦  «        | j        | j        z   ¬¦  «        }t          j        | j        | j        ft           j        ¬¦  «        }||z  |z  }|S )Nr6  r   )Údiagonal)	rÒ   Útrilrå   rR  rN  r‡   ÚTrP  rO  )rh   Úlower_causal_maskÚupper_causal_maskrL  s       r^   r]  z4Gemma3nAudioAttention.create_local_causal_valid_maskÒ  sÃ   € Ý!œJÝŒJ˜Ô)¨4¬?Ð;Å5Ä:ÐNÑNÔNØð
ñ 
ô 
ô ð 	õ "œJÝŒJ˜œ¨Ô):Ð;Å5Ä:ÐNÑNÔNØÔ*¨TÔ-DÑDð
ñ 
ô 
Ðõ #(¤*¨d¬o¸tÔ?PÐ-QÕY^ÔYcÐ"dÑ"dÔ"dÐØ"9Ð<MÑ"MÐPaÑ"aÐØ&Ð&r`   ÚxÚpad_leftÚ	pad_rightrð   c                 ó´   — |j         ^}}}|                     ||g|¢R ¦  «        }|                     ||g|¢R ¦  «        }t          j        |||gd¬¦  «        }|S )NrW   r  )r7  Ú	new_zerosrÒ   r  )	rh   re  rf  rg  Úbatchr;  Ú
tail_shapeÚleftÚrights	            r^   Ú	_pad_dim1zGemma3nAudioAttention._pad_dim1ß  sl   € Ø !¤Ðˆˆq�:Ø�{Š{˜E 8Ð9¨jÐ9Ð9Ñ:Ô:ˆØ—’˜U IÐ;°
Ð;Ð;Ñ<Ô<ˆÝŒI�t˜Q Ð&¨AÐ.Ñ.Ô.ˆØˆr`   rç   c                 ó$  — |j         }|dd…         \  }}|| j        z   dz
  | j        z  }|| j        z  |z
  x}dk    r|                      |d|¦  «        }||| j        f|dd…         z   }|                     |¦  «                             ¦   «         }|S )aE  Turns a sequence to non overlapping blocks.

        Args:
            hidden_states: a tensor of [batch, time, ...].

        Returns:
            A tensor of [batch, num_blocks, block_size, ...], with necessary
            paddings,
            where output[:, i, ...] are x[:, i*block_size:(i+1)*block_size, ...].
        Nr   rW   r   )r7  rN  rn  r*  Ú
contiguous)rh   rç   r7  ÚbÚtÚ
num_blocksÚpadding_lenÚpermute_dimss           r^   Ú_convert_to_blockz'Gemma3nAudioAttention._convert_to_blockæ  s¦   € ð Ô#ˆØ�R�a�RŒy‰ˆˆ1Ø˜$œ/Ñ)¨AÑ-°$´/ÑAˆ
à%¨¬Ñ7¸!Ñ;Ð;ˆK¸qÒ@Ð@Ø ŸNšN¨=¸!¸[ÑIÔIˆMà˜: t¤Ð7¸%ÀÀÀ¼)ÑCˆØ%×-Ò-¨lÑ;Ô;×FÒFÑHÔHˆØÐr`   c                 ó0  — | j         }| j        | j        z   dz
  }|                      |||¦  «        }| j        }| j        }|                     d||¬¦  «        }|j        dk    r"|j        dk    rt          j        |dd¬¦  «        }| 	                    ¦   «         S )aã  Extracts temporal context for every block.

        Args:
            hidden_states: a tensor of [batch, time, ...].

        Returns:
            A tensor of [batch, num_blocks, context_size, ...], with necessary
            paddings,
            where context_size = block_size + left_context + right_context,
            and output[:, i, ...] are x[:, start-left_context:end+right_context,
            ...],
            start = i * block_size, end = (i + 1) * block_size.
        rW   )Ú	dimensionÚsizeÚstepr   r	   ré   )ÚsourceÚdestination)
rP  rO  rN  rn  rR  ÚunfoldÚndimrÒ   Úmovedimrp  )rh   rç   rf  rg  Ú	frame_lenÚ
frame_stepÚ
x_unfoldeds          r^   Ú_extract_block_contextz,Gemma3nAudioAttention._extract_block_contextü  s¨   € ð Ô(ˆð Ô+¨d¬oÑ=ÀÑAˆ	ØŸš }°hÀ	ÑJÔJˆàÔ%ˆ	Ø”_ˆ
ð #×)Ò)°A¸IÈJÐ)ÑWÔWˆ
ð Ô Ò!Ð! j¤o¸Ò&9Ð&9õ œ z¸"È!ÐLÑLÔLˆJà×$Ò$Ñ&Ô&Ð&r`   Úmaskc                 óò  — g |j         d d…         ¢| j        ‘| j        ‘R }|                      |¦  «                             |¦  «                             ¦   «         }|                      |¦  «                             |¦  «                             ¦   «         }|                      |¦  «                             |¦  «                             ¦   «         }t          j	        j
                             | j        ¦  «        }ddd| j        f}|                     |¦  «        }	|| j        z  |	z  }|j         d d…         \  }
}|                      |¦  «        }|                      |¦  «        }|                      |¦  «        }|j         d         }| }|                      |¦  «        }|j        dk    r@|j         d         |j         d         z  | j        k    r|                     |
|| j        ¦  «        }|j         |
|| j        fk    r&t'          d|j         › d|
› d|› d| j        › d	�	¦  «        ‚|                     d¦  «                             d
¦  «        }| j                             d¦  «                             d¦  «                             d¦  «        }t          j        ||                     |j        ¦  «        ¦  «        }|                      ||¦  «        }| j                             |j        ¦  «        }||z  }t          j        |¦  «        }||z  }t          j        ||t          j        |j        ¦  «        j        ¦  «        }t          j	        j
                              |dt          j!        ¬¦  «                             |j        ¬¦  «        }|j         \  }}}}}|j         d         }| "                    ddddd¦  «                             d||¦  «        }| "                    ddddd¦  «                             d||¦  «        }t          j#        ||¦  «        } |                      |||||¦  «         "                    ddddd¦  «        }!|!                     |
|| j$        z  | j        | j        f¦  «        }!|!d d …d |…f         }!|!S )Nré   rW   r   rM   r	   z%Shape of extracted_valid_mask_blocks z	 is not (z, z) after potential reshape.éþÿÿÿr   )rÝ   r  r6  )%r7  r  r  rV  r*  rp  rW  rX  rÒ   rã   r(  rY  rU  ÚviewrK  rv  rƒ  r~  rR  re   r  rL  Úlogical_andr  r  rS  rM  ÚtanhÚwhereÚfinfor  ÚminÚsoftmaxr  r9  ÚbmmrN  )"rh   rç   r„  Ú	qkv_shapeÚquery_statesÚ
key_statesÚvalue_statesÚper_dim_scale_spÚbroadcast_shapeÚper_dim_scale_sp_broadcastr"  Úq_timeÚquery_blocksÚ
key_blocksÚvalue_blocksr#  Úoriginal_valid_maskÚextracted_valid_mask_blocksÚcondition_from_input_validityÚcondition_from_causalityÚfinal_condition_for_whereÚlogitsÚsoftcap_valÚprobabilitiesÚb_dimÚn_dimÚu_dimÚw_dimÚc_dimÚh_dimÚprob_bunÚv_bunÚ
result_bmmÚcontext_vectorss"                                     r^   rõ   zGemma3nAudioAttention.forward,  sž  € àN�mÔ)¨#¨2¨#Ô.ÐN°´ÐNÀÄÐNÐNˆ	Ø—{’{ =Ñ1Ô1×9Ò9¸)ÑDÔD×OÒOÑQÔQˆØ—[’[ Ñ/Ô/×7Ò7¸	ÑBÔB×MÒMÑOÔOˆ
Ø—{’{ =Ñ1Ô1×9Ò9¸)ÑDÔD×OÒOÑQÔQˆå œ8Ô.×7Ò7¸Ô8JÑKÔKÐà˜a  D¤MÐ2ˆØ%5×%:Ò%:¸?Ñ%KÔ%KÐ"Ø# d¤lÑ2Ð5OÑOˆà)Ô/°°°Ô3Ñˆ
�Fà×-Ò-¨lÑ;Ô;ˆØ×0Ò0°Ñ<Ô<ˆ
Ø×2Ò2°<Ñ@Ô@ˆØ'Ô-¨aÔ0Ðð  $˜eÐð '+×&AÒ&AÐBUÑ&VÔ&VÐ#ð (Ô,°Ò1Ð1Ø+Ô1°!Ô4Ð7RÔ7XÐYZÔ7[Ñ[Ð_cÔ_pÒpÐpà*E×*MÒ*MØÐ,¨dÔ.?ñ+ô +Ð'ð 'Ô,ØØØÔð1
ò 
ð 
õ
 ðVØ/Ô5ðVð VØ@JðVð Và$ðVð Và(,Ô(9ðVð Vð Vñô ð ð )D×(MÒ(MÈaÑ(PÔ(P×(ZÒ(ZÐ[]Ñ(^Ô(^Ð%ð $(Ô#?×#IÒ#IÈ!Ñ#LÔ#L×#VÒ#VÐWXÑ#YÔ#Y×#cÒ#cÐdeÑ#fÔ#fÐ õ
 %*Ô$5Ø)Ø$×'Ò'Ð(EÔ(LÑMÔMñ%
ô %
Ð!ð ×1Ò1°,À
ÑKÔKˆð ”l—o’o f¤mÑ4Ô4ˆØ˜+Ñ%ˆÝ”˜FÑ#Ô#ˆØ˜+Ñ%ˆõ ”Ð6¸ÅÄÈFÌLÑ@YÔ@YÔ@]Ñ^Ô^ˆÝœÔ+×3Ò3°FÀÍ%Ì-Ð3ÑXÔX×[Ò[ÐbnÔbtÐ[ÑuÔuˆð -:Ô,?Ñ)ˆˆu�e˜U EØÔ" 2Ô&ˆØ ×(Ò(¨¨A¨q°!°QÑ7Ô7×?Ò?ÀÀEÈ5ÑQÔQˆØ×$Ò$ Q¨¨1¨a°Ñ3Ô3×;Ò;¸BÀÀuÑMÔMˆÝ”Y˜x¨Ñ/Ô/ˆ
Ø$×,Ò,¨U°E¸5À%ÈÑOÔO×WÒWÐXYÐ[\Ð^_ÐabÐdeÑfÔfˆØ)×1Ò1àØ  4¤?Ñ2Ø”Ø”ð	ñ
ô 
ˆð *¨!¨!¨!¨W¨f¨W¨*Ô5ˆàÐr`   )r|   r}   r~   r�   râ   r]  rÒ   rö   r‚   rn  rv  rƒ  rÓ   rõ   rÍ   rÎ   s   @r^   rI  rI  ¯  s  ø€ € € € € ð 
Ð1ð  
ð  
ð  
ð  
ð  
ð  
ðD'ð 'ð 'ð˜5œ<ð °3ð À3ð È5Ì<ð ð ð ð ð¨u¬|ð ÀÄð ð ð ð ð,.'°E´Lð .'ÀUÄ\ð .'ð .'ð .'ð .'ð`d U¤\ð d¸Ô9Ið dÈeÌlð dð dð dð dð dð dð dð dr`   rI  c                   ód   ‡ — e Zd ZdZ	 d
dedee         defˆ fd„Zdej	        dej	        fd	„Z
ˆ xZS )ÚGemma3nAudioCumulativeGroupNormaè  Applies Group Normalization cumulatively over the time dimension.

    This layer normalizes the input by calculating the mean and variance
    cumulatively over the time dimension (dim 1). The statistics are computed
    over all feature dimensions (specified by `feature_dims` and `num_channels`)
    for elements marked as valid by the optional `mask`.

    If a `mask` is provided (True for valid, False for invalid/padded),
    invalid time steps do not contribute to the statistics calculation, and
    their corresponding output values are zeroed out.

    Scale and bias, if enabled, are applied per-channel (last dimension).
    This behavior is similar to JAX's `GroupNormalization` with `num_groups=1`
    and `cumulative=True`.
    r¥   Únum_channelsÚfeature_dimsrÞ   c           	      óV  •— t          ¦   «                              ¦   «          || _        t          |¦  «        | _        || _        t          j        t          j	        |¦  «        ¦  «        | _
        t          t          ddt          | j        ¦  «        z   dz   ¦  «        ¦  «        | _        d S )Nr   rW   )rÊ   râ   r®  r«   r¯  rÞ   rã   rä   rÒ   rå   ræ   rf   rd   Úreduction_axes)rh   r®  r¯  rÞ   rË   s       €r^   râ   z(Gemma3nAudioCumulativeGroupNorm.__init__¤  sŽ   ø€ õ 	‰Œ×ÒÑÔÐØ(ˆÔÝ! ,Ñ/Ô/ˆÔØˆŒõ ”l¥5¤:¨lÑ#;Ô#;Ñ<Ô<ˆŒõ $¥E¨!¨Qµ°TÔ5FÑ1GÔ1GÑ-GÈ!Ñ-KÑ$LÔ$LÑMÔMˆÔÐÐr`   rç   rð   c                 óÈ  — | j         | j        fz   }|j        dd…         |k    r"t          d|j        dd…         › d|› �¦  «        ‚|j        }t
          j        }|                     |¦  «        }t          j        ||¬¦  «        }t          j	        || j
        d¬¦  «        }t          j        |d¬	¦  «        }t          j	        || j
        d¬¦  «        }	t          j        |	d¬	¦  «        }
t          j        |
d
¬¦  «        }||z  }||z
                       d¦  «        }t          j	        || j
        d¬¦  «        }t          j        |d¬	¦  «        }||z  }||z
  t          j        || j        z   ¦  «        z  }| j                             |¦  «        }dg|                     ¦   «         dz
  z  | j        gz   }||                     |¦  «        z  }||z  }|                     |¦  «        S )zÞApplies cumulative group norm, optionally using a mask.

        Args:
          hidden_states: Input tensor, shape [B, T, *feature_dims, C].

        Returns:
          Normalized tensor with the same shape as x.
        r   NzInput tensor shape suffix z> does not match expected suffix (feature_dims + num_channels) r6  T©rÝ   rê   rW   r  rý   )rŒ  )r¯  r®  r7  re   r  rÒ   r  r  Ú	ones_likeÚsumr±  ÚcumsumÚclamprì   ÚrsqrtrÞ   ræ   rÝ   r‡  )rh   rç   Úexpected_input_suffixÚinput_dtypeÚ
calc_dtypeÚx_calcÚ	mask_calcÚsum_values_at_tÚcum_sum_valuesÚelements_in_group_at_tÚcum_count_elementsÚsafe_cum_count_elementsÚcum_meanÚsquared_diff_from_meanÚsum_sq_diff_at_tÚcum_sum_sq_diffÚcum_varianceÚnormalized_xÚscaleÚscale_view_shapeÚfinal_outputs                        r^   rõ   z'Gemma3nAudioCumulativeGroupNorm.forward¶  s   € ð !%Ô 1°TÔ5FÐ4HÑ HÐØÔ˜q˜r˜rÔ"Ð&;Ò;Ð;ÝðQ¨]Ô-@ÀÀÀÔ-Dð Qð QØ9NðQð Qñô ð ð
 $Ô)ˆå”]ˆ
Ø×!Ò! *Ñ-Ô-ˆõ ”O F°*Ð=Ñ=Ô=ˆ	õ  œ) F°Ô0CÈTÐRÑRÔRˆåœ o¸1Ð=Ñ=Ô=ˆõ "'¤¨9¸$Ô:MÐW[Ð!\Ñ!\Ô!\Ðå"œ\Ð*@ÀaÐHÑHÔHÐå"'¤+Ð.@ÀcÐ"JÑ"JÔ"JÐð "Ð$;Ñ;ˆð
 #)¨8Ñ"3×!8Ò!8¸Ñ!;Ô!;ÐÝ œ9Ð%;ÀÔATÐ^bÐcÑcÔcÐõ  œ,Ð'7¸QÐ?Ñ?Ô?ˆð 'Ð)@Ñ@ˆð  Ñ)­U¬[¸ÈÌÑ9PÑ-QÔ-QÑQˆð ”—’˜zÑ*Ô*ˆà˜3 -×"3Ò"3Ñ"5Ô"5¸Ñ"9Ñ:¸dÔ>OÐ=PÑPÐØ# e§j¢jÐ1AÑ&BÔ&BÑBˆð $ iÑ/ˆà�Š˜{Ñ+Ô+Ð+r`   )r¥   )r|   r}   r~   r   r‚   r   r†   râ   rÒ   rö   rõ   rÍ   rÎ   s   @r^   r­  r­  “  s«   ø€ € € € € ðð ð( ð	Nð NàðNð ˜s”mðNð ð	Nð Nð Nð Nð Nð Nð$G, U¤\ð G,°e´lð G,ð G,ð G,ð G,ð G,ð G,ð G,ð G,r`   r­  c                   óp   ‡ — e Zd ZdZ	 ddedededeeeeef         fˆ fd„Zdej	        d	ej	        fd
„Z
ˆ xZS )ÚGemma3nAudioSSCPConvBlockzÙA single convolution block for the SubSampleConvProjection.

    This block consists of a 2D convolution, followed by CumulativeGroupNorm,
    and a ReLU activation. It handles manual padding for the convolution.
    ©r   r   r   r   rù   ÚidxÚinput_freq_dimÚmanual_paddingc                 ó"  •— t          ¦   «                              ¦   «          || _        || _        |dk    rdn| j        j        |dz
           }| j        j        |         }| j        j        |         \  }}| j        j        |         \  }	}
t          j        ||||f|	|
fdd¬¦  «        | _	        || j        d         z   | j        d         z   }||z
  |
z  dz   }t          ||f| j        j        ¬¦  «        | _        t          j        ¦   «         | _        d S )Nr   rW   )r   r   F)Úin_channelsÚout_channelsÚkernel_sizeÚstrideÚpaddingrü   )r®  r¯  rÞ   )rÊ   râ   rù   rÑ  r¤   r¨   rª   rã   ÚConv2dÚconvr­  r¦   ÚnormÚReLUÚ
activation)rh   rù   rÏ  rÐ  rÑ  rÓ  rÔ  Úkernel_hÚkernel_wÚstride_hÚstride_wÚf_in_paddedÚ
f_out_convrË   s                €r^   râ   z"Gemma3nAudioSSCPConvBlock.__init__  s+  ø€ õ 	‰Œ×ÒÑÔÐØˆŒØ,ˆÔð  !š8˜8�a�a¨¬Ô)KÈCÐRSÉGÔ)TˆØ”{Ô9¸#Ô>ˆØ!œ[Ô>¸sÔCÑˆ�(Ø!œ[Ô>¸sÔCÑˆ�(å”IØ#Ø%àØðð ˜hÐ'ØØð

ñ 

ô 

ˆŒ	ð % tÔ':¸1Ô'=Ñ=ÀÔ@SÐTUÔ@VÑVˆØ! HÑ,°Ñ9¸AÑ=ˆ
å3Ø%Ø$˜Ø”Ô4ð
ñ 
ô 
ˆŒ	õ œ'™)œ)ˆŒˆˆr`   Úaudio_encodingsrð   c                 ó¦  — t          j        || j        dd¬¦  «                             | j        j        j        ¦  «        }|                      |¦  «        }|                     dddd¦  «                             ¦   «         }|  	                    |¦  «        }|                     dddd¦  «                             ¦   «         }|  
                    |¦  «        S )NÚconstantrb   )ÚmodeÚvaluer   r   r	   rW   )ÚFr)  rÑ  r  rÙ  ræ   r  r9  rp  rÚ  rÜ  )rh   rã  Úaudio_encodings_paddedÚaudio_encodings_convÚ
x_for_normÚx_normedÚaudio_encodings_normeds          r^   rõ   z!Gemma3nAudioSSCPConvBlock.forward2  sÃ   € õ "#¤ ¸Ô8KÐR\ÐdgÐ!hÑ!hÔ!h×!kÒ!kØŒIÔÔ"ñ"
ô "
Ðð
  $ŸyšyÐ)?Ñ@Ô@Ðð *×1Ò1°!°Q¸¸1Ñ=Ô=×HÒHÑJÔJˆ
Ø—9’9˜ZÑ(Ô(ˆà!)×!1Ò!1°!°Q¸¸1Ñ!=Ô!=×!HÒ!HÑ!JÔ!JÐØ�ŠÐ5Ñ6Ô6Ð6r`   )rÎ  )r|   r}   r~   r   r�   r‚   r«   râ   rÒ   rö   rõ   rÍ   rÎ   s   @r^   rÍ  rÍ     s­   ø€ € € € € ðð ð 5Að)$ð )$à"ð)$ð ð)$ð ð	)$ð
 ˜c 3¨¨SÐ0Ô1ð)$ð )$ð )$ð )$ð )$ð )$ðV7 u¤|ð 7¸¼ð 7ð 7ð 7ð 7ð 7ð 7ð 7ð 7r`   rÍ  c                   óH   ‡ — e Zd Zdefˆ fd„Zdej        dej        fd„Zˆ xZS )Ú#Gemma3nAudioSubSampleConvProjectionrù   c                 ó’  •— t          ¦   «                              ¦   «          || _        |j        }g }g }t	          d¦  «        D ]r}|j        |         \  }}|j        |         \  }}	d}
|dz
  }d}d}|||
|f}|                     |¦  «         ||z   |z   }||z
  |	z  dz   }|                     |¦  «         |}Œst          d|j        ||d         ¬¦  «        | _	        t          d|d         ||d         ¬¦  «        | _
        |j        d         }|d         }||z  | _        t          j        | j        | j        j        d¬¦  «        | _        d S )Nr   r   rW   )rÏ  rÐ  rù   rÑ  ré   Frû   )rÊ   râ   rù   r’   rf   r¨   rª   ÚappendrÍ  Úconv_0Úconv_1r¤   Úinput_proj_in_featuresrã   r  r?   Úinput_proj_linear)rh   rù   Úcurrent_f_for_block_inputÚcalculated_block_paddingÚcalculated_f_out_dimsr]   rÝ  rÞ  rß  rà  Ú	pad_t_topÚpad_t_bottomÚ
pad_f_leftÚpad_f_rightÚmanual_padding_tuplerá  Úf_out_after_convÚfinal_c_outÚfinal_f_outrË   s                      €r^   râ   z,Gemma3nAudioSubSampleConvProjection.__init__F  s¡  ø€ Ý‰Œ×ÒÑÔÐØˆŒà$*Ô$:Ð!Ø#%Ð Ø "Ðå�q‘”ð 	9ð 	9ˆAØ!'Ô!=¸aÔ!@ÑˆH�hØ!'Ô!=¸aÔ!@ÑˆH�hð ˆIØ# a™<ˆLð ˆJØˆKð ØØØð	$Ð ð %×+Ò+Ð,@ÑAÔAÐAð 4°jÑ@À;ÑNˆKØ +¨hÑ 6¸8ÑCÀaÑGÐØ!×(Ò(Ð)9Ñ:Ô:Ð:Ø(8Ð%Ð%å/ØØ!Ô1ØØ3°AÔ6ð	
ñ 
ô 
ˆŒõ 0ØØ0°Ô3ØØ3°AÔ6ð	
ñ 
ô 
ˆŒð Ô3°BÔ7ˆØ+¨BÔ/ˆØ&1°KÑ&?ˆÔ#Ý!#¤¨4Ô+FÈÌÔH_ÐfkÐ!lÑ!lÔ!lˆÔÐÐr`   rã  rð   c                 óN  — |                      d¦  «        }|                      |¦  «        }|                      |¦  «        }|j        \  }}}}|                     dddd¦  «                             ¦   «         }|                     ||||z  ¦  «        }	|                      |	¦  «        }
|
S )NrW   r   r   r	   )r  rò  ró  r7  r9  rp  r‡  rõ  )rh   rã  Úaudio_encodings_reshapedre  rq  Úc_outÚt_outÚf_outÚ
x_permutedÚoutput_flattenedÚoutputs              r^   rõ   z+Gemma3nAudioSubSampleConvProjection.forward  sŸ   € ð $3×#<Ò#<¸QÑ#?Ô#?Ð Ø�KŠKÐ0Ñ1Ô1ˆØ�KŠK˜‰NŒNˆà!"¤Ñˆˆ5�%˜à—Y’Y˜q ! Q¨Ñ*Ô*×5Ò5Ñ7Ô7ˆ
Ø%Ÿ?š?¨1¨e°U¸U±]ÑCÔCÐØ×'Ò'Ð(8Ñ9Ô9ˆØˆr`   ©	r|   r}   r~   r�   râ   rÒ   rö   rõ   rÍ   rÎ   s   @r^   rï  rï  E  ss   ø€ € € € € ð7mÐ1ð 7mð 7mð 7mð 7mð 7mð 7mðr u¤|ð ¸¼ð ð ð ð ð ð ð ð r`   rï  c                   óV   ‡ — e Zd Zdefˆ fd„Zdej        dej        dej        fd„Zˆ xZ	S )ÚGemma3nAudioConformerAttentionrù   c                 óÖ  •— t          ¦   «                              ¦   «          || _        | j        j        | _        |                      dt          j        | j        j        ¦  «        d¬¦  «         t          | j        j        ¦  «        | _
        t          |¦  «        | _        t          j        | j        | j        j        d¬¦  «        | _        t          | j        j        ¦  «        | _        d S )Nr•   Frÿ   rû   )rÊ   râ   rù   r?   Úpost_in_featuresr  rÒ   rZ  r•   rÜ   Úpre_attn_normrI  Úattnrã   r  ÚpostÚ	post_norm©rh   rù   rË   s     €r^   râ   z'Gemma3nAudioConformerAttention.__init__�  s¶   ø€ Ý‰Œ×ÒÑÔÐØˆŒØ $¤Ô 7ˆÔØ×ÒÐ0µ%´,¸t¼{Ô?\Ñ2]Ô2]ÐjoÐÑpÔpÐpÝ+¨D¬KÔ,CÑDÔDˆÔÝ)¨&Ñ1Ô1ˆŒ	Ý”I˜dÔ3°T´[Ô5LÐSXÐYÑYÔYˆŒ	Ý'¨¬Ô(?Ñ@Ô@ˆŒˆˆr`   rã  rÑ   rð   c                 ó†  — |}t          j        || j         | j        ¦  «        }|                      |¦  «        }|                      ||¦  «        }|j        \  }}}}	|                     ||||	z  ¦  «        }
|                      |
¦  «        }t          j        || j         | j        ¦  «        }||                      |¦  «        z   S rò   )	rÒ   r·  r•   r  r  r7  r*  r  r  )rh   rã  rÑ   Úaudio_encodings_input_to_attnÚaudio_encodings_normÚaudio_encodings_attn_outrq  rr  r  r  r  s              r^   rõ   z&Gemma3nAudioConformerAttention.forward™  sÅ   € Ø(7Ð%Ýœ+ o¸Ô8NÐ7NÐPTÔPfÑgÔgˆØ#×1Ò1°/ÑBÔBÐà#'§9¢9Ð-AÀ>Ñ#RÔ#RÐ ð %=Ô$BÑ!ˆˆ1ˆi˜Ø#;×#CÒ#CÀAÀqÈ)ÐV^ÑJ^Ñ#_Ô#_Ð àŸ)š)Ð$<Ñ=Ô=ˆÝœ+ o¸Ô8NÐ7NÐPTÔPfÑgÔgˆØ,¨t¯~ª~¸oÑ/NÔ/NÑNÐNr`   ©
r|   r}   r~   r�   râ   rÒ   rö   rÓ   rõ   rÍ   rÎ   s   @r^   r  r  Ž  s‰   ø€ € € € € ðAÐ1ð Að Að Að Að Að AðO u¤|ð OÀUÔEUð OÐZ_ÔZfð Oð Oð Oð Oð Oð Oð Oð Or`   r  c                   óH   ‡ — e Zd Zdefˆ fd„Zdej        dej        fd„Zˆ xZS )Ú Gemma3nAudioConformerFeedForwardrù   c                 ó$  •— t          ¦   «                              ¦   «          || _        |                      dt	          j        | j        j        ¦  «        d¬¦  «         t          | j        j        ¦  «        | _	        t          j        | j        j        | j        j        dz  d¬¦  «        | _        t          j        | j        j        dz  | j        j        d¬¦  «        | _        t          | j        j        ¦  «        | _        | j        j        | _        d S )Nr•   Frÿ   rM   rû   )rÊ   râ   rù   r  rÒ   rZ  r•   rÜ   r?   Úpre_layer_normrã   r  Úffw_layer_1Úffw_layer_2Úpost_layer_normr¢   Úpost_layer_scaler  s     €r^   râ   z)Gemma3nAudioConformerFeedForward.__init__«  sÜ   ø€ Ý‰Œ×ÒÑÔÐØˆŒà×ÒÐ0µ%´,¸t¼{Ô?\Ñ2]Ô2]ÐjoÐÑpÔpÐpå,¨T¬[Ô-DÑEÔEˆÔÝœ9 T¤[Ô%<¸d¼kÔ>UÐXYÑ>YÐ`eÐfÑfÔfˆÔÝœ9 T¤[Ô%<¸qÑ%@À$Ä+ÔBYÐ`eÐfÑfÔfˆÔÝ-¨d¬kÔ.EÑFÔFˆÔØ $¤Ô @ˆÔÐÐr`   rã  rð   c                 óŠ  — |}t          j        || j         | j        ¦  «        }|                      |¦  «        }|                      |¦  «        }t
          j                             |¦  «        }|                      |¦  «        }t          j        || j         | j        ¦  «        }|  	                    |¦  «        }||| j
        z  z   S rò   )rÒ   r·  r•   r  r  rã   r(  Úsilur  r  r  )rh   rã  Úresiduals      r^   rõ   z(Gemma3nAudioConformerFeedForward.forward·  sµ   € Ø"ˆÝœ+ o¸Ô8NÐ7NÐPTÔPfÑgÔgˆØ×-Ò-¨oÑ>Ô>ˆØ(,×(8Ò(8¸Ñ(IÔ(IˆÝœ-×,Ò,¨_Ñ=Ô=ˆØ(,×(8Ò(8¸Ñ(IÔ(IˆÝœ+ o¸Ô8NÐ7NÐPTÔPfÑgÔgˆØ×.Ò.¨Ñ?Ô?ˆØ˜?¨TÔ-BÑBÑCÐCr`   r	  rÎ   s   @r^   r  r  ª  s|   ø€ € € € € ð
AÐ1ð 
Að 
Að 
Að 
Að 
Að 
Að	D u¤|ð 	D¸¼ð 	Dð 	Dð 	Dð 	Dð 	Dð 	Dð 	Dð 	Dr`   r  c                   óH   ‡ — e Zd Zdefˆ fd„Zdej        dej        fd„Zˆ xZS )Ú Gemma3nAudioConformerLightConv1drù   c           	      óä  •— t          ¦   «                              ¦   «          || _        t          | j        j        | j        j        ¬¦  «        | _        t          j        | j        j        | j        j        dz  d¬¦  «        | _	        t          j
        | j        j        | j        j        | j        j        dd| j        j        d¬¦  «        | _        |                      dt          j        | j        j        ¦  «        d¬	¦  «         t          | j        j        | j        j        ¬¦  «        | _        t          j        | j        j        | j        j        d¬¦  «        | _        | j        j        dz
  | _        d S )
N©rÞ   r   Frû   rW   r   )rÓ  rÔ  rÕ  rÖ  r×  Úgroupsrü   r•   rÿ   )rÊ   râ   rù   rÜ   r?   r”   r  rã   r  Úlinear_startÚConv1drŸ   Údepthwise_conv1dr  rÒ   rZ  r•   Ú	conv_normÚ
linear_endÚcausal_paddingr  s     €r^   râ   z)Gemma3nAudioConformerLightConv1d.__init__Ä  s2  ø€ Ý‰Œ×ÒÑÔÐØˆŒå,¨T¬[Ô-DÈ$Ì+ÔJbÐcÑcÔcˆÔÝœI d¤kÔ&=¸t¼{Ô?VÐYZÑ?ZÐafÐgÑgÔgˆÔÝ "¤	ØœÔ/ØœÔ0ØœÔ9ØØØ”;Ô*Øð!
ñ !
ô !
ˆÔð 	×ÒÐ0µ%´,¸t¼{Ô?\Ñ2]Ô2]ÐjoÐÑpÔpÐpÝ'¨¬Ô(?ÀTÄ[ÔE]Ð^Ñ^Ô^ˆŒÝœ) D¤KÔ$;¸T¼[Ô=TÐ[`ÐaÑaÔaˆŒà"œkÔ?À!ÑCˆÔÐÐr`   rã  rð   c                 óF  — |}|                       |¦  «        }|                      |¦  «        }t          j        j                             |d¬¦  «        }|                     ddd¦  «        }t          j        || j	        df¦  «        }|  
                    |¦  «        }|                     ddd¦  «        }t          j        || j         | j        ¦  «        }|                      |¦  «        }t          j                             |¦  «        }|                      |¦  «        }||z   }|S )Nré   r  r   r   rW   )r  r(  rÒ   rã   r(  Úglur9  rè  r)  r-  r*  r·  r•   r+  r!  r,  )rh   rã  Úaudio_encodings_residualÚaudio_encodings_permutedÚaudio_encodings_permuted_paddedr  s         r^   rõ   z(Gemma3nAudioConformerLightConv1d.forwardÙ  s  € Ø#2Ð à×-Ò-¨oÑ>Ô>ˆØ×+Ò+¨OÑ<Ô<ˆÝœ(Ô-×1Ò1°/ÀrÐ1ÑJÔJˆà#2×#:Ò#:¸1¸aÀÑ#CÔ#CÐ å*+¬%Ð0HÈ4ÔK^Ð`aÐJbÑ*cÔ*cÐ'Ø×/Ò/Ð0OÑPÔPˆà)×1Ò1°!°Q¸Ñ:Ô:ˆÝœ+ o¸Ô8NÐ7NÐPTÔPfÑgÔgˆØŸ.š.¨Ñ9Ô9ˆÝœ-×,Ò,¨_Ñ=Ô=ˆØŸ/š/¨/Ñ:Ô:ˆØ Ð#;Ñ;ˆØˆr`   r	  rÎ   s   @r^   r$  r$  Ã  sr   ø€ € € € € ðDÐ1ð Dð Dð Dð Dð Dð Dð* u¤|ð ¸¼ð ð ð ð ð ð ð ð r`   r$  c                   óV   ‡ — e Zd Zdefˆ fd„Zdej        dej        dej        fd„Zˆ xZ	S )ÚGemma3nAudioConformerBlockrù   c                 óÂ  •— t          ¦   «                              ¦   «          || _        t          | j        ¦  «        | _        t          | j        ¦  «        | _        t          | j        ¦  «        | _        t          | j        ¦  «        | _	        |  
                    dt          j        | j        j        ¦  «        d¬¦  «         t          | j        j        ¦  «        | _        d S )Nr•   Frÿ   )rÊ   râ   rù   r  Úffw_layer_startr  Ú	attentionr$  Úlconv1dÚffw_layer_endr  rÒ   rZ  r•   rÜ   r?   rÚ  r  s     €r^   râ   z#Gemma3nAudioConformerBlock.__init__ï  sª   ø€ Ý‰Œ×ÒÑÔÐØˆŒå?ÀÄÑLÔLˆÔÝ7¸¼ÑDÔDˆŒÝ7¸¼ÑDÔDˆŒÝ=¸d¼kÑJÔJˆÔØ×ÒÐ0µ%´,¸t¼{Ô?\Ñ2]Ô2]ÐjoÐÑpÔpÐpÝ" 4¤;Ô#:Ñ;Ô;ˆŒ	ˆ	ˆ	r`   rã  rÑ   rð   c                 ó‚  — |                       |¦  «        }|                      ||¦  «        }| }||                     d¦  «                             |j        ¦  «        z  }|                      |¦  «        }|                      |¦  «        }t          j        || j	         | j	        ¦  «        }|  
                    |¦  «        }|S )Nré   )r6  r7  r  r  r  r8  r9  rÒ   r·  r•   rÚ  )rh   rã  rÑ   Úvalidity_mask_for_lconvÚaudio_encodings_for_lconv_inputr  s         r^   rõ   z"Gemma3nAudioConformerBlock.forwardú  s½   € Ø×.Ò.¨Ñ?Ô?ˆØŸ.š.¨¸.ÑIÔIˆØ#1 /ÐØ*9Ð<S×<]Ò<]Ð^`Ñ<aÔ<a×<dÒ<dØÔ!ñ=
ô =
ñ +
Ð'ð Ÿ,š,Ð'FÑGÔGˆà×,Ò,¨_Ñ=Ô=ˆÝœ+ o¸Ô8NÐ7NÐPTÔPfÑgÔgˆØ—’˜?Ñ+Ô+ˆØˆr`   r  rÎ   s   @r^   r4  r4  î  sw   ø€ € € € € ð	<Ð1ð 	<ð 	<ð 	<ð 	<ð 	<ð 	<ð u¤|ð ÀUÔEUð ÐZ_ÔZfð ð ð ð ð ð ð ð r`   r4  c                   ó   — e Zd ZdS )ÚGemma3nTextScaledWordEmbeddingN©r|   r}   r~   r[   r`   r^   r>  r>    ó   € € € € € Ø€Dr`   r>  c                   óL   ‡ — e Zd ZdZdefˆ fd„Zdej        dej        fd„Zˆ xZ	S )ÚGemma3nTextLaurelBlockz Learned Augmented Residual Layerrù   c                 ój  •— t          ¦   «                              ¦   «          || _        t          j        | j        j        | j        j        d¬¦  «        | _        t          j        | j        j        | j        j        d¬¦  «        | _        t          | j        j        | j        j
        ¬¦  «        | _        d S )NFrû   r&  )rÊ   râ   rù   rã   r  r?   rR   Úlinear_leftÚlinear_rightrÜ   r”   Úpost_laurel_normr  s     €r^   râ   zGemma3nTextLaurelBlock.__init__  s�   ø€ Ý‰Œ×ÒÑÔÐØˆŒåœ9 T¤[Ô%<¸d¼kÔ>UÐ\aÐbÑbÔbˆÔÝœI d¤kÔ&=¸t¼{Ô?VÐ]bÐcÑcÔcˆÔÝ .¨t¬{Ô/FÈDÌKÔLdÐ eÑ eÔ eˆÔÐÐr`   rç   rð   c                 óŠ   — |                       |¦  «        }|                      |¦  «        }|                      |¦  «        }||z   S rò   )rD  rE  rF  )rh   rç   Úlaurel_hidden_statesÚnormed_laurel_hidden_statess       r^   rõ   zGemma3nTextLaurelBlock.forward  sL   € Ø-1×-=Ò-=¸mÑ-LÔ-LÐØ-1×->Ò->Ð?SÑ-TÔ-TÐØ&*×&;Ò&;Ð<PÑ&QÔ&QÐ#ØÐ:Ñ:Ð:r`   )
r|   r}   r~   r   r3   râ   rÒ   rö   rõ   rÍ   rÎ   s   @r^   rB  rB    sx   ø€ € € € € Ø*Ð*ðfÐ0ð fð fð fð fð fð fð; U¤\ð ;°e´lð ;ð ;ð ;ð ;ð ;ð ;ð ;ð ;r`   rB  c                   ór   ‡ — e Zd Zd
dedefˆ fd„Zdej        dej        fd„Zdej        dej        fd	„Z	ˆ xZ
S )ÚGemma3nTextMLPr   rù   Ú	layer_idxc                 ó’   •— t          ¦   «                              |¦  «         |j        |         | _        |j        |         | _        d S rò   )rÊ   râ   rB   rS   Úactivation_sparsity©rh   rù   rL  rË   s      €r^   râ   zGemma3nTextMLP.__init__#  sA   ø€ Ý‰Œ×Ò˜Ñ Ô Ð Ø!'Ô!9¸)Ô!DˆÔØ#)Ô#EÀiÔ#PˆÔ Ð Ð r`   rç   rð   c                 óô   — |                       |¦  «        }| j        dk    r|                      |¦  «        }|                      |¦  «        }|                      |¦  «        }|                      ||z  ¦  «        }|S )Nrb   )Ú	gate_projrN  Ú_gaussian_topkÚact_fnÚup_projÚ	down_proj)rh   rç   rQ  ÚactivationsrT  rU  s         r^   rõ   zGemma3nTextMLP.forward(  sr   € Ø—N’N =Ñ1Ô1ˆ	ØÔ# cÒ)Ð)Ø×+Ò+¨IÑ6Ô6ˆIØ—k’k )Ñ,Ô,ˆØ—,’,˜}Ñ-Ô-ˆØ—N’N ;°Ñ#8Ñ9Ô9ˆ	ØÐr`   Úinputsc                 ó²  — t          j        | j        t           j        |j        ¬¦  «        }t           j        j                             dd¦  «        }|                     |¦  «        }| 	                    |j
        ¦  «        }t          j        |dd¬¦  «        }t          j        |ddd¬¦  «        }|||z  z   }t          j                             ||z
  ¦  «        S )	N©r  r  r   rW   ré   Tr³  F)rÝ   rê   Úunbiased)rÒ   rZ  rN  r  r  ÚdistributionsÚnormalÚNormalÚicdfr  r  rí   Ústdrã   r(  Úrelu)rh   rW  Útarget_sparsity_tensorÚnormal_distÚstd_multiplierÚinputs_meanÚ
inputs_stdÚcutoff_xs           r^   rR  zGemma3nTextMLP._gaussian_topk1  sÀ   € Ý!&¤¨dÔ.FÍeÌmÐdjÔdqÐ!rÑ!rÔ!rÐõ Ô)Ô0×7Ò7¸¸1Ñ=Ô=ˆØ'2×'7Ò'7Ð8NÑ'OÔ'OˆØ'×,Ò,¨V¬\Ñ:Ô:ˆÝ”j ¨R¸Ð>Ñ>Ô>ˆÝ”Y˜v¨2°tÀeÐLÑLÔLˆ
Ø ¨nÑ!<Ñ<ˆÝŒ}×!Ò! &¨8Ñ"3Ñ4Ô4Ð4r`   )r   )r|   r}   r~   r3   r‚   râ   rÒ   rö   rõ   rR  rÍ   rÎ   s   @r^   rK  rK  "  s¦   ø€ € € € € ðQð QÐ0ð Q¸Sð Qð Qð Qð Qð Qð Qð
 U¤\ð °e´lð ð ð ð ð5 U¤\ð 5°e´lð 5ð 5ð 5ð 5ð 5ð 5ð 5ð 5r`   rK  c                   óê   ‡ — e Zd ZdZdefˆ fd„Zdej        dej        fd„Zdej        dej        fd„Z	d	ej        d
ej        dej        fd„Z
dej        dej        fd„Zdej        dej        fd„Zˆ xZS )ÚGemma3nTextAltUpa�  Alternating Updates (AltUp)

    The AltUp module wraps transformer layers. The `predict` step modifies the
    input to the transformer layer, and the `correct` step propagates the output
    of the transformer layer to the sparsely updated dimensions.

    See more in the research paper:

    https://proceedings.neurips.cc/paper_files/paper/2023/file/f2059277ac6ce66e7e5543001afa8bb5-Paper-Conference.pdf
    rù   c                 ó¨  •— t          ¦   «                              ¦   «          || _        t          j        t          j        | j        j        ¦  «        ¦  «        | _        t          j	        | j        j
        | j        j
        d¬¦  «        | _        t          j	        | j        j
        | j        j
        dz  d¬¦  «        | _        t          j	        | j        j        | j        j
        d¬¦  «        | _        t          | j        j        | j        j        ¬¦  «        | _        |                      dt          j        | j        j        dz  ¦  «        d¬¦  «         d S )NFrû   r   r&  Úrouter_input_scaleç      ð¿rÿ   )rÊ   râ   rù   rã   rä   rÒ   rT  r?   Úcorrect_output_scaler  rN   Úcorrection_coefsÚprediction_coefsÚmodality_routerrÜ   r”   Úrouter_normr  rZ  r  s     €r^   râ   zGemma3nTextAltUp.__init__N  s
  ø€ Ý‰Œ×ÒÑÔÐØˆŒÝ$&¤Lµ´¸T¼[Ô=TÑ1UÔ1UÑ$VÔ$VˆÔ!Ý "¤	¨$¬+Ô*FÈÌÔHdÐkpÐ qÑ qÔ qˆÔÝ "¤	¨$¬+Ô*FÈÌÔHdÐfgÑHgÐnsÐ tÑ tÔ tˆÔÝ!œy¨¬Ô)@À$Ä+ÔB^ÐejÐkÑkÔkˆÔÝ)¨$¬+Ô*AÀtÄ{ÔG_Ð`Ñ`Ô`ˆÔØ×ÒÐ1µ5´<ÀÄÔ@WÐY]Ñ@]Ñ3^Ô3^ÐkpÐÑqÔqÐqÐqÐqr`   re  rð   c                 óØ   — |                       |¦  «        | j        z  }|                      |¦  «        }t          j        |                     ¦   «         ¦  «                             |¦  «        S rò   )rp  rj  ro  rÒ   r‰  r†   ró   )rh   re  Úrouter_inputsÚrouteds       r^   Úcompute_router_modalitiesz*Gemma3nTextAltUp.compute_router_modalitiesX  sV   € Ø×(Ò(¨Ñ+Ô+¨dÔ.EÑEˆØ×%Ò% mÑ4Ô4ˆÝŒz˜&Ÿ,š,™.œ.Ñ)Ô)×1Ò1°!Ñ4Ô4Ð4r`   rç   c                 óz  — |                       || j        j                 ¦  «        }| j        rF| j        j        �:| j        j        j                             | j        j         | j        j        ¦  «          |                      |¦  «        j	        g |j
        dd…         ¢| j        j        ‘| j        j        ‘R Ž                      dddd¦  «        }t          j        |                     dddd¦  «        |¦  «        }|                     dddd¦  «        }||z  }|                     ¦   «                              |¦  «        S )aµ  Predicts the output of a layer using a trainable map.

        Args:
            hidden_states: A 4D tensor of shape `[num_altup_inputs, batch_size, num_tokens, hidden_size]` derived by
                stacking the input embeddings and preprocessing the last `num_altup_inputs - 1` matrices.

        Returns:
            A 4D tensor of shape `[num_altup_inputs, batch_size, num_tokens, hidden_size]` containing the predictions.
        Nré   r   rW   r	   r   )rt  rù   rJ   ÚtrainingrK   rn  ræ   ÚdataÚclamp_r*  r7  rN   r9  rÒ   r:  rp  ró   )rh   rç   Ú
modalitiesÚ	all_coefsÚpredictionss        r^   ÚpredictzGemma3nTextAltUp.predict]  sA  € ð ×3Ò3°MÀ$Ä+ÔB^Ô4_Ñ`Ô`ˆ
àŒ=ð 	p˜Tœ[Ô8ÐDØÔ!Ô(Ô-×4Ò4°d´kÔ6QÐ5QÐSWÔS^ÔSnÑoÔoÐoðˆD×!Ò! *Ñ-Ô-ÜðiØ Ô& s¨ sÔ+ðiØ-1¬[Ô-IðiØKOÌ;ÔKgðið ið içŠW�Q˜˜1˜aÑ Ô ð 	õ ”l =×#8Ò#8¸¸A¸qÀ!Ñ#DÔ#DÀiÑPÔPˆØ!×)Ò)¨!¨Q°°1Ñ5Ô5ˆØ�}Ñ$ˆØ×%Ò%Ñ'Ô'×/Ò/°Ñ>Ô>Ð>r`   r{  Ú	activatedc                 ó†  — |                       |¦  «        }||| j        j                 z
  }|                     | j        j        ddd¦  «        }| j        rl| j        j        �`| j        j         	                    | j        j         | j        j        ¦  «        }t          j        j                             ||d¬¦  «        dz   }n|                      |¦  «        dz   }|                     ddd¦  «                             d¦  «        }t          j        ||¦  «        }||z  }|                     ¦   «                              |¦  «        S )a_  Corrects the predictions relative to the

        Args:
            predictions: A 4D tensor of shape `[num_altup_inputs, batch_size, num_tokens, hidden_size]` derived by
                stacking the input embeddings and preprocessing the last `num_altup_inputs - 1` matrices.
            activated: A 3D tensor of shape `[batch_size, num_tokens, hidden_size]` containing the activated inputs.

        Returns:
            A 4D tensor of shape `[num_altup_inputs, batch_size, num_tokens, hidden_size]` correcting the original
                predictions relative to the activated input embeddings.
        rW   Nrû   rý   r   r   ré   )rt  rù   rJ   ÚrepeatrN   rv  rK   rm  ræ   r·  rÒ   rã   r(  Úlinearr9  r  Úmulrp  ró   )rh   r{  r}  ry  Ú
innovationræ   rz  Ú	correcteds           r^   ÚcorrectzGemma3nTextAltUp.correcty  s)  € ð ×3Ò3°IÑ>Ô>ˆ
Ø ¨T¬[Ô-IÔ!JÑJˆ
Ø×&Ò& t¤{Ô'CÀQÈÈ1ÑMÔMˆ
àŒ=ð 	@˜Tœ[Ô8ÐDØÔ*Ô1×7Ò7¸¼Ô9TÐ8TÐVZÔVaÔVqÑrÔrˆFÝœÔ+×2Ò2°:¸vÈDÐ2ÑQÔQÐTWÑWˆIˆIà×-Ò-¨jÑ9Ô9¸CÑ?ˆIð
 ×%Ò% a¨¨AÑ.Ô.×8Ò8¸Ñ<Ô<ˆ	å”I˜j¨)Ñ4Ô4ˆ	Ø�[Ñ ˆ	Ø×#Ò#Ñ%Ô%×-Ò-¨iÑ8Ô8Ð8r`   rƒ  c                 ól   — |                      | j        ¦  «        | j        z                        |¦  «        S )a	  
        This is only defined as the `forward` so that accelerate hooks can move correctly `correct_output_scale`
        (which is a nn.Parameter, not a Module) between devices when offloading. It is otherwise only used in
        `scale_corrected_output`
        )ró   rl  ©rh   rƒ  s     r^   rõ   zGemma3nTextAltUp.forward˜  s2   € ð ×!Ò! $Ô";Ñ<Ô<¸tÔ?XÑX×aÒaÐbkÑlÔlÐlr`   c                 ó,   — |                       |¦  «        S )zMScales the provided 3D tensor of shape [batch_size, num_tokens, hidden_size].)rõ   r†  s     r^   Úscale_corrected_outputz'Gemma3nTextAltUp.scale_corrected_output   s   € à�|Š|˜IÑ&Ô&Ð&r`   )r|   r}   r~   r   r3   râ   rÒ   rö   rt  r|  r„  rõ   rˆ  rÍ   rÎ   s   @r^   rh  rh  B  s#  ø€ € € € € ð	ð 	ðrÐ0ð rð rð rð rð rð rð5¨5¬<ð 5¸E¼Lð 5ð 5ð 5ð 5ð
? U¤\ð ?°e´lð ?ð ?ð ?ð ?ð89 5¤<ð 9¸E¼Lð 9ÈUÌ\ð 9ð 9ð 9ð 9ð>m ¤ð m°%´,ð mð mð mð mð'°´ð 'ÀÄð 'ð 'ð 'ð 'ð 'ð 'ð 'ð 'r`   rh  rW   re  r  r  Úunsqueeze_dimc                 ó†   — |                      |¦  «        }|                      |¦  «        }| |z  t          | ¦  «        |z  z   S )a\  Applies Rotary Position Embedding to the query and key tensors.

    Args:
        x (`torch.Tensor`): The tensor to embed.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        unsqueeze_dim (`int`, *optional*, defaults to 1):
            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    )r  r$   )re  r  r  r‰  s       r^   Úapply_rotary_pos_embr‹  ¥  s@   € ð" �-Š-˜Ñ
&Ô
&€CØ
�-Š-˜Ñ
&Ô
&€CØ�‰G� A™œ¨Ñ,Ñ-Ð-r`   c                   óì   ‡ — e Zd Zdedefˆ fd„Z	 	 ddej        dej        dej        dz  dedz  d	e	ee
ej        ej        f         f         dz  d
ee         de
ej        ej        dz  f         fd„Zˆ xZS )ÚGemma3nTextAttentionrù   rL  c                 ó¼  •— t          ¦   «                              ¦   «          || _        || _        t	          |d¦  «        r|j        |         nd | _        | j        dk    | _        | j        r|j        nd | _        t          |d|j
        |j        z  ¦  «        | _        |j        |j        z  | _        d| _        | j        j        | _        d| _        | j        j        | j        j        z
  }||cxk    odk    nc | _        |j        d |…         }| j        rIt+          |¦  «        dz
  |d d d…                              |j        |         ¦  «        z
  | _        d	| _        nLd | _        |t+          |¦  «        dz
  |d d d…                              |j        |         ¦  «        z
  k    | _        t3          j        |j
        |j        | j        z  |j        ¬
¦  «        | _        t;          |j        |j        ¬¦  «        | _        | j        s§t3          j        |j
        |j        | j        z  |j        ¬
¦  «        | _         t3          j        |j
        |j        | j        z  |j        ¬
¦  «        | _!        t;          |j        |j        ¬¦  «        | _"        t;          |j        |j        d	¬¦  «        | _#        t3          j        |j        | j        z  |j
        |j        ¬
¦  «        | _$        d S )NrH   rZ   r  rý   Tr   rW   ré   Frû   )rÝ   rÞ   )rÝ   rÞ   rß   )%rÊ   râ   rù   rL  ÚhasattrrH   Ú
layer_typeÚ
is_slidingrG   Úgetattrr?   Únum_attention_headsr  rE   Únum_key_value_groupsÚscalingÚattention_dropoutÚ	is_causalrD   rP   Úis_kv_shared_layerrd   ÚindexÚkv_shared_layer_indexÚstore_full_length_kvrã   r  Úattention_biasrV  rÜ   r”   Úq_normrW  rX  Úk_normÚv_normÚo_proj)rh   rù   rL  Úfirst_kv_shared_layer_idxÚprev_layersrË   s        €r^   râ   zGemma3nTextAttention.__init__¼  sÝ  ø€ Ý‰Œ×ÒÑÔÐØˆŒØ"ˆŒÝ;BÀ6È=Ñ;YÔ;YÐc˜&Ô,¨YÔ7Ð7Ð_cˆŒØœ/Ð-@Ò@ˆŒØ7;´ÐP˜fÔ3Ð3ÈDˆÔå ¨
°FÔ4FÈ&ÔJdÑ4dÑeÔeˆŒØ$*Ô$>À&ÔB\Ñ$\ˆÔ!ØˆŒØ!%¤Ô!>ˆÔØˆŒà$(¤KÔ$AÀDÄKÔDdÑ$dÐ!Ø"+Ð/HÐ"LÐ"LÒ"LÐ"LÈ1Ò"LÐ"LÐ"LÐ"LˆÔØÔ(Ð)CÐ*CÐ)CÔDˆØÔ"ð 		å),¨[Ñ)9Ô)9¸AÑ)=ÀÈDÈDÈbÈDÔ@Q×@WÒ@WÐX^ÔXjÐktÔXuÑ@vÔ@vÑ)vˆDÔ&Ø(-ˆDÔ%Ð%à)-ˆDÔ&à(1µS¸Ñ5EÔ5EÈÑ5IÈKÐX\ÐX\ÐZ\ÐX\ÔL]×LcÒLcØÔ" 9Ô-ñMô Mñ 6ò )ˆDÔ%õ ”iØÔ Ô :¸T¼]Ñ JÐQWÔQfð
ñ 
ô 
ˆŒõ %¨¬¸fÔ>QÐRÑRÔRˆŒð Ô&ð 	iÝœ)ØÔ" FÔ$>ÀÄÑ$NÐU[ÔUjðñ ô ˆDŒKõ œ)ØÔ" FÔ$>ÀÄÑ$NÐU[ÔUjðñ ô ˆDŒKõ )¨V¬_À&ÔBUÐVÑVÔVˆDŒKÝ(¨V¬_À&ÔBUÐbgÐhÑhÔhˆDŒKå”iØÔ&¨¬Ñ6¸Ô8JÐQWÔQfð
ñ 
ô 
ˆŒˆˆr`   Nrç   Úposition_embeddingsÚattention_maskÚpast_key_valuesÚshared_kv_statesri   rð   c                 ó°  — |j         d d…         }g |¢d‘| j        j        ‘R }|\  }	}
|                      |¦  «                             |¦  «        }|                      |¦  «        }t          ||	|
d¬¦  «        }|                     dd¦  «        }| j        rE|| j	                 \  }}| 
                    |j        ¦  «        }| 
                    |j        ¦  «        }n¹|                      |¦  «                             |¦  «        }|                      |¦  «        }t          ||	|
d¬¦  «        }|                     dd¦  «        }|                      |¦  «                             |¦  «        }|                      |¦  «        }|                     dd¦  «        }|�&| j        s|                     ||| j        ¦  «        \  }}| j        r||f|| j        <   t'          j        | j        j        t,          ¦  «        } || ||||f| j        r| j        nd| j        | j        dœ|¤Ž\  }} |j        g |¢d‘R Ž                      ¦   «         }|                      |¦  «        }||fS )Nré   r   )r‰  rW   rb   )Údropoutr•  rG   )r7  rù   r  rV  r‡  r�  r‹  Ú	transposer˜  rš  r  r  rW  rž  rX  rŸ  ru   rL  r›  r   Úget_interfaceÚ_attn_implementationr#   rv  r–  r•  rG   r*  rp  r   )rh   rç   r£  r¤  r¥  r¦  ri   Úinput_shapeÚhidden_shaper  r  r�  r‘  r’  Úattention_interfaceÚattn_outputÚattn_weightss                    r^   rõ   zGemma3nTextAttention.forwardì  sŒ  € ð $Ô)¨#¨2¨#Ô.ˆØ?˜Ð? bÐ?¨$¬+Ô*>Ð?Ð?ˆà&‰ˆˆSØ—{’{ =Ñ1Ô1×6Ò6°|ÑDÔDˆØ—{’{ <Ñ0Ô0ˆÝ+¨L¸#¸sÐRSÐTÑTÔTˆØ#×-Ò-¨a°Ñ3Ô3ˆð
 Ô"ð 	8Ø'7¸Ô8RÔ'SÑ$ˆJ˜à#Ÿš |Ô':Ñ;Ô;ˆJØ'Ÿ?š?¨<Ô+>Ñ?Ô?ˆLˆLàŸš ]Ñ3Ô3×8Ò8¸ÑFÔFˆJØŸš ZÑ0Ô0ˆJÝ-¨j¸#¸sÐRSÐTÑTÔTˆJØ#×-Ò-¨a°Ñ3Ô3ˆJàŸ;š; }Ñ5Ô5×:Ò:¸<ÑHÔHˆLØŸ;š; |Ñ4Ô4ˆLØ'×1Ò1°!°QÑ7Ô7ˆLàÐ&¨tÔ/FÐ&Ø'6×'=Ò'=¸jÈ,ÐX\ÔXfÑ'gÔ'gÑ$ˆJ˜ØÔ$ð 	HØ/9¸<Ð/GÐ˜Tœ^Ñ,å(?Ô(MØŒKÔ,Õ.Eñ)
ô )
Ðð %8Ð$7ØØØØØð
%
ð /3¬mÐD�DÔ*Ð*ÀØ”LØÔ.ð
%
ð 
%
ð ð
%
ð 
%
Ñ!ˆ�\ð *�kÔ)Ð;¨;Ð;¸Ð;Ð;Ð;×FÒFÑHÔHˆØ—k’k +Ñ.Ô.ˆØ˜LÐ(Ð(r`   ©NN)r|   r}   r~   r3   r‚   râ   rÒ   rö   r   rµ   r«   r   r   rõ   rÍ   rÎ   s   @r^   r�  r�  »  s÷   ø€ € € € € ð.
Ð0ð .
¸Sð .
ð .
ð .
ð .
ð .
ð .
ðj )-ØPTð;)ð ;)à”|ð;)ð #œ\ð;)ð œ tÑ+ð	;)ð
  ™ð;)ð ˜s E¨%¬,¸¼Ð*DÔ$EÐEÔFÈÑMð;)ð Ð+Ô,ð;)ð 
ˆuŒ|˜Uœ\¨DÑ0Ð0Ô	1ð;)ð ;)ð ;)ð ;)ð ;)ð ;)ð ;)ð ;)r`   r�  c                   ó0  ‡ — e Zd Zdedefˆ fd„Z	 	 	 	 	 	 ddej        dej        dej        deee	ej        ej        f         f         dz  d	ej        dz  d
ej
        dz  dedz  dee         de	ej        e	ej        ej        f         dz  f         fd„Zˆ xZS )ÚGemma3nTextDecoderLayerrù   rL  c                 ó  •— t          ¦   «                              ||¦  «         t          ||¬¦  «        | _        |j        | _        t
          |j                 | _        t          |¦  «        | _	        t          |¦  «        | _        t          ||¦  «        | _        t          j        | j        | j        d¬¦  «        | _        t          j        | j        | j        d¬¦  «        | _        t'          | j        |j        ¬¦  «        | _        d S )N)rL  Frû   r&  )rÊ   râ   rK  ÚmlprA   r   Úhidden_activationrS  rh  ÚaltuprB  Úlaurelr�  Ú	self_attnrã   r  r?   Úper_layer_input_gateÚper_layer_projectionrÜ   r”   Úpost_per_layer_input_normrO  s      €r^   râ   z Gemma3nTextDecoderLayer.__init__+  sÜ   ø€ Ý‰Œ×Ò˜ Ñ+Ô+Ð+Ý! &°IÐ>Ñ>Ô>ˆŒà+1Ô+MˆÔ(Ý˜VÔ5Ô6ˆŒå% fÑ-Ô-ˆŒ
Ý,¨VÑ4Ô4ˆŒÝ-¨f°iÑ@Ô@ˆŒÝ$&¤I¨dÔ.>ÀÔ@`ÐglÐ$mÑ$mÔ$mˆÔ!Ý$&¤I¨dÔ.NÐPTÔP`ÐglÐ$mÑ$mÔ$mˆÔ!Ý)7¸Ô8HÈfÔNaÐ)bÑ)bÔ)bˆÔ&Ð&Ð&r`   Nrç   r£  Úper_layer_inputr¦  r¤  Úposition_idsr¥  ri   rð   c           
      ó�  — | j                              |¦  «        }	|	| j        j                 }
|                      |
¦  «        }|                      |¦  «        } | j        d||||||dœ|¤Ž\  }}|                      |¦  «        }|
|z   }||z   t          j	        d¦  «        z  }|  
                    |¦  «        }|                      |¦  «        }|                      |¦  «        }||z   }| j                              |	|¦  «        }|| j        j                                      ¦   «         }| j        j        r| j                              |¦  «        }|                      |¦  «        }|                      |¦  «        }t'          j        ||¦  «        }|                      |¦  «        }|                      |¦  «        }|dd …xx         |z  cc<   |S )N)rç   r¤  r¦  r¾  r£  r¥  r   rW   r[   )r·  r|  rù   rJ   Úinput_layernormr¸  r¹  Úpost_attention_layernormr	  ÚsqrtÚpre_feedforward_layernormrµ  Úpost_feedforward_layernormr„  r[  rL   rˆ  rº  rS  rÒ   Úmultiplyr»  r¼  )rh   rç   r£  r½  r¦  r¤  r¾  r¥  ri   r{  Úactive_predictionÚactive_prediction_normedÚlaurel_outputr  r;  Ú
attn_gatedÚattn_laurelÚ	attn_normÚattn_ffwÚattn_ffw_normÚattn_ffw_laurel_gatedÚcorrected_predictionsÚfirst_predictions                          r^   rõ   zGemma3nTextDecoderLayer.forward9  sï  € ð ”j×(Ò(¨Ñ7Ô7ˆØ'¨¬Ô(DÔEÐà#'×#7Ò#7Ð8IÑ#JÔ#JÐ ØŸšÐ$<Ñ=Ô=ˆà �$”.ð 
Ø2Ø)Ø-Ø%Ø 3Ø+ð
ð 
ð ð
ð 
‰ˆˆað ×,Ò,¨TÑ2Ô2ˆà&¨Ñ-ˆ
Ø! MÑ1µT´Y¸q±\´\ÑAˆà×2Ò2°;Ñ?Ô?ˆ	Ø—8’8˜IÑ&Ô&ˆØ×7Ò7¸ÑAÔAˆØ +¨mÑ ;ÐØ $¤
× 2Ò 2°;Ð@UÑ VÔ VÐà0°´Ô1MÔN×TÒTÑVÔVÐØŒ;Ô*ð 	SØ#œz×@Ò@ÐAQÑRÔRÐð  ×4Ò4Ð5EÑFÔFÐØŸ;š;Ð'7Ñ8Ô8ÐÝ œ>Ð*:¸OÑLÔLÐð  ×4Ò4Ð5EÑFÔFÐØ×9Ò9Ð:JÑKÔKÐØ˜a˜b˜bÐ!Ð!Ô!Ð%5Ñ5Ð!Ð!Ñ!à$Ð$r`   )NNNNNN)r|   r}   r~   r3   r‚   râ   rÒ   rö   rµ   r«   Ú
LongTensorr   r   r   rØ   rõ   rÍ   rÎ   s   @r^   r³  r³  *  s6  ø€ € € € € ðcÐ0ð c¸Sð cð cð cð cð cð cð" -1Ø(,ØPTØ.2Ø04Ø(,ð3%ð 3%à”|ð3%ð #œ\ð3%ð œð	3%ð
 ˜s E¨%¬,¸¼Ð*DÔ$EÐEÔFÈÑMð3%ð œ tÑ+ð3%ð Ô&¨Ñ-ð3%ð  ™ð3%ð Ð+Ô,ð3%ð 
ˆuŒ|˜U 5Ô#4°eÔ6GÐ#GÔHÈ4ÑOÐOÔ	Pð3%ð 3%ð 3%ð 3%ð 3%ð 3%ð 3%ð 3%r`   r³  c            	       óÚ   ‡ — e Zd ZU eed<   dZddgZdgZee	dœZ
 ej        ¦   «         d„ ¦   «         Zd„ Zd	„ Z	 	 	 dded
z  ded
z  dedej        fˆ fd„Z	 	 	 dded
z  ded
z  defd„Zˆ xZS )ÚGemma3nPreTrainedModelrù   )ÚimageÚtextÚaudior¥  r¦  r³  )rç   Ú
attentionsc                 óÎ  — t          j        | |¦  «         t          |t          ¦  «        rt	          j        |j        ¦  «         �nit          |t          ¦  «        rÆt	          j        |j	        ¦  «         |j
        dz  }dt          j        j                             t          j        d¦  «        ¦  «        z  }t	          j        |j        ||z  ¦  «         t	          j        |j        |j        ¦  «         t	          j        |j        |                     ¦   «         ¦  «         �nŽt          |t.          ¦  «        r!t	          j        |j        |j        ¦  «         �nXt          |t4          ¦  «        rBt	          j        |j        ¦  «         t	          j        |j        | j        j        dz  ¦  «         �nt          |t>          ¦  «        rÕd\  }}|j         dz  }tC          j"        tG          |¦  «        tG          |¦  «        z  ¦  «        tI          |dz
  d¦  «        z  }|t          j%        t          j&        |¦  «        | z  ¦  «        z  }t	          j        |j'        | #                    ¦   «          (                    d¦  «         (                    d¦  «        ¦  «         �nt          |tR          ¦  «        rRt	          j        |j*        | j        dz  ¦  «         t	          j        |j+        dtC          j,        d	¦  «        z  ¦  «         n°t          |tZ          ¦  «        r›|j.        D ]“}	|j/        }
|j0        |	         d
k    rtb          |j0        |	                  }
 |
|j        |	¬¦  «        \  }}t	          j        te          ||	› d�¦  «        |¦  «         t	          j        te          ||	› d�¦  «        |¦  «         Œ”tg          |d¦  «        r&t	          j        |j4        | j        j4        ¦  «         d S d S )Nrë   rý   rb   rk  )rý   r8   r   rW   r   ç       @rp   )r�  Ú	_inv_freqÚ_original_inv_freqr•   )5r   Ú_init_weightsrc   r­  ÚinitÚones_ræ   rI  Úzeros_rU  r  rÒ   rã   r(  rY  rZ  Úcopy_rK  Ú	constant_rM  rQ  rL  r]  r>  Úembed_scaleÚscalar_embed_scalerh  rl  rj  rù   r?   rø   r  r	  r
  r†   r  r  r  rþ   r  ÚGemma3nTextModelÚper_layer_projection_scaleÚper_layer_input_scalerÂ  ÚGemma3nRotaryEmbeddingrH   Úcompute_default_rope_parametersro   r   r’  r�  r•   )rh   ÚmodulerK  r^  r  r  r  r  rþ   r�  Úrope_init_fnÚcurr_inv_freqr;  s                r^   rÜ  z$Gemma3nPreTrainedModel._init_weightsy  s�  € åÔ% d¨FÑ3Ô3Ð3Ý�fÕ=Ñ>Ô>ð  	^ÝŒJ�v”}Ñ%Ô%Ð%Ñ%Ý˜Õ 5Ñ6Ô6ð 	^ÝŒK˜Ô,Ñ-Ô-Ð-Ø”o tÑ+ˆGØ¥¤Ô!4×!=Ò!=½e¼lÈ3Ñ>OÔ>OÑ!PÔ!PÑPˆLÝŒJ�v”~ w°Ñ'=Ñ>Ô>Ð>ÝŒN˜6œ>¨6Ô+KÑLÔLÐLÝŒJ�vÔ5°v×7\Ò7\Ñ7^Ô7^Ñ_Ô_Ð_Ñ_Ý˜Õ >Ñ?Ô?ð 	^ÝŒN˜6Ô-¨vÔ/HÑIÔIÐIÑIÝ˜Õ 0Ñ1Ô1ð 	^ÝŒK˜Ô3Ñ4Ô4Ð4ÝŒN˜6Ô4°d´kÔ6MÈtÑ6SÑTÔTÐTÑTÝ˜Õ EÑFÔFð 	^Ø+5Ñ(ˆM˜=Ø#œ_°Ñ1ˆNÝ&*¤h­u°]Ñ/CÔ/CÅeÈMÑFZÔFZÑ/ZÑ&[Ô&[Õ^aØ Ñ" Añ_ô _ñ 'Ð#ð +­U¬Yµu´|ÀNÑ7SÔ7SÐWnÐVnÑ7nÑ-oÔ-oÑoˆNÝŒJ�vÔ,¨n×.BÒ.BÑ.DÔ.D×.NÒ.NÈqÑ.QÔ.Q×.[Ò.[Ð\]Ñ.^Ô.^Ñ_Ô_Ð_Ñ_Ý˜Õ 0Ñ1Ô1ð 
	^ÝŒN˜6Ô<¸dÔ>NÐPTÑ>TÑUÔUÐUÝŒN˜6Ô7¸½T¼YÀs¹^¼^Ñ9KÑLÔLÐLÐLÝ˜Õ 6Ñ7Ô7ð 	^Ø$Ô0ð ^ð ^�
Ø%ÔE�ØÔ# JÔ/°9Ò<Ð<Ý#6°vÔ7GÈ
Ô7SÔ#T�LØ#/ <°´È*Ð#UÑ#UÔ#UÑ �˜qÝ”
�7 6¨jÐ+CÐ+CÐ+CÑDÔDÀmÑTÔTÐTÝ”
�7 6¨jÐ+LÐ+LÐ+LÑMÔMÈ}Ñ]Ô]Ð]Ð]å�6Ð.Ñ/Ô/ð 	TÝŒN˜6Ô3°T´[Ô5RÑSÔSÐSÐSÐSð	Tð 	Tr`   c                 ó   — | j         j        S rò   ©Ú
base_modelÚembed_tokens_per_layer©rh   s    r^   Úget_per_layer_input_embeddingsz5Gemma3nPreTrainedModel.get_per_layer_input_embeddings¡  s   € ØŒÔ5Ð5r`   c                 ó   — || j         _        d S rò   rí  ©rh   rç  s     r^   Úset_per_layer_input_embeddingsz5Gemma3nPreTrainedModel.set_per_layer_input_embeddings¤  s   € Ø16ˆŒÔ.Ð.Ð.r`   NTÚnew_num_tokensÚpad_to_multiple_ofÚmean_resizingrð   c                 ó~   •— t          ¦   «                              |||¬¦  «        }|                      |||¦  «         |S )N)rõ  rö  r÷  )rÊ   Úresize_token_embeddingsÚ_resize_per_layer_embeddings)rh   rõ  rö  r÷  Úinputs_embedsrË   s        €r^   rù  z.Gemma3nPreTrainedModel.resize_token_embeddings§  sM   ø€ õ ™œ×7Ò7Ø)Ø1Ø'ð 8ñ 
ô 
ˆð
 	×)Ò)¨.Ð:LÈmÑ\Ô\Ð\ØÐr`   c                 óš  — | j         | j                             ¦   «         _        | j                             ¦   «         j        r‰|                      ¦   «         }|                      ||||¦  «        }t          |d¦  «        r|j        }t          ||¦  «         | 
                    |j        j        ¦  «         |                      |¦  «         d S d S )NÚ_hf_hook)r;   rù   Úget_text_configr=   rA   rñ  Ú_get_resized_embeddingsr�  rý  r0   Úrequires_grad_ræ   rá   rô  )rh   rõ  rö  r÷  rï  Únew_embeddings_per_layerÚhooks          r^   rú  z3Gemma3nPreTrainedModel._resize_per_layer_embeddingsµ  sÛ   € ð DHÄ?ˆŒ×#Ò#Ñ%Ô%Ô@ØŒ;×&Ò&Ñ(Ô(ÔDð 		JØ%)×%HÒ%HÑ%JÔ%JÐ"Ø'+×'CÒ'CØ&¨Ð8JÈMñ(ô (Ð$õ Ð-¨zÑ:Ô:ð CØ-Ô6�Ý"Ð#;¸TÑBÔBÐBØ$×3Ò3Ð4JÔ4QÔ4_Ñ`Ô`Ð`Ø×/Ò/Ð0HÑIÔIÐIÐIÐIð		Jð 		Jr`   )NNT)r|   r}   r~   r·   rƒ   Úinput_modalitiesÚ_skip_keys_device_placementÚ_no_split_modulesr³  r�  Ú_can_record_outputsrÒ   Úno_gradrÜ  rñ  rô  r‚   r‡   rã   Ú	Embeddingrù  rú  rÍ   rÎ   s   @r^   rÓ  rÓ  o  sR  ø€ € € € € € ØÐÐÑØ1ÐØ#4Ð6HÐ"IÐØ2Ð3Ðà0Ø*ðð Ðð
 €U„]�_„_ð%Tð %Tñ „_ð%TðN6ð 6ð 6ð7ð 7ð 7ð
 &*Ø)-Ø"ð	ð à˜d™
ðð   $™Jðð ð	ð
 
Œðð ð ð ð ð ð  &*Ø)-Ø"ð	Jð Jà˜d™
ðJð   $™JðJð ð	Jð Jð Jð Jð Jð Jð Jð Jr`   rÓ  c                   óš   ‡ — e Zd ZU dZeed<   dZdZdefˆ fd„Ze	e
dej        dej        dee         deez  fd	„¦   «         ¦   «         Zˆ xZS )
ÚGemma3nAudioEncoderzx
    An audio encoder based on the [Universal Speech Model](https://huggingface.co/papers/2303.01037) architecture.
    rù   Ú	audio_melrÖ  c                 ó  •‡— t          ¦   «                              ‰¦  «         ‰| _        t          ‰¦  «        | _        t          j        ˆfd„t          ‰j        ¦  «        D ¦   «         ¦  «        | _	        |  
                    ¦   «          d S )Nc                 ó.   •— g | ]}t          ‰¦  «        ‘ŒS r[   )r4  )r\   r;  rù   s     €r^   r_   z0Gemma3nAudioEncoder.__init__.<locals>.<listcomp>Ø  s"   ø€ Ð^Ð^Ð^°AÕ'¨Ñ/Ô/Ð^Ð^Ð^r`   )rÊ   râ   rù   rï  Úsubsample_conv_projectionrã   Ú
ModuleListrf   rž   Ú	conformerÚ	post_initr  s    `€r^   râ   zGemma3nAudioEncoder.__init__Ò  s€   øø€ Ý‰Œ×Ò˜Ñ Ô Ð ØˆŒå)LÈVÑ)TÔ)TˆÔ&ÝœØ^Ð^Ð^Ð^½¸vÔ?\Ñ9]Ô9]Ð^Ñ^Ô^ñ
ô 
ˆŒð 	�ŠÑÔÐÐÐr`   rÑ   ri   rð   c                 óì  — |                       |¦  «        }|j        d         }d}t          t          | j        j        ¦  «        ¦  «        D ]}|| j        j        |         d         z  }Œt          j        ||j        ¬¦  «        |z  }t          j	        ||j        d         dz
  ¬¦  «        }|j
        dk    r@|j
        dk    r5|                     d¦  «                             |j        d         d¦  «        }nX|j
        |j
        k    rH|j        d         dk    r7|j        d         dk    r&||j        d         k    r|                     d¦  «        }t          j        |d|¦  «        }	| j        D ]}
 |
||	¦  «        }Œ| j        j        dk    r2|dd…dd| j        j        …f         }|	dd…dd| j        j        …f         }	|                     |	                     d¦  «        d¦  «        }t#          ||	¬¦  «        S )	a¬  Encodes a batch of MELs.

        Args:
            audio_mel: a torch.Tensor of shape [batch, num_frames, num_channels,
              mel_bins].

        Returns:
            audio_encodings: a torch.Tensor of shape
                `[batch_size, self.config.audio_soft_tokens_per_image,
                self.config.audio_config.hidden_size]`
            audio_mel_mask: a torch.BoolTensor of shape [batch, num_frames].
        rW   r   r5  )r  ré   Nrb   )Úlast_hidden_staterÑ   )r  r7  rf   rd   rù   rª   rÒ   r  r  r·  r~  r  ÚexpandÚgatherr  r    Úmasked_fillrÐ   )rh   r  rÑ   ri   rã  Út_subÚtime_stride_productÚstride_pair_idxÚindicesÚcurrent_maskÚblocks              r^   rõ   zGemma3nAudioEncoder.forwardÜ  s!  € ð" ×8Ò8¸ÑCÔCˆð  Ô% aÔ(ˆàÐÝ$¥S¨¬Ô)JÑ%KÔ%KÑLÔLð 	Yð 	YˆOØ 4¤;Ô#DÀ_Ô#UÐVWÔ#XÑXÐÐõ
 ”,˜u¨^Ô-BÐCÑCÔCÐFYÑYˆÝ”+˜g¨>Ô+?ÀÔ+BÀQÑ+FÐGÑGÔGˆð Ô Ò"Ð" w¤|°qÒ'8Ð'8Ø×'Ò'¨Ñ*Ô*×1Ò1°.Ô2FÀqÔ2IÈ2ÑNÔNˆGˆGàÔ 7¤<Ò/Ð/ØÔ$ QÔ'¨1Ò,Ð,Ø”˜aÔ  AÒ%Ð%Ø˜œ qÔ)Ò)Ð)ð ×'Ò'¨Ñ*Ô*ˆGå”| N°A°wÑ?Ô?ˆà”^ð 	Cð 	CˆEØ#˜e O°\ÑBÔBˆOˆOàŒ;Ô,¨qÒ0Ð0Ø-¨a¨a¨aÐ1UÐ1U°D´KÔ4UÐ1UÐ.UÔVˆOà'¨¨¨Ð+OÐ+O¨d¬kÔ.OÐ+OÐ(OÔPˆLà)×5Ò5°l×6LÒ6LÈRÑ6PÔ6PÐRUÑVÔVˆÝ-Ø-Ø'ð
ñ 
ô 
ð 	
r`   )r|   r}   r~   r   r�   rƒ   Úmain_input_namer  râ   r   r   rÒ   rö   rÓ   r   r   r«   rÐ   rõ   rÍ   rÎ   s   @r^   r
  r
  È  sÆ   ø€ € € € € € ðð ð ÐÐÑà!€OØÐðÐ1ð ð ð ð ð ð ð  Øð8
Øœð8
Ø7<Ô7Gð8
ØSYÐZlÔSmð8
à	Ð/Ñ	/ð8
ð 8
ð 8
ñ „_ñ  Ôð8
ð 8
ð 8
ð 8
ð 8
r`   r
  c                   ó   — e Zd ZdS )rç  Nr?  r[   r`   r^   rç  rç    r@  r`   rç  zBThe base Gemma 3n language model without a language modeling head.©Úcustom_introc                   óx  ‡ — e Zd ZU eed<   defˆ fd„Zdej        dej        fd„Z		 ddej        dej        dz  dej        fd	„Z
e ed
¬¦  «        e	 	 	 	 	 	 	 ddej        dz  dej        dz  dej        dz  dej        dz  dedz  dej        dz  dedz  dee         defd„¦   «         ¦   «         ¦   «         Zˆ xZS )rä  rù   c                 ó®  •‡ ‡‡— t          ¦   «                              ‰¦  «         ‰j        ‰ _        ‰j        ‰ _        t	          ‰j        ‰j        ‰j        z  ‰ j        ‰j        dz  ¬¦  «        ‰ _        t          j
        ‰ j        ‰j        ‰j        z  d¬¦  «        ‰ _        t          ‰j        ‰j        ¬¦  «        ‰ _        t          j        ˆfd„t!          ‰j        ¦  «        D ¦   «         ¦  «        ‰ _        t          ‰j        ‰j        ¬¦  «        ‰ _        t          j        ˆ fd„t!          d‰ j        j        ¦  «        D ¦   «         ¦  «        ‰ _        t          j        ˆ fd	„t!          d‰ j        j        ¦  «        D ¦   «         ¦  «        ‰ _        ‰                      d
t1          j        ‰ j        dz  ¦  «        d¬¦  «         ‰                      dt1          j        t1          j        d¦  «        ¦  «        d¬¦  «         g ‰ _        t9          ‰ j        ¦  «        D ]7\  Š}|j        j        r&‰ j                             ˆfd„dD ¦   «         ¦  «         Œ8d S )Nr¡   )râ  Frû   r&  c                 ó0   •— g | ]}t          ‰|¦  «        ‘ŒS r[   )r³  )r\   rL  rù   s     €r^   r_   z-Gemma3nTextModel.__init__.<locals>.<listcomp>6  s$   ø€ ÐiÐiÐi¸IÕ$ V¨YÑ7Ô7ÐiÐiÐir`   c                 óR   •— g | ]#}t          j        ‰j        ‰j        d ¬¦  «        ‘Œ$S ©Frû   ©rã   r  r?   ©r\   r;  rh   s     €r^   r_   z-Gemma3nTextModel.__init__.<locals>.<listcomp><  ó0   ø€ ÐwÐwÐwÈ1�RŒY�tÔ'¨Ô)9ÀÐFÑFÔFÐwÐwÐwr`   rW   c                 óR   •— g | ]#}t          j        ‰j        ‰j        d ¬¦  «        ‘Œ$S r%  r&  r'  s     €r^   r_   z-Gemma3nTextModel.__init__.<locals>.<listcomp>@  r(  r`   rå  rë   rÿ   ræ  rÙ  c                 ó    •— g | ]
}d ‰› d|› �‘ŒS )zlayers.z.self_attn.r[   )r\   Únamer]   s     €r^   r_   z-Gemma3nTextModel.__init__.<locals>.<listcomp>K  s*   ø€ ÐiÐiÐi¸Ð3˜qÐ3Ð3¨TÐ3Ð3ÐiÐiÐir`   )rW  rX  rž  rŸ  ) rÊ   râ   r?   rA   r>  r=   rD   Úpadding_idxrï  rã   r  Úper_layer_model_projectionrÜ   r”   Úper_layer_projection_normr  rf   ÚlayersrÚ  rù   rN   Úaltup_projectionsÚaltup_unembed_projectionsr  rÒ   rZ  r¸  Ú"_keys_to_ignore_on_load_unexpectedÚ	enumerater¹  r˜  Úextend)rh   rù   Úlayerr]   rË   s   `` @€r^   râ   zGemma3nTextModel.__init__!  sd  øøøø€ Ý‰Œ×Ò˜Ñ Ô Ð à!Ô-ˆÔØ+1Ô+MˆÔ(å&DØÔ-ØÔ$ vÔ'IÑIØÔØÔ:¸CÑ?ð	'
ñ '
ô '
ˆÔ#õ +-¬)ØÔØÔ$ vÔ'IÑIØð+
ñ +
ô +
ˆÔ'õ *8¸Ô8ZÐ`fÔ`sÐ)tÑ)tÔ)tˆÔ&Ý”mØiÐiÐiÐiÍÈvÔOgÑIhÔIhÐiÑiÔiñ
ô 
ˆŒõ # 6Ô#5¸6Ô;NÐOÑOÔOˆŒ	å!#¤ØwÐwÐwÐwÕPUÐVWÐY]ÔYdÔYuÑPvÔPvÐwÑwÔwñ"
ô "
ˆÔõ *,¬ØwÐwÐwÐwÕPUÐVWÐY]ÔYdÔYuÑPvÔPvÐwÑwÔwñ*
ô *
ˆÔ&ð 	×ÒÐ9½5¼<ÈÔHXÐZ^ÑH^Ñ;_Ô;_ÐlqÐÑrÔrÐrØ×ÒÐ4µe´kÅ%Ä,ÈsÑBSÔBSÑ6TÔ6TÐafÐÑgÔgÐgð 35ˆÔ/Ý! $¤+Ñ.Ô.ð 	ð 	‰HˆAˆuØŒÔ1ð ØÔ7×>Ò>ØiÐiÐiÐiÐ@hÐiÑiÔiñô ð øð	ð 	r`   Ú	input_idsrð   c                 ór   —  |                       |¦  «        j        g |j        ¢| j        j        ‘| j        ‘R Ž S rò   )rï  r*  r7  rù   rD   rA   )rh   r6  s     r^   Úget_per_layer_inputsz%Gemma3nTextModel.get_per_layer_inputsN  sP   € Ø=ˆt×*Ò*¨9Ñ5Ô5Ô=ð 
ØŒ_ð
àŒKÔ)ð
ð Ô,ð
ð 
ð 
ð 	
r`   Nrû  Úper_layer_inputsc                 ó´  — |                       |¦  «        }|| j                             |j        |j        ¬¦  «        z  } |j        g |j        d d…         ¢| j        j        ‘| j	        ‘R Ž }|  
                    |¦  «        }|€|S |j        |j        k    r|dd | j        j        …d d …f         }||z   | j                             |j        |j        ¬¦  «        z  S )NrY  ré   .)r-  rå  r  r  r  r*  r7  rù   rD   rA   r.  ræ  )rh   rû  r9  r»  s       r^   Úproject_per_layer_inputsz)Gemma3nTextModel.project_per_layer_inputsU  s(  € ð
 .2×-LÒ-LÈ]Ñ-[Ô-[ÐØ Ô ?× BÒ BØÔ%Ð.BÔ.Ið !Cñ !
ô !
ñ 	
Ðð  <Ð3Ô;ð  
ØÔ   " Ô%ð 
àŒKÔ)ð 
ð Ô,ð 
ð  
ð  
Ðð
  $×=Ò=Ð>RÑSÔSÐàÐ#Ø'Ð'àÔ%Ð)9Ô)?Ò?Ð?à/°Ð5T°t´{Ô7TÐ5TÐVWÐVWÐVWÐ0WÔXÐà$Ð'7Ñ7¸4Ô;U×;XÒ;XØÔ%Ð.BÔ.Ið <Yñ <
ô <
ñ 
ð 	
r`   F)Útie_last_hidden_statesr¤  r¾  r¥  rÆ   ri   c           	      óf  — |du |duz  rt          d¦  «        ‚|�*|                      |¦  «        }|                      |¦  «        }|                      ||¦  «        }|r|€t	          | j        ¬¦  «        }|€V|�|                     ¦   «         nd}	t          j        |j	        d         |j
        ¬¦  «        |	z   }|                     d¦  «        }t          |x}
t          ¦  «        s&| j        ||||dœ}t          di |¤Žt          di |¤Ždœ}
|}t          j        |d	z  d
d¬¦  «        dz  }t          j        d¦  «        }|g}t%          d| j        j        ¦  «        D ]²} | j        |dz
           |¦  «        }|                     |j        |j
        ¬¦  «        }t          j        |d	z  d
d¬¦  «        }t          j        t          j        ||                     |j
        ¦  «        ¦  «        ¦  «        }||z  |z  }|                     |¦  «         Œ³t          j        |d¬¦  «        }i }t7          | j        j        ¦  «        D ]}|                      |||¦  «        ||<   Œt=          ¦   «         }t?          | j         d| j        j!        …         ¦  «        D ]U\  }}|
| j        j        |                  }|dd…dd…|dd…f         } |||| j        j        |                  |f||||dœ|¤Ž}ŒVt          j        |d         d	z  d
d¬¦  «        dz  }|d         g}t%          d| j        j        ¦  «        D ]¸} | j"        |dz
           ||         ¦  «        }|                     |j        |j
        ¬¦  «        }t          j        |d	z  d
d¬¦  «        }t          j        t          j        ||                     |j
        ¦  «        ¦  «        ¦  «        }||z  |z  }|                     |¦  «         Œ¹t          j        |¦  «        }t          j        |d¬¦  «        }|  #                    |¦  «        }tI          ||¬¦  «        S )zµ
        per_layer_inputs (torch.Tensor, *optional*, defaults to None):
            Pre-computed per-layer embeddings. If None, they are derived from input_ids if provided.
        Nú:You must specify exactly one of input_ids or inputs_embeds)rù   r   rW   r5  )rù   rû  r¤  r¥  r¾  )rY   rZ   r   ré   Tr³  r¡   gñhãˆµøä>rY  r  )r¦  r¤  r¾  r¥  )r  r¥  r[   )%re   Úembed_tokensr8  r;  r   rù   Úget_seq_lengthrÒ   r  r7  r  r  rc   rµ   r   r   rí   rZ  rf   rN   r0  r  r  rÂ  Úmaximumrñ  ÚstackÚsetrH   Ú
rotary_embr   r3  r/  rD   r1  rÚ  r   )rh   r6  r9  r¤  r¾  r¥  rû  rÆ   ri   Úpast_seen_tokensÚcausal_mask_mappingÚmask_kwargsÚhidden_states_0Útarget_magnitudeÚepsilon_tensorÚtemp_hidden_statesr]   Ú
altup_projÚcurrent_hidden_stateÚnew_magnituderç   r£  r�  r¦  Údecoder_layerÚcausal_maskr½  Úaltup_unemb_projs                               r^   rõ   zGemma3nTextModel.forwardq  sÎ  € ð$ ˜Ð -°tÐ";Ñ<ð 	[ÝÐYÑZÔZÐZàÐ Ø ×-Ò-¨iÑ8Ô8ˆMØ#×8Ò8¸ÑCÔCÐà×8Ò8¸ÐHXÑYÔYÐàð 	?˜Ð0Ý*°$´+Ð>Ñ>Ô>ˆOàÐØCRÐC^˜×=Ò=Ñ?Ô?Ð?ÐdeÐÝ œ<¨Ô(;¸AÔ(>À}ÔG[Ð\Ñ\Ô\Ð_oÑoˆLØ'×1Ò1°!Ñ4Ô4ˆLõ °Ð?Ð-ÅÑFÔFð 	ð œ+Ø!.Ø"0Ø#2Ø ,ðð ˆKõ #5Ð"CÐ"C°{Ð"CÐ"CÝ%FÐ%UÐ%UÈÐ%UÐ%Uð#ð #Ðð (ˆõ !œ: o°qÑ&8¸bÈ$ÐOÑOÔOÐSVÑVÐÝœ dÑ+Ô+ˆà-Ð.ÐÝ�q˜$œ+Ô6Ñ7Ô7ð 	<ð 	<ˆAà6˜Ô/°°A±Ô6°ÑGÔGˆJØ#-§=¢=°Ô7LÐUeÔUl =Ñ#mÔ#mÐ Ý!œJÐ';¸QÑ'>ÀBÐPTÐUÑUÔUˆMÝ!œJ¥u¤}°]ÀN×DUÒDUÐVfÔVmÑDnÔDnÑ'oÔ'oÑpÔpˆMØ#7Ð:JÑ#JÈ]Ñ#ZÐ Ø×%Ò%Ð&:Ñ;Ô;Ð;Ð;åœÐ$6¸AÐ>Ñ>Ô>ˆØ ÐÝ˜dœkÔ5Ñ6Ô6ð 	gð 	gˆJØ.2¯oªo¸mÈ\Ð[eÑ.fÔ.fÐ 
Ñ+Ð+õ
 $™:œ:Ðå )¨$¬+Ð6U¸¼Ô8UÐ6UÔ*VÑ WÔ Wð 	ð 	ÑˆAˆ}Ø-¨d¬kÔ.EÀaÔ.HÔIˆKØ.¨q¨q¨q°!°!°!°Q¸¸¸¨zÔ:ˆOà)˜MØØ# D¤KÔ$;¸AÔ$>Ô?Øð	ð "2Ø*Ø)Ø /ð	ð 	ð ð	ð 	ˆMˆMõ !œ: m°AÔ&6¸!Ñ&;ÀÈTÐRÑRÔRÐVYÑYÐØ+¨AÔ.Ð/ÐÝ�q˜$œ+Ô6Ñ7Ô7ð 	<ð 	<ˆAà-R¨TÔ-KÈAÐPQÉEÔ-RÐS`ÐabÔScÑ-dÔ-dÐØ#3×#6Ò#6¸_Ô=RÐ[kÔ[rÐ#6Ñ#sÔ#sÐ Ý!œJÐ';¸QÑ'>ÀBÐPTÐUÑUÔUˆMÝ!œJ¥u¤}°]ÀN×DUÒDUÐVfÔVmÑDnÔDnÑ'oÔ'oÑpÔpˆMØ#7Ð:JÑ#JÈ]Ñ#ZÐ Ø×%Ò%Ð&:Ñ;Ô;Ð;Ð;åœÐ$6Ñ7Ô7ˆÝœ
 =°aÐ8Ñ8Ô8ˆØŸ	š	 -Ñ0Ô0ˆå&Ø+Ø+ð
ñ 
ô 
ð 	
r`   rò   )NNNNNNN)r|   r}   r~   r3   rƒ   râ   rÒ   rÑ  rö   r8  r;  r   r   r   r   rØ   r‡   r   r   r   rõ   rÍ   rÎ   s   @r^   rä  rä    s½  ø€ € € € € € àÐÐÑð+Ð0ð +ð +ð +ð +ð +ð +ðZ
¨eÔ.>ð 
À5Ä<ð 
ð 
ð 
ð 
ð 15ð
ð 
à”|ð
ð  œ,¨Ñ-ð
ð 
Œð	
ð 
ð 
ð 
ð8  Ø€_¨EÐ2Ñ2Ô2Øð .2Ø04Ø.2Ø04Ø(,Ø26Ø!%ðm
ð m
àÔ# dÑ*ðm
ð  œ,¨Ñ-ðm
ð œ tÑ+ð	m
ð
 Ô&¨Ñ-ðm
ð  ™ðm
ð Ô(¨4Ñ/ðm
ð ˜$‘;ðm
ð Ð+Ô,ðm
ð 
!ðm
ð m
ð m
ñ „^ñ 3Ô2ñ  Ôðm
ð m
ð m
ð m
ð m
r`   rä  z?The base Gemma 3n language model with a language modeling head.c                   ó   — e Zd ZdS )ÚGemma3nForCausalLMNr?  r[   r`   r^   rS  rS  ä  s   € € € € € à€Dr`   rS  c                   óv   ‡ — e Zd ZdZdeez  defˆ fd„Z	 	 d
dej	        dz  dej
        dz  dej
        fd	„Zˆ xZS )ÚGemma3nMultimodalEmbedderzQEmbeds token ids or soft tokens for multimodal content into language model space.Úmultimodal_configr¹   c                 ó  •— t          ¦   «                              ¦   «          |j        | _        |j        | _        |j        | _        |j        | _        |j        | _        t          j
        | j        | j        ¦  «        | _        t          | j        | j        ¬¦  «        | _        t          | j        | j        ¬¦  «        | _        t          j        | j        | j        d¬¦  «        | _        t          | j        | j        d¬¦  «        | _        d S )Nr&  Frû   )rÞ   rß   )rÊ   râ   r?   Úmultimodal_hidden_sizer”   rÞ   r‘   r;   Útext_hidden_sizerã   r  Ú	embeddingrÜ   Úhard_embedding_normÚsoft_embedding_normr  Úembedding_projectionÚembedding_post_projection_norm)rh   rV  r¹   rË   s      €r^   râ   z"Gemma3nMultimodalEmbedder.__init__ì  så   ø€ õ
 	‰Œ×ÒÑÔÐà&7Ô&CˆÔ#Ø$Ô1ˆŒØ-Ô:ˆÔØ+Ô6ˆŒØ +Ô 7ˆÔåœ d¤o°tÔ7RÑSÔSˆŒÝ#1°$Ô2MÐSWÔS[Ð#\Ñ#\Ô#\ˆÔ Ý#1°$Ô2MÐSWÔS[Ð#\Ñ#\Ô#\ˆÔ Ý$&¤I¨dÔ.IÈ4ÔK`ÐglÐ$mÑ$mÔ$mˆÔ!Ý.<¸TÔ=RÐX\ÔX`ÐmrÐ.sÑ.sÔ.sˆÔ+Ð+Ð+r`   Nr6  rû  rð   c                 ó  — |du |duz  rt          d¦  «        ‚|�|                      |¦  «        }n2|                      || j        z
  ¦  «        }|                      |¦  «        }|                      |¦  «        }|                      |¦  «        S )aå  Embeds token ids or soft tokens for multimodal content into language model space.

        Args:
            input_ids: A torch.LongTensor containing the token ids to embed. Values should be in the range
                `[vocab_offset, vocab_offset + vocab_size)`.
            inputs_embeds: A torch.Tensor containing the soft tokens to embed.

        Returns:
            A torch.Tensor of embeddings with  shape `[batch_size, seq_len, self.config.text_config.hidden_size]`.
        Nr>  )re   r\  rZ  r‘   r[  r]  r^  )rh   r6  rû  Úemb_normÚhard_embÚemb_norm_projs         r^   rõ   z!Gemma3nMultimodalEmbedder.forwardÿ  s™   € ð ˜Ð -°tÐ";Ñ<ð 	[ÝÐYÑZÔZÐZàÐ$Ø×/Ò/°Ñ>Ô>ˆHˆHà—~’~ i°$Ô2CÑ&CÑDÔDˆHØ×/Ò/°Ñ9Ô9ˆHà×1Ò1°(Ñ;Ô;ˆØ×2Ò2°=ÑAÔAÐAr`   r±  )r|   r}   r~   r   r�   r­   r3   râ   rÒ   rÑ  rö   rõ   rÍ   rÎ   s   @r^   rU  rU  é  s½   ø€ € € € € Ø[Ð[ðtà-Ð0CÑCðtð 'ðtð tð tð tð tð tð* .2Ø-1ðBð BàÔ# dÑ*ðBð ”| dÑ*ðBð 
Œð	Bð Bð Bð Bð Bð Bð Bð Br`   rU  z�
    The base Gemma 3n model comprising a vision backbone, an audio backbone, and a language model without a
    language modeling head.
    c                   ót  ‡ — e Zd Zdefˆ fd„Zd„ Zd„ Ze ed¬¦  «        de	j
        dee         d	eez  fd
„¦   «         ¦   «         Z	 	 	 	 dde	j        dz  de	j
        dz  de	j
        dz  de	j
        dz  fd„Ze	 	 	 	 	 	 	 	 	 	 	 dde	j        dz  de	j
        dz  de	j
        dz  de	j        dz  de	j        dz  de	j        dz  dedz  de	j        dz  de	j
        dz  de	j        dz  dedz  dee         d	efd„¦   «         Ze ed¬¦  «        de	j        de	j        dee         d	eez  fd„¦   «         ¦   «         Zˆ xZS )ÚGemma3nModelrù   c                 ó,  •— t          ¦   «                              |¦  «         | `| `|j        j        | _        t          j        |j        ¦  «        | _	        t          |j        |j        ¦  «        | _        t          |j        |j        ¦  «        | _        d S rò   )rÊ   râ   Úmulti_modal_projectorÚtext_config_dtyper¹   r=   r    Úfrom_configr»   Úaudio_towerrU  rº   Úembed_visionÚembed_audior  s     €r^   râ   zGemma3nModel.__init__"  s�   ø€ Ý‰Œ×Ò˜Ñ Ô Ð ØÐ&ØÐ"Ø*0Ô*<Ô*WˆÔ'Ý$Ô0°Ô1DÑEÔEˆÔÝ5°fÔ6JÈFÔL^Ñ_Ô_ˆÔÝ4°VÔ5HÈ&ÔJ\Ñ]Ô]ˆÔÐÐr`   c                 ó   — | j         j        S rò   ©Úlanguage_modelrï  rð  s    r^   rñ  z+Gemma3nModel.get_per_layer_input_embeddings+  s   € ØÔ"Ô9Ð9r`   c                 ó   — || j         _        d S rò   rm  ró  s     r^   rô  z+Gemma3nModel.set_per_layer_input_embeddings.  s   € Ø5:ˆÔÔ2Ð2Ð2r`   zOProjects the last hidden state from the vision model into language model space.r  Úpixel_valuesri   rð   c                 ó:  —  | j         d	|dddœ|¤Ž}|j        }|                     |j        d         | j        j        j        | j        j        ¦  «                             ddd¦  «        }|| j        j        j        dz  z  }|  	                    |¬¦  «        |_
        |S )
NFT)rp  r±   Úreturn_dictr   r   rW   r¡   ©rû  r[   )Úvision_towerr  r*  r7  rù   rº   r?   r¾   r9  rj  Úpooler_output)rh   rp  ri   Úvision_outputsr  s        r^   Úget_image_featureszGemma3nModel.get_image_features1  s¶   € ð +˜Ô*Ðs¸ÐQVÐdhÐsÐsÐlrÐsÐsˆØ*Ô<Ðð .×5Ò5ØÔ# AÔ&ØŒKÔ%Ô1ØŒKÔ4ñ
ô 
÷ Š'�!�Q˜Ñ
Ô
ð	 	ð 	˜Tœ[Ô6ÔBÀCÑGÑGÐØ'+×'8Ò'8ÐGXÐ'8Ñ'YÔ'YˆÔ$àÐr`   Nr6  rû  Úimage_featuresÚaudio_featuresc           	      ó4  — |€Ç| |                       ¦   «         t          j        | j        j        t          j        |j        ¬¦  «        ¦  «        k    }|                     d¦  «        }| |                       ¦   «         t          j        | j        j        t          j        |j        ¬¦  «        ¦  «        k                         d¦  «        }n || j        j        k    }|| j        j        k    }| 	                    ¦   «         }| 
                    d¦  «                             |j        ¦  «        }|�St          ||j        d         z  |                     ¦   «         k    d|› d|j        d         |j        d         z  › �¦  «         | 	                    ¦   «         }| 
                    d¦  «                             |j        ¦  «        }|�St          ||j        d         z  |                     ¦   «         k    d|› d|j        d         |j        d         z  › �¦  «         ||fS )	zï
        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
        equal to the length of multimodal features. If the lengths are different, an error is raised.
        NrY  ré   z6Image features and image tokens do not match, tokens: z, features: r   rW   z6Audio features and audio tokens do not match, tokens: )Úget_input_embeddingsrÒ   rZ  rù   rÁ   Úlongr  ÚallrÄ   rµ  r  r  r   r7  Únumel)	rh   r6  rû  rx  ry  Úspecial_image_maskÚspecial_audio_maskÚn_image_tokensÚn_audio_tokenss	            r^   Úget_placeholder_maskz!Gemma3nModel.get_placeholder_maskG  sj  € ð ÐØ!.Ð2M°$×2KÒ2KÑ2MÔ2MÝ”˜Tœ[Ô7½u¼zÐR_ÔRfÐgÑgÔgñ3ô 3ò "Ðð "4×!7Ò!7¸Ñ!;Ô!;ÐàØ.�4×,Ò,Ñ.Ô.Ý”L ¤Ô!;Å5Ä:ÐVcÔVjÐkÑkÔkñô ò÷ Šc�"‰gŒgð Ðð "+¨d¬kÔ.HÒ!HÐØ!*¨d¬kÔ.HÒ!HÐà+×/Ò/Ñ1Ô1ˆØ/×9Ò9¸"Ñ=Ô=×@Ò@ÀÔAUÑVÔVÐØÐ%Ý"Ø Ô!4°RÔ!8Ñ8¸N×<PÒ<PÑ<RÔ<RÒRð YÈð  Yð  YÐesÔeyÐz{Ôe|ð  @Nô  @Tð  UVô  @Wñ  fWð  Yð  Yñô ð ð
 ,×/Ò/Ñ1Ô1ˆØ/×9Ò9¸"Ñ=Ô=×@Ò@ÀÔAUÑVÔVÐØÐ%Ý"Ø Ô!4°RÔ!8Ñ8¸N×<PÒ<PÑ<RÔ<RÒRð YÈð  Yð  YÐesÔeyÐz{Ôe|ð  @Nô  @Tð  UVô  @Wñ  fWð  Yð  Yñô ð ð
 "Ð#5Ð5Ð5r`   Úinput_featuresr¤  Úinput_features_maskr¾  r¥  Útoken_type_idsÚlabelsrÆ   Ú	lm_kwargsc                 óN  — |du |	duz  rt          d¦  «        ‚|�� |                      ¦   «         |¦  «        }	t          j        |dk    || j        k     ¦  «        }t          j        ||t          j        |¦  «        ¦  «        }| j                             |¦  «        }t          j        || j	        j
        k    || j        j
        k     ¦  «        }| j	        j
        | j	        j        z   dz
  }t          j        |||¦  «                             |	j        ¦  «        }|  	                    |¬¦  «        }|                     |	j        |	j        ¦  «        }|                     d¦  «        }t          j        |||	¦  «        }	|| j        j
        k    }| j        j
        | j        j        z   dz
  }t          j        |||¦  «                             |	j        ¦  «        }|                      |¬¦  «        }|                     |	j        |	j        ¦  «        }|                     d¦  «        }t          j        |||	¦  «        }	nd}|�m|                      |d¬¦  «        j        }|                     |	j        |	j        ¦  «        }|                      ||	|¬	¦  «        \  }}|	                     ||¦  «        }	|��3|��0|                      || d¬¦  «        }|j        }|j        }t          j        | j        dz
  ggt          j        |j        ¬
¦  «        }|                      |¬¦  «        } t          j        |                     d¦  «        | |¦  «        }|j        \  }!}"}#| j        j        |"z
  }$|                      |!|$|#¦  «        }%t          j        ||%fd¬¦  «        }|                     |	j        |	j        ¦  «        }|                      ||	|¬¦  «        \  }}&|	                     |&|¦  «        }	 | j        dd|||||	|ddœ|¤Ž}'t=          |'j        |r|'j         nd|'j!        |'j"        |�|nd|�|nd¬¦  «        S )a}  
        input_features_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
            Attention mask for `input_features` where non-zero values mark valid audio frames.
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`.

        Example:

        ```python
        >>> from PIL import Image
        >>> import httpx
        >>> from io import BytesIO
        >>> from transformers import AutoProcessor, Gemma3nForConditionalGeneration

        >>> model = Gemma3nForConditionalGeneration.from_pretrained("google/gemma3n2-3b-mix-224")
        >>> processor = AutoProcessor.from_pretrained("google/gemma3n2-3b-mix-224")

        >>> prompt = "Where is the cat standing?"
        >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
        >>> with httpx.stream("GET", url) as response:
        ...     image = Image.open(BytesIO(response.read()))

        >>> inputs = processor(images=image, text=prompt,  return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(**inputs,)
        >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Where is the cat standing?\nsnow"
        ```
        Nr>  r   rW   )r6  ré   T)rr  )rû  rx  rY  r  )rû  ry  )r6  r9  r¤  r¾  r¥  rû  rÆ   rr  )r  r¥  rç   r×  Úimage_hidden_statesrÖ   r[   )#re   r{  rÒ   rˆ  r=   rŠ  Ú
zeros_likern  r8  rj  r‘   rk  r;   r  r  r  r  rw  ru  rƒ  Úmasked_scatterÚget_audio_featuresrÑ   rZ  r|  r7  rù   r½   r  r  rÕ   r  r¥  rç   r×  )(rh   r6  rp  r„  r¤  r…  r¾  r¥  r†  rû  r‡  rÆ   rˆ  Úper_layer_inputs_maskÚper_layer_inputs_tokensr9  Úvision_maskÚdummy_vision_token_idÚvision_input_idsÚvision_embedsÚexpanded_vision_maskÚ
audio_maskÚdummy_audio_token_idÚaudio_input_idsÚaudio_embedsÚexpanded_audio_maskrx  r  r;  Úaudio_outputsry  Úaudio_padding_toksÚaudio_padding_embsÚaudio_batch_sizeÚaudio_seq_lenÚaudio_embed_dimÚextra_padding_tokensÚextra_padding_featuresr€  Úoutputss(                                           r^   rõ   zGemma3nModel.forwards  s’  € ð` ˜Ð -°tÐ";Ñ<ð 	[ÝÐYÑZÔZÐZàÑ Ø7˜D×5Ò5Ñ7Ô7¸	ÑBÔBˆMõ %*Ô$5°iÀ1²nÀiÐRVÔRqÒFqÑ$rÔ$rÐ!Ý&+¤kÐ2GÈÕTYÔTdÐenÑToÔToÑ&pÔ&pÐ#Ø#Ô2×GÒGÐH_Ñ`Ô`Ðõ  Ô+Ø˜TÔ.Ô;Ò;¸YÈÔIYÔIfÒ=fñô ˆKð %)Ô$5Ô$BÀTÔEVÔEaÑ$aÐdeÑ$eÐ!Ý$œ{¨;¸	ÐCXÑYÔY×\Ò\Ð]jÔ]qÑrÔrÐØ ×-Ò-Ð8HÐ-ÑIÔIˆMØ)×,Ò,¨]Ô-AÀ=ÔCVÑWÔWˆMØ#.×#8Ò#8¸Ñ#<Ô#<Ð Ý!œKÐ(<¸mÈ]Ñ[Ô[ˆMð # dÔ&6Ô&CÒCˆJØ#'Ô#3Ô#@À4ÔCSÔC^Ñ#^ÐabÑ#bÐ Ý#œk¨*°iÐAUÑVÔV×YÒYÐZgÔZnÑoÔoˆOØ×+Ò+°oÐ+ÑFÔFˆLØ'Ÿ?š?¨=Ô+?ÀÔATÑUÔUˆLØ",×"6Ò"6°rÑ":Ô":ÐÝ!œKÐ(;¸\È=ÑYÔYˆMˆMà#Ðð Ð#Ø!×4Ò4°\ÈtÐ4ÑTÔTÔbˆNØ+×.Ò.¨}Ô/CÀ]ÔEXÑYÔYˆNØ$(×$=Ò$=Ø¨À~ð %>ñ %ô %Ñ!Ð ð *×8Ò8Ð9KÈ^Ñ\Ô\ˆMð Ñ%Ð*=Ñ*IØ ×3Ò3°NÐEXÐDXÐfjÐ3ÑkÔkˆMØ*Ô8ˆNØ&Ô5ˆJõ "'¤°´À!Ñ0CÐ/DÐ.EÍUÌZÐ`nÔ`uÐ!vÑ!vÔ!vÐØ!%×!1Ò!1Ð<NÐ!1Ñ!OÔ!OÐÝ"œ[¨×)=Ò)=¸bÑ)AÔ)AÐCUÐWeÑfÔfˆNà?MÔ?SÑ<Ð˜m¨_Ø#'¤;Ô#JÈ]Ñ#ZÐ Ø%7×%>Ò%>Ð?OÐQeÐgvÑ%wÔ%wÐ"å"œY¨Ð8NÐ'OÐUVÐWÑWÔWˆNØ+×.Ò.¨}Ô/CÀ]ÔEXÑYÔYˆNØ$(×$=Ò$=Ø¨À~ð %>ñ %ô %Ñ!ˆAÐ!ð *×8Ò8Ð9KÈ^Ñ\Ô\ˆMà%�$Ô%ð 

ØØ-Ø)Ø%Ø+Ø'ØØð

ð 

ð ð

ð 

ˆõ *Ø%Ô7Ø7@ÐJ˜GÔ3Ð3ÀdØ!Ô/ØÔ)Ø2>Ð2J  ÐPTØ2@Ð2L  ÐRVð
ñ 
ô 
ð 	
r`   zPProjects the last hidden state from the audio encoder into language model space.c                 ól   —  | j         ||fddi|¤Ž}|                      |j        ¬¦  «        }||_        |S )a0  
        input_features (`torch.FloatTensor]` of shape `(num_images, seq_length, num_features)`):
            The tensors corresponding to the input audio.
        input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`):
            The attention mask for the input audio.
        rr  Trs  )ri  rk  r  ru  )rh   r„  r…  ri   rš  r˜  s         r^   r�  zGemma3nModel.get_audio_featuresü  s]   € ð 9I¸Ô8HØÐ/ð9
ð 9
Ø=Að9
ØEKð9
ð 9
ˆð ×'Ò'°mÔ6UÐ'ÑVÔVˆØ&2ˆÔ#àÐr`   )NNNN)NNNNNNNNNNN)r|   r}   r~   r·   râ   rñ  rô  r   r   rÒ   rØ   r   r   r«   r   rw  rÑ  rƒ  rö   r   r‡   rÕ   rõ   rÐ   r�  rÍ   rÎ   s   @r^   rd  rd    sÞ  ø€ € € € € ð^˜}ð ^ð ^ð ^ð ^ð ^ð ^ð:ð :ð :ð;ð ;ð ;ð Ø€^Ð!rÐsÑsÔsðàÔ'ðð Ð+Ô,ðð 
Ð+Ñ	+ð	ð ð ñ tÔsñ Ôðð, .2Ø26Ø37Ø37ð*6ð *6àÔ# dÑ*ð*6ð Ô(¨4Ñ/ð*6ð Ô)¨DÑ0ð	*6ð
 Ô)¨DÑ0ð*6ð *6ð *6ð *6ðX ð .2Ø15Ø37Ø.2Ø37Ø04Ø(,Ø26Ø26Ø*.Ø!%ðF
ð F
àÔ# dÑ*ðF
ð Ô'¨$Ñ.ðF
ð Ô)¨DÑ0ð	F
ð
 œ tÑ+ðF
ð #œ\¨DÑ0ðF
ð Ô&¨Ñ-ðF
ð  ™ðF
ð Ô(¨4Ñ/ðF
ð Ô(¨4Ñ/ðF
ð Ô  4Ñ'ðF
ð ˜$‘;ðF
ð Ð.Ô/ðF
ð 
$ðF
ð F
ð F
ñ ÔðF
ðP Ø€^Ð!sÐtÑtÔtðàœðð #œ\ðð Ð+Ô,ð	ð
 
Ð/Ñ	/ðð ð ñ uÔtñ Ôðð ð ð ð r`   rd  z†
    The base Gemma 3n model comprising a vision backbone, an audio backbone, a language model, and a language modeling
    head.
    c                   ó€  ‡ — e Zd ZdZd„ Zd„ Zee	 	 	 	 	 	 	 	 	 	 	 	 ddej	        dz  dej
        dz  dej
        dz  d	ej        dz  d
ej        dz  dej	        dz  dedz  dej	        dz  dej
        dz  dej	        dz  dedz  deej        z  dee         defd„¦   «         ¦   «         Z	 	 	 	 	 	 	 	 	 	 	 	 dˆ fd„	Zd„ Zˆ xZS )ÚGemma3nForConditionalGenerationFc                 ó4   — | j                              ¦   «         S rò   )Úmodelrñ  rð  s    r^   rñ  z>Gemma3nForConditionalGeneration.get_per_layer_input_embeddings	  s   € ØŒz×8Ò8Ñ:Ô:Ð:r`   c                 ó:   — | j                              |¦  «         d S rò   )r§  rô  ró  s     r^   rô  z>Gemma3nForConditionalGeneration.set_per_layer_input_embeddings	  s   € ØŒ
×1Ò1°%Ñ8Ô8Ð8Ð8Ð8r`   Nr   r6  rp  r„  r¤  r…  r¾  r¥  r†  rû  r‡  rÆ   Úlogits_to_keeprˆ  rð   c                 ó  —  | j         d|||||||||	|
|ddœ|¤Ž}|j        }t          |t          ¦  «        rt	          | d¦  «        n|}|                      |dd…|dd…f         ¦  «        }| j                             ¦   «         j        x}�||z  }t          j
        |¦  «        }||z  }d}|
�, | j        ||
| j                             ¦   «         j        fi |¤Ž}t          |||j        |j        |j        |j        |j        ¬¦  «        S )aŒ  
        input_features_mask (torch.Tensor, *optional*, defaults to None):
            The attention mask for the input audio.
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are
            ignored (masked), the loss is only computed for the tokens with labels in
            `[0, ..., config.text_config.vocab_size]`.

        Example:

        ```python
        >>> from PIL import Image
        >>> import httpx
        >>> from io import BytesIO
        >>> from transformers import AutoProcessor, Gemma3ForConditionalGeneration

        >>> model = Gemma3ForConditionalGeneration.from_pretrained("google/gemma-3-4b-it")
        >>> processor = AutoProcessor.from_pretrained("google/gemma-3-4b-it")

        >>> messages = [
        ...     {
        ...         "role": "system",
        ...         "content": [
        ...             {"type": "text", "text": "You are a helpful assistant."}
        ...         ]
        ...     },
        ...     {
        ...         "role": "user", "content": [
        ...             {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"},
        ...             {"type": "text", "text": "Where is the cat standing?"},
        ...         ]
        ...     },
        ... ]

        >>> inputs = processor.apply_chat_template(
        ...     messages,
        ...     tokenizer=True,
        ...     return_dict=True,
        ...     return_tensors="pt",
        ...     add_generation_prompt=True
        ... )
        >>> # Generate
        >>> generate_ids = model.generate(**inputs)
        >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "user\nYou are a helpful assistant.\n\n\n\n\n\nWhere is the cat standing?\nmodel\nBased on the image, the cat is standing in a snowy area, likely outdoors. It appears to"
        ```
        T)r6  rp  r„  r¤  r…  r¾  r¥  r†  rû  r‡  rÆ   rr  N)ÚlossrŸ  r¥  rç   r×  rŠ  rÖ   r[   )r§  r  rc   r‚   ÚsliceÚlm_headrù   rþ  rI   rÒ   r‰  Úloss_functionr;   rÚ   r¥  rç   r×  rŠ  rÖ   )rh   r6  rp  r„  r¤  r…  r¾  r¥  r†  rû  r‡  rÆ   r©  rˆ  r¢  rç   Úslice_indicesrŸ  rI   r«  s                       r^   rõ   z'Gemma3nForConditionalGeneration.forward"	  s^  € ðD �$”*ð 
ØØ%Ø)Ø)Ø 3Ø%Ø+Ø)Ø'ØØØð
ð 
ð ð
ð 
ˆð   Ô1ˆå8BÀ>ÕSVÑ8WÔ8WÐk�˜~˜o¨tÑ4Ô4Ð4Ð]kˆØ—’˜m¨A¨A¨A¨}¸a¸a¸aÐ,?Ô@ÑAÔAˆØ'+¤{×'BÒ'BÑ'DÔ'DÔ'\Ð\Ð#ÐiØÐ5Ñ5ˆFÝ”Z Ñ'Ô'ˆFØÐ5Ñ5ˆFàˆØÐØ%�4Ô% f¨f°d´k×6QÒ6QÑ6SÔ6SÔ6^ÐlÐlÐbkÐlÐlˆDå,ØØØ#Ô3Ø!Ô/ØÔ)Ø 'Ô ;Ø 'Ô ;ð
ñ 
ô 
ð 	
r`   Tc                 ót   •—  t          ¦   «         j        |f|||||
||	|dœ|¤Ž}|s|
s||d<   ||d<   ||d<   |S )N)r¥  rû  r¤  r¾  rÆ   r©  r†  Úis_first_iterationrp  r„  r…  )rÊ   Úprepare_inputs_for_generation)rh   r6  r¥  rû  r¾  rp  r„  r¤  r…  r†  rÆ   r©  r‡  r±  ri   Úmodel_inputsrË   s                   €r^   r²  z=Gemma3nForConditionalGeneration.prepare_inputs_for_generation‹	  s…   ø€ ð$ =•u‘w”wÔ<Øð
à+Ø'Ø)Ø%ØØ)Ø)Ø1ð
ð 
ð ð
ð 
ˆð  ð 	F Yð 	FØ+7ˆL˜Ñ(Ø-;ˆLÐ)Ñ*Ø2EˆLÐ.Ñ/àÐr`   c                 ó    — t          d¦  «        ‚)Nz7Do not inherit create_masks_for_generate from PaliGemma)rˆ   )rh   Úsuper_kwargss     r^   Úcreate_masks_for_generatez9Gemma3nForConditionalGeneration.create_masks_for_generate´	  s   € ÝÐVÑWÔWÐWr`   )NNNNNNNNNNNr   )NNNNNNNNTNNF)r|   r}   r~   Úaccepts_loss_kwargsrñ  rô  r   r   rÒ   rÑ  rØ   rö   r   r‡   r‚   r   r   rÚ   rõ   r²  r¶  rÍ   rÎ   s   @r^   r¥  r¥  	  sï  ø€ € € € € ð  Ðð;ð ;ð ;ð9ð 9ð 9ð Øð .2Ø15Ø37Ø.2Ø37Ø04Ø(,Ø26Ø26Ø*.Ø!%Ø-.ðe
ð e
àÔ# dÑ*ðe
ð Ô'¨$Ñ.ðe
ð Ô)¨DÑ0ð	e
ð
 œ tÑ+ðe
ð #œ\¨DÑ0ðe
ð Ô&¨Ñ-ðe
ð  ™ðe
ð Ô(¨4Ñ/ðe
ð Ô(¨4Ñ/ðe
ð Ô  4Ñ'ðe
ð ˜$‘;ðe
ð ˜eœlÑ*ðe
ð Ð.Ô/ðe
ð 
'ðe
ð e
ð e
ñ „^ñ Ôðe
ðT ØØØØØØ ØØØØØ ð'ð 'ð 'ð 'ð 'ð 'ðRXð Xð Xð Xð Xð Xð Xr`   r¥  )
r�   r
  r·   rS  r¥  rd  rÓ  r3   rä  r­   )rW   )sr	  Úcollectionsr   Úcollections.abcr   r   Údataclassesr   Útypingr   rÒ   Útorch.nnrã   Útorch.nn.functionalr(  rè  Úhuggingface_hub.dataclassesr   Ú r
   rÝ  rV  r   Úcache_utilsr   r   Úconfiguration_utilsr   Úmasking_utilsr   r   Úmodeling_outputsr   r   Úmodeling_rope_utilsr   Úmodeling_utilsr   r   Úprocessing_utilsr   Úutilsr   r   r   r   r   r   Úutils.genericr   Úutils.output_capturingr   Úautor    Úgemma2.modeling_gemma2r!   r"   r#   r$   Úgemma3.configuration_gemma3r%   Úgemma3.modeling_gemma3r&   r'   r(   r)   r*   Úpaligemma.modeling_paligemmar+   r,   r-   r.   Ú'timm_wrapper.configuration_timm_wrapperr/   Úaccelerate.hooksr0   Ú
get_loggerr|   rÈ   r3   r�   r­   r·   rÐ   rÕ   rÚ   ÚModulerÜ   rø   rI  r­  rÍ  rï  r  r  r$  r4  r>  rB  rK  rh  rö   r‚   r‹  r�  r³  rÓ  r
  rç  rä  rS  rU  rd  r¥  Ú__all__r[   r`   r^   ú<module>rÔ     s
  ðð €€€Ø  Ð  Ð  Ð  Ð  Ð  Ø .Ð .Ð .Ð .Ð .Ð .Ð .Ð .Ø !Ð !Ð !Ð !Ð !Ð !Ø Ð Ð Ð Ð Ð à €€€Ø Ð Ð Ð Ð Ð Ø Ð Ð Ð Ð Ð Ð Ð Ð Ø .Ð .Ð .Ð .Ð .Ð .à &Ð &Ð &Ð &Ð &Ð &Ø !Ð !Ð !Ð !Ð !Ð !Ø .Ð .Ð .Ð .Ð .Ð .Ð .Ð .Ø 3Ð 3Ð 3Ð 3Ð 3Ð 3Ø RÐ RÐ RÐ RÐ RÐ RÐ RÐ RØ SÐ SÐ SÐ SÐ SÐ SÐ SÐ SØ 6Ð 6Ð 6Ð 6Ð 6Ð 6Ø FÐ FÐ FÐ FÐ FÐ FÐ FÐ FØ &Ð &Ð &Ð &Ð &Ð &ðð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð 8Ð 7Ð 7Ð 7Ð 7Ð 7Ø 5Ð 5Ð 5Ð 5Ð 5Ð 5Ø Ð Ð Ð Ð Ð ðð ð ð ð ð ð ð ð ð ð ð ð ;Ð :Ð :Ð :Ð :Ð :ðð ð ð ð ð ð ð ð ð ð ð ð ð ðð ð ð ð ð ð ð ð ð ð ð ð HÐ GÐ GÐ GÐ GÐ Gð ÐÑÔð 4Ø3Ð3Ð3Ð3Ð3Ð3à	ˆÔ	˜HÑ	%Ô	%€ð €Ð0Ð1Ñ1Ô1ØðGð Gð Gð Gð GÐ(ñ Gô Gñ „ñ 2Ô1ðGðT €Ð0Ð1Ñ1Ô1Øð_ð _ð _ð _ð _Ð)ñ _ô _ñ „ñ 2Ô1ð_ðD €Ð0Ð1Ñ1Ô1Øð%#ð %#ð %#ð %#ð %#Ð+ñ %#ô %#ñ „ñ 2Ô1ð%#ðP €Ð0Ð1Ñ1Ô1ØðP(ð P(ð P(ð P(ð P(Ð$ñ P(ô P(ñ „ñ 2Ô1ðP(ðf Ø
ð3ð 3ð 3ð 3ð 3Ð%?ñ 3ô 3ñ „ñ „ð3ð9ð 9ð 9ð 9ð 9Ð!=ñ 9ô 9ð 9ð$9ð 9ð 9ð 9ð 9Ð$Cñ 9ô 9ð 9ð,4ð 4ð 4ð 4ð 4�R”Yñ 4ô 4ð 4ð0g)ð g)ð g)ð g)ð g)¨B¬Iñ g)ô g)ð g)ðTað að að að a˜BœIñ aô að aðHj,ð j,ð j,ð j,ð j, b¤iñ j,ô j,ð j,ðZB7ð B7ð B7ð B7ð B7 ¤	ñ B7ô B7ð B7ðJFð Fð Fð Fð F¨"¬)ñ Fô Fð FðROð Oð Oð Oð O R¤Yñ Oô Oð Oð8Dð Dð Dð Dð D r¤yñ Dô Dð Dð2(ð (ð (ð (ð ( r¤yñ (ô (ð (ðVð ð ð ð  ¤ñ ô ð ð<	ð 	ð 	ð 	ð 	Ð%Bñ 	ô 	ð 	ð;ð ;ð ;ð ;ð ;˜RœYñ ;ô ;ð ;ð$5ð 5ð 5ð 5ð 5�Yñ 5ô 5ð 5ð@`'ð `'ð `'ð `'ð `'�r”yñ `'ô `'ð `'ðF.ð .˜EœLð .¨u¬|ð .À%Ä,ð .Ð_bð .ð .ð .ð .ð,l)ð l)ð l)ð l)ð l)˜2œ9ñ l)ô l)ð l)ð^B%ð B%ð B%ð B%ð B%Ð0ñ B%ô B%ð B%ðJVJð VJð VJð VJð VJÐ2ñ VJô VJð VJðrN
ð N
ð N
ð N
ð N
Ð0ñ N
ô N
ð N
ðb	ð 	ð 	ð 	ð 	Ð2ñ 	ô 	ð 	ð €ÐaÐbÑbÔbðC
ð C
ð C
ð C
ð C
�ñ C
ô C
ñ cÔbðC
ðL €Ð^Ð_Ñ_Ô_ð	ð 	ð 	ð 	ð 	Ð*ñ 	ô 	ñ `Ô_ð	ð/Bð /Bð /Bð /Bð /B ¤	ñ /Bô /Bð /Bðd €ððñ ô ðoð oð oð oð o�>ñ oô oñô ðoðd €ððñ ô ð\Xð \Xð \Xð \Xð \XÐ&Gñ \Xô \Xñô ð\Xð~ð ð €€€r`   