§
    Ó˜ºhWg  ã            &       óÐ  — d Z ddlmZmZmZ ddlZddlmZ ddlmZm	Z	m
Z
mZmZmZmZmZmZmZmZmZmZmZmZmZ ddgZ G d	„ de¦  «        Zd
de› de› de› de› de
› d�z   e_         dee         dee         dee         dee         dee         dee         dededededededededededef"d „Zdee         dee         dee         dee         dee         dee         dededededededededededef"d!„Z ee¬"¦  «        	 	 	 	 	 	 d&dee         dee         dee         dee         dee         dee         ded$ee         dedededededededededef$d%„¦   «         ZdS )'z'Implementation for the NAdam algorithm.é    )ÚcastÚOptionalÚUnionN)ÚTensoré   )Ú_capturable_docÚ_default_to_fused_or_foreachÚ_differentiable_docÚ_disable_dynamo_if_unsupportedÚ_foreach_docÚ!_get_capturable_supported_devicesÚ_get_scalar_dtypeÚ
_get_valueÚ_maximize_docÚ_params_docÚ_stack_if_compilingÚ
_to_scalarÚ_use_grad_for_differentiableÚ_view_as_realÚ	OptimizerÚParamsTÚNAdamÚnadamc                   óº   ‡ — e Zd Z	 	 	 	 	 	 ddddddœd	ed
eeef         deeef         dedededede	e         dededefˆ fd„Z
ˆ fd„Zd„ Zedd„¦   «         Zˆ xZS )r   çü©ñÒMb`?©gÍÌÌÌÌÌì?g+‡ÙÎ÷ï?ç:Œ0âŽyE>r   çü©ñÒMbp?FN)ÚforeachÚmaximizeÚ
capturableÚdifferentiableÚparamsÚlrÚbetasÚepsÚweight_decayÚmomentum_decayÚdecoupled_weight_decayr   r    r!   r"   c                ój  •— t          |t          ¦  «        r'|                     ¦   «         dk    rt          d¦  «        ‚d|k    st          d|› �¦  «        ‚d|k    st          d|› �¦  «        ‚d|d         cxk    rdk     sn t          d|d         › �¦  «        ‚d|d         cxk    rdk     sn t          d	|d         › �¦  «        ‚d|k    st          d
|› �¦  «        ‚d|k    st          d|› �¦  «        ‚t	          |||||||	||
|¬¦
  «
        }t          ¦   «                              ||¦  «         d S )Nr   zTensor lr must be 1-elementç        zInvalid learning rate: zInvalid epsilon value: r   ç      ð?z#Invalid beta parameter at index 0: z#Invalid beta parameter at index 1: zInvalid weight_decay value: zInvalid momentum_decay value: )
r$   r%   r&   r'   r(   r)   r    r   r!   r"   )Ú
isinstancer   ÚnumelÚ
ValueErrorÚdictÚsuperÚ__init__)Úselfr#   r$   r%   r&   r'   r(   r)   r   r    r!   r"   ÚdefaultsÚ	__class__s                €úY/var/www/html/services/rag-service/venv/lib/python3.11/site-packages/torch/optim/nadam.pyr2   zNAdam.__init__!   sŠ  ø€ õ �b�&Ñ!Ô!ð 	< b§h¢h¡j¤j°A¢o oÝÐ:Ñ;Ô;Ð;Ø�bŠyˆyÝÐ;°rÐ;Ð;Ñ<Ô<Ð<Ø�cŠzˆzÝÐ<°sÐ<Ð<Ñ=Ô=Ð=Ø�e˜A”hÐ$Ð$Ò$Ð$ Ò$Ð$Ð$Ð$ÝÐMÀ5ÈÄ8ÐMÐMÑNÔNÐNØ�e˜A”hÐ$Ð$Ò$Ð$ Ò$Ð$Ð$Ð$ÝÐMÀ5ÈÄ8ÐMÐMÑNÔNÐNØ�lÒ"Ð"ÝÐJ¸LÐJÐJÑKÔKÐKØ�nÒ$Ð$ÝÐN¸nÐNÐNÑOÔOÐOÝØØØØ%Ø)Ø#9ØØØ!Ø)ð
ñ 
ô 
ˆõ 	‰Œ×Ò˜ Ñ*Ô*Ð*Ð*Ð*ó    c                 óª  •— t          ¦   «                              |¦  «         | j        D �]§}|                     dd¦  «         |                     dd ¦  «         |                     dd¦  «         |                     dd¦  «         |                     dd¦  «         |d         D �],}| j                             |g ¦  «        }t          |¦  «        dk    rût          j        |d	         ¦  «        sjt          |d	         ¦  «        }|d         r(t          j
        |t          ¦   «         |j        ¬
¦  «        n!t          j
        |t          ¦   «         ¬¦  «        |d	<   t          j        |d         ¦  «        s]|d         }|d         r(t          j
        |t          ¦   «         |j        ¬
¦  «        n!t          j
        |t          ¦   «         ¬¦  «        |d<   �Œ.�Œ©d S )Nr    Fr   r!   r"   r)   r#   r   Ústep©ÚdtypeÚdevice©r;   Ú
mu_product)r1   Ú__setstate__Úparam_groupsÚ
setdefaultÚstateÚgetÚlenÚtorchÚ	is_tensorÚfloatÚtensorr   r<   )r3   rB   ÚgroupÚpÚp_stateÚstep_valÚmu_prod_valr5   s          €r6   r?   zNAdam.__setstate__L   sé  ø€ Ý‰Œ×Ò˜UÑ#Ô#Ð#ØÔ&ð 	ñ 	ˆEØ×Ò˜Z¨Ñ/Ô/Ð/Ø×Ò˜Y¨Ñ-Ô-Ð-Ø×Ò˜\¨5Ñ1Ô1Ð1Ø×ÒÐ-¨uÑ5Ô5Ð5Ø×ÒÐ5°uÑ=Ô=Ð=Ø˜8”_ð ñ �Øœ*Ÿ.š.¨¨BÑ/Ô/�Ý�w‘<”< 1Ò$Ð$Ý œ?¨7°6¬?Ñ;Ô;ð Ý#(¨°¬Ñ#9Ô#9˜ð
  % \Ô2ðS�EœLØ (Õ0AÑ0CÔ0CÈAÌHðñ ô ð õ "'¤¨hÕ>OÑ>QÔ>QÐ!RÑ!RÔ!Rð   ™õ !œ?¨7°<Ô+@ÑAÔAð Ø&-¨lÔ&;˜ð
  % \Ô2ðV�EœLØ +Õ3DÑ3FÔ3FÈqÌxðñ ô ð õ "'¤¨kÕARÑATÔATÐ!UÑ!UÔ!Uð   Ñ-ùñð	ð 	r7   c                 óî  — d}|d         D �]è}	|	j         ��Ý|t          j        |	¦  «        z  }|                     |	¦  «         |	j         j        rt          d¦  «        ‚|                     |	j         ¦  «         | j        |	         }
t          |
¦  «        dk    rð|d         r(t          j        dt          ¦   «         |	j
        ¬¦  «        n!t          j        dt          ¦   «         ¬	¦  «        |
d
<   |d         r(t          j        dt          ¦   «         |	j
        ¬¦  «        n!t          j        dt          ¦   «         ¬	¦  «        |
d<   t          j        |	t          j        ¬¦  «        |
d<   t          j        |	t          j        ¬¦  «        |
d<   |                     |
d         ¦  «         |                     |
d         ¦  «         |                     |
d         ¦  «         |                     |
d
         ¦  «         �Œê|S )NFr#   z'NAdam does not support sparse gradientsr   r!   © r:   r+   r=   r9   r,   r>   )Úmemory_formatÚexp_avgÚ
exp_avg_sq)ÚgradrE   Ú
is_complexÚappendÚ	is_sparseÚRuntimeErrorrB   rD   Úzerosr   r<   rH   ÚonesÚ
zeros_likeÚpreserve_format)r3   rI   Úparams_with_gradÚgradsÚexp_avgsÚexp_avg_sqsÚmu_productsÚstate_stepsÚhas_complexrJ   rB   s              r6   Ú_init_groupzNAdam._init_groupj   sñ  € ð ˆØ�x”ð $	2ñ $	2ˆAØŒvÑ!Ø�uÔ/°Ñ2Ô2Ñ2�Ø ×'Ò'¨Ñ*Ô*Ð*Ø”6Ô#ð RÝ&Ð'PÑQÔQÐQØ—’˜QœVÑ$Ô$Ð$àœ
 1œ�å�u‘:”: ’?�?ð ! Ô.ðJ�œ BÕ.?Ñ.AÔ.AÈ!Ì(ÐSÑSÔSÐSå"œ\¨#Õ5FÑ5HÔ5HÐIÑIÔIð ˜&‘Mð ! Ô.ðJ�œ
 2Õ->Ñ-@Ô-@ÈÌÐRÑRÔRÐRå"œ\¨#Õ5FÑ5HÔ5HÐIÑIÔIð ˜,Ñ'õ (-Ô'7Ø­Ô)>ð(ñ (ô (�E˜)Ñ$õ +0Ô*:Ø­Ô)>ð+ñ +ô +�E˜,Ñ'ð —’  iÔ 0Ñ1Ô1Ð1Ø×"Ò" 5¨Ô#6Ñ7Ô7Ð7Ø×"Ò" 5¨Ô#6Ñ7Ô7Ð7Ø×"Ò" 5¨¤=Ñ1Ô1Ð1ùØÐr7   c                 ó  — |                       ¦   «          d}|�5t          j        ¦   «         5   |¦   «         }ddd¦  «         n# 1 swxY w Y   | j        D ]°}g }g }g }g }g }g }	t	          t
          t          t          f         |d         ¦  «        \  }
}|                      |||||||	¦  «        }t          ||||||	|
||d         |d         |d         |d         |d         |d         |d	         |d
         |d         |¬¦  «         Œ±|S )z°Perform a single optimization step.

        Args:
            closure (Callable, optional): A closure that reevaluates the model
                and returns the loss.
        Nr%   r$   r'   r(   r&   r    r)   r   r!   r"   )Úbeta1Úbeta2r$   r'   r(   r&   r    r)   r   r!   r"   rb   )	Ú _cuda_graph_capture_health_checkrE   Úenable_gradr@   r   ÚtuplerG   rc   r   )r3   ÚclosureÚlossrI   r\   r]   r^   r_   r`   ra   re   rf   rb   s                r6   r9   z
NAdam.stepœ   s–  € ð 	×-Ò-Ñ/Ô/Ð/àˆØÐÝÔ"Ñ$Ô$ð !ð !Ø�w‘y”y�ð!ð !ð !ñ !ô !ð !ð !ð !ð !ð !ð !øøøð !ð !ð !ð !ð Ô&ð &	ð &	ˆEØ-/ÐØ"$ˆEØ%'ˆHØ(*ˆKØ(*ˆKØ(*ˆKÝ¥¥e­U lÔ 3°U¸7´^ÑDÔD‰LˆE�5à×*Ò*ØØ ØØØØØñô ˆKõ Ø ØØØØØØØØ˜”;Ø" >Ô2Ø$Ð%5Ô6Ø˜%”LØ˜zÔ*Ø',Ð-EÔ'FØ˜iÔ(Ø  Ô.Ø$Ð%5Ô6Ø'ð%ñ ô ð ð ð* ˆs   ¬AÁAÁ
A)r   r   r   r   r   F©N)Ú__name__Ú
__module__Ú__qualname__r   r   rG   r   ri   Úboolr   r2   r?   rc   r   r9   Ú__classcell__)r5   s   @r6   r   r       s;  ø€ € € € € ð $(Ø%1ØØØ $Ø',ð)+ð #'ØØ Ø$ð)+ð )+ð )+àð)+ð �%˜�-Ô ð)+ð �U˜E�\Ô"ð	)+ð
 ð)+ð ð)+ð ð)+ð !%ð)+ð ˜$”ð)+ð ð)+ð ð)+ð ð)+ð )+ð )+ð )+ð )+ð )+ðVð ð ð ð ð<0ð 0ð 0ðd "ð6ð 6ð 6ñ "Ô!ð6ð 6ð 6ð 6ð 6r7   a¼  Implements NAdam algorithm.

    .. math::
       \begin{aligned}
            &\rule{110mm}{0.4pt}                                                                 \\
            &\textbf{input}      : \gamma_t \text{ (lr)}, \: \beta_1,\beta_2 \text{ (betas)},
                \: \theta_0 \text{ (params)}, \: f(\theta) \text{ (objective)}                   \\
            &\hspace{13mm} \: \lambda \text{ (weight decay)}, \:\psi \text{ (momentum decay)}    \\
            &\hspace{13mm} \: \textit{decoupled\_weight\_decay}, \:\textit{maximize}             \\
            &\textbf{initialize} :  m_0 \leftarrow 0 \text{ ( first moment)},
                v_0 \leftarrow 0 \text{ ( second moment)}                                 \\[-1.ex]
            &\rule{110mm}{0.4pt}                                                                 \\
            &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
            &\hspace{5mm}\textbf{if} \: \textit{maximize}:                                       \\
            &\hspace{10mm}g_t           \leftarrow   -\nabla_{\theta} f_t (\theta_{t-1})         \\
            &\hspace{5mm}\textbf{else}                                                           \\
            &\hspace{10mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})          \\
            &\hspace{5mm} \theta_t \leftarrow \theta_{t-1}                                       \\
            &\hspace{5mm} \textbf{if} \: \lambda \neq 0                                          \\
            &\hspace{10mm}\textbf{if} \: \textit{decoupled\_weight\_decay}                       \\
            &\hspace{15mm} \theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1}                    \\
            &\hspace{10mm}\textbf{else}                                                          \\
            &\hspace{15mm} g_t \leftarrow g_t + \lambda \theta_{t-1}                             \\
            &\hspace{5mm} \mu_t \leftarrow \beta_1 \big(1 - \frac{1}{2}  0.96^{t \psi} \big)     \\
            &\hspace{5mm} \mu_{t+1} \leftarrow \beta_1 \big(1 - \frac{1}{2} 0.96^{(t+1)\psi}\big)\\
            &\hspace{5mm}m_t           \leftarrow   \beta_1 m_{t-1} + (1 - \beta_1) g_t          \\
            &\hspace{5mm}v_t           \leftarrow   \beta_2 v_{t-1} + (1-\beta_2) g^2_t          \\
            &\hspace{5mm}\widehat{m_t} \leftarrow \mu_{t+1} m_t/(1-\prod_{i=1}^{t+1}\mu_i)\\[-1.ex]
            & \hspace{11mm} + (1-\mu_t) g_t /(1-\prod_{i=1}^{t} \mu_{i})                         \\
            &\hspace{5mm}\widehat{v_t} \leftarrow   v_t/\big(1-\beta_2^t \big)                   \\
            &\hspace{5mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t}/
                \big(\sqrt{\widehat{v_t}} + \epsilon \big)                                       \\
            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
            &\bf{return} \:  \theta_t                                                     \\[-1.ex]
            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
       \end{aligned}

    For further details regarding the algorithm we refer to `Incorporating Nesterov Momentum into Adam`_.
    z
    Args:
        aø  
        lr (float, Tensor, optional): learning rate (default: 2e-3)
        betas (Tuple[float, float], optional): coefficients used for computing
            running averages of gradient and its square (default: (0.9, 0.999))
        eps (float, optional): term added to the denominator to improve
            numerical stability (default: 1e-8)
        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
        momentum_decay (float, optional): momentum momentum_decay (default: 4e-3)
        decoupled_weight_decay (bool, optional): whether to decouple the weight
            decay as in AdamW to obtain NAdamW. If True, the algorithm does not
            accumulate weight decay in the momentum nor variance. (default: False)
        z	
        zÏ

    .. _Incorporating Nesterov Momentum into Adam:
        https://openreview.net/forum?id=OM0jvwB8jIp57ZJjtNEZ
    .. _Decoupled Weight Decay Regularization:
        https://arxiv.org/abs/1711.05101

    r#   r]   r^   r_   r`   ra   re   rf   r$   r'   r(   r&   r)   r    r!   r"   rb   c          
      ód  — t           j                             ¦   «         st          |¦  «        }t	          | ¦  «        D �]ñ\  }}|s||         n||          }||         }||         }||         }||         }t          j        |¦  «        rPt          j        |¦  «        }t          j        |¦  «        }t          j        |¦  «        }t          j        |¦  «        }t           j                             ¦   «         sZ|rXt          ¦   «         }|j
        j        |j
        j        cxk    r|j
        j        k    rn n|j
        j        |v sJ d|› d�¦   «         ‚|dz  }|r|}nt          |¦  «        }d||z  z
  }|	dk    r5|r|                     d||	z  z
  ¦  «         n|                     ||	¬¦  «        }|ddd||
z  z  z  z
  z  }|ddd|dz   |
z  z  z  z
  z  }||z  }|                     |d|z
  ¦  «         |                     |¦  «                             ||d|z
  ¬	¦  «         |                     |¦  «                             ¦   «         }|s|ri|                     |¦  «        }||z  }|| d|z
  z  d|z
  z  z  }|| |z  d|z
  z  z  }|                     ||¦  «         |                     ||¦  «         �Œvt          |¦  «        |z  }|                     |¦  «         |                     ||| d|z
  z  dt          |¦  «        z
  z  ¬	¦  «         |                     ||| |z  d|z
  z  ¬	¦  «         �Œód S )
NzVIf capturable=True, params, mu_products and state_steps must be on supported devices: ú.r   r   ©Úalphar,   ç      à?ç¸…ëQ¸î?)Úvalue)rE   ÚjitÚis_scriptingr   Ú	enumeraterT   Úview_as_realÚcompilerÚis_compilingr   r<   Útyper   Úmul_ÚaddÚlerp_Úaddcmul_ÚdivÚsqrtÚaddcdiv_Úadd_)r#   r]   r^   r_   r`   ra   re   rf   r$   r'   r(   r&   r)   r    r!   r"   rb   ÚiÚparamrS   rQ   rR   r>   Ústep_tÚcapturable_supported_devicesr9   Úbias_correction2ÚmuÚmu_nextÚdenomÚmu_product_nexts                                  r6   Ú_single_tensor_nadamr‘     sÊ  € õ( Œ9×!Ò!Ñ#Ô#ð Ý˜‰^Œ^ˆå˜fÑ%Ô%ð Gñ G‰ˆˆ5Ø'Ð6ˆu�QŒxˆx¨e°A¬h¨YˆØ˜1”+ˆØ  ”^ˆ
Ø  ”^ˆ
Ø˜Q”ˆåÔ˜EÑ"Ô"ð 	8ÝÔ& uÑ-Ô-ˆEÝÔ% dÑ+Ô+ˆDÝÔ(¨Ñ1Ô1ˆGÝÔ+¨JÑ7Ô7ˆJõ Œ~×*Ò*Ñ,Ô,ð 	°ð 	Ý+LÑ+NÔ+NÐ(à”Ô! ZÔ%6Ô%;ÐQÐQÒQÐQ¸v¼}Ô?QÒQÐQÐQÐQÐQØ”LÔ%Ð)EÐEÐEÐEðIØ)EðIð Ið Iñ FÔEÐEð 	�!‰ˆàð 	&ØˆDˆDå˜fÑ%Ô%ˆDà˜u d™{™?Ðà˜1ÒÐØ%ð ;à—
’
˜1˜r LÑ0Ñ0Ñ1Ô1Ð1Ð1à—x’x ¨\�xÑ:Ô:�ð �c˜C 4¨D°>Ñ,AÑ#BÑCÑCÑDˆØ˜3 ¨°$¸±(¸nÑ1LÑ(MÑ!NÑNÑOˆð 	�bÑˆ
ð 	�Š�d˜A ™IÑ&Ô&Ð&Ø�Š˜ÑÔ×'Ò'¨¨d¸!¸e¹)Ð'ÑDÔDÐDØ—’Ð/Ñ0Ô0×5Ò5Ñ7Ô7ˆàð 	˜Zð 	Ø—I’I˜c‘N”NˆEð )¨7Ñ2ˆOØ˜B˜3 #¨¡(Ñ+¨s°ZÑ/?Ñ@ÑAˆDØ "  w¡°#¸Ñ2GÑ!HÑIˆGØ�NŠN˜4 Ñ'Ô'Ð'Ø�NŠN˜7 EÑ*Ô*Ð*Ñ*å(¨Ñ4Ô4°wÑ>ˆOØ�JŠJ�s‰OŒOˆOØ�NŠNØ�e R C¨3°©8Ñ$4¸½jÈÑ>TÔ>TÑ8TÑ$Uð ñ ô ð ð �NŠNØ˜¨ s¨W¡}¸¸Ñ9NÑ&Oð ñ ô ð ñ ðKGð Gr7   c                ó  ‡‡‡‡
‡(— t          | ¦  «        dk    rd S |r
J d¦   «         ‚t          j                             ¦   «         sJ|rHt	          d¬¦  «        Š(t          ˆ(fd„t          | ||¦  «        D ¦   «         ¦  «        sJ d‰(› d�¦   «         ‚t          ‰¦  «        Št          j	        | |||||g¦  «        }| 
                    ¦   «         D �]1\  \  }}}}}}}t          t          t                   |¦  «        }t          t          t                   |¦  «        }t          t          t                   |¦  «        }t          t          t                   |¦  «        }t          t          t                   |¦  «        }t          t          t                   |¦  «        }|rt          ||||¦  «         |rt          j        |¦  «        }t          j                             ¦   «         s9|d         j        r,t          j        |t          j        dd	¬
¦  «        d¬¦  «         nt          j        |d¦  «         |	dk    rO|rt          j        |d‰|	z  z
  ¦  «         n1|rt          j        |||	¬¦  «         nt          j        |||	¬¦  «        }t          j        ||d‰z
  ¦  «         t          j        |‰¦  «         t          j        |||d‰z
  ¦  «         t          j        |¦  «        }|�r&t          j        |‰
¦  «        } t          j        d| ¦  «        }!t          j        |!d¦  «         t          j        |!d¦  «         t          j        |!‰¦  «         t          j        | ‰
¦  «         t          j        d| ¦  «        }"t          j        |"d¦  «         t          j        |"d¦  «         t          j        |"‰¦  «         ~ t          j        ‰|¦  «        }#t          j        |#d¦  «         t          j        |#¦  «         t          j        |#¦  «         n,ˆfd„|D ¦   «         }#ˆˆ
fd„|D ¦   «         }!ˆˆ
fd„|D ¦   «         }"t          j        ||!¦  «         t          j        ||#¦  «         t          j        ||¦  «         ~#|�rt          j        |!d¦  «         t          j        |!‰¦  «         t          j        |d¦  «        }$t          j        |$¦  «         t          j        |!|$¦  «         |!}%~$t          j        ||"¦  «        }$t          j        |"‰¦  «         t          j        |$d¦  «         t          j        |"|$¦  «         |"}&~$t          j        |%|¦  «        }'t          j        |'|&|¦  «         t          j        ||'|¦  «         �Œ±tA          ˆfd„t          ||!¦  «        D ¦   «         ¦  «        }%tA          ˆfd„t          ||"¦  «        D ¦   «         ¦  «        }&t          j        ||||%¦  «         t          j        ||||&¦  «         �Œ3d S )Nr   z#_foreach ops don't support autogradF)Úsupports_xlac              3   ó˜   •K  — | ]D\  }}}|j         j        |j         j        cxk    o|j         j        k    nc o|j         j        ‰v V — ŒEd S rl   )r<   r   )Ú.0rJ   Úmpr9   r‹   s       €r6   ú	<genexpr>z&_multi_tensor_nadam.<locals>.<genexpr>˜  sz   øè è € ð 
ð 
ñ ��2�tð ŒHŒM˜RœYœ^Ð?Ð?Ò?Ð?¨t¬{Ô/?Ò?Ð?Ð?Ð?ð >Ø””Ð!=Ð=ð
ð 
ð 
ð 
ð 
ð 
r7   zWIf capturable=True, params, mu_products, and state_steps must be on supported devices: rs   r,   Úcpu)r<   rt   r   rw   g      à¿c                 ó@   •— g | ]}d ‰t          |¦  «        z  z
  dz  ‘ŒS )r   rv   ©r   )r•   r9   rf   s     €r6   ú
<listcomp>z'_multi_tensor_nadam.<locals>.<listcomp>þ  s=   ø€ ð $ð $ð $Ø;?��U�j¨Ñ.Ô.Ñ.Ñ.°3Ñ6ð$ð $ð $r7   c           	      óL   •— g | ] }‰d ddt          |¦  «        ‰z  z  z  z
  z  ‘Œ!S )r,   rv   rw   rš   ©r•   r9   re   r(   s     €€r6   r›   z'_multi_tensor_nadam.<locals>.<listcomp>  sJ   ø€ ð ð ð àð ˜˜s d­z¸$Ñ/?Ô/?À.Ñ/PÑ&QÑRÑRÑSðð ð r7   c           	      óR   •— g | ]#}‰d ddt          |¦  «        dz   ‰z  z  z  z
  z  ‘Œ$S )r,   rv   rw   r   rš   r�   s     €€r6   r›   z'_multi_tensor_nadam.<locals>.<listcomp>  sU   ø€ ð ð ð ð ð Ø˜ ­*°TÑ*:Ô*:¸QÑ*>À.Ñ)PÑ QÑRÑRñTðð ð r7   c                 ól   •— g | ]0\  }}t          ‰¦  «        d |z
  z  d t          |¦  «        z
  z  dz  ‘Œ1S ©r,   éÿÿÿÿrš   )r•   r>   r�   r$   s      €r6   r›   z'_multi_tensor_nadam.<locals>.<listcomp>6  sT   ø€ ð ð ð á&˜
 Bõ   ‘^”^ s¨R¡xÑ0°C½*ÀZÑ:PÔ:PÑ4PÑQÐUWÑWðð ð r7   c                 ól   •— g | ]0\  }}t          ‰¦  «        |z  d t          |¦  «        |z  z
  z  dz  ‘Œ1S r    rš   )r•   r>   rŽ   r$   s      €r6   r›   z'_multi_tensor_nadam.<locals>.<listcomp><  sb   ø€ ð ð ð ñ ,˜
 Gõ # 2™œØ!ñ"à¥¨JÑ!7Ô!7¸'Ñ!AÑAñCð ñðð ð r7   )!rD   rE   r}   r~   r   ÚallÚzipr   r   Ú"_group_tensors_by_device_and_dtypeÚvaluesr   Úlistr   r   Ú_foreach_negÚis_cpuÚ_foreach_add_rH   Ú_foreach_mul_Ú_foreach_addÚ_foreach_lerp_Ú_foreach_addcmul_Ú_foreach_sqrtÚ_foreach_mulÚ_foreach_powÚ_foreach_sub_Ú_foreach_neg_Ú_foreach_sqrt_Ú_foreach_div_Ú_foreach_subÚ_foreach_addcdiv_r   ))r#   r]   r^   r_   r`   ra   re   rf   r$   r'   r(   r&   r)   r    r!   r"   rb   Úgrouped_tensorsÚgrouped_params_Úgrouped_grads_Úgrouped_exp_avgs_Úgrouped_exp_avg_sqs_Úgrouped_mu_products_Úgrouped_state_steps_Ú_Úgrouped_paramsÚgrouped_gradsÚgrouped_exp_avgsÚgrouped_exp_avg_sqsÚgrouped_mu_productsÚgrouped_state_stepsÚexp_avg_sq_sqrtÚexponentÚmusÚmu_nextsÚbias_correction_sqrtr�   Ústep_size_gradsÚstep_size_expavgÚ	numeratorr‹   s)         ``` `                             @r6   Ú_multi_tensor_nadamrÎ   z  s*  øøøøø€ õ( ˆ6�{„{�aÒÐØˆàÐDÐDÐDÑDÔDÐDõ Œ>×&Ò&Ñ(Ô(ð 
¨Zð 
Ý'HØð(
ñ (
ô (
Ð$õ ð 
ð 
ð 
ð 
õ  # 6¨;¸ÑDÔDð
ñ 
ô 
ñ 
ô 
ð 	
ð 	
ð
/à+ð/ð /ð /ñ	
ô 	
ð 	
õ 
�B‰Œ€BåÔBØ	�˜ +¨{¸KÐHñô €Oð ×"Ò"Ñ$Ô$ðkñ kñ 		ñ 	ØØØØØØØÝ�d¥6œl¨OÑ<Ô<ˆÝ�T¥&œ\¨>Ñ:Ô:ˆÝ¥¥V¤Ð.?Ñ@Ô@ÐÝ"¥4­¤<Ð1EÑFÔFÐÝ"¥4­¤<Ð1EÑFÔFÐÝ"¥4­¤<Ð1EÑFÔFÐð ð 	ÝØ Ð/?ÐATñô ð ð ð 	>Ý!Ô.¨}Ñ=Ô=ˆMõ Œ~×*Ò*Ñ,Ô,ð 	8Ð1DÀQÔ1GÔ1Nð 	8ÝÔØ#¥U¤\°#¸eÐ%DÑ%DÔ%DÈCðñ ô ð ð õ ÔÐ 3°QÑ7Ô7Ð7à˜1ÒÐØ%ð åÔ# N°A¸¸\Ñ8IÑ4IÑJÔJÐJÐJð ð ÝÔ'Ø% ~¸\ðñ ô ð ð õ %*Ô$6Ø% ~¸\ð%ñ %ô %�Mõ
 	ÔÐ-¨}¸aÀ%¹iÑHÔHÐHåÔÐ/°Ñ7Ô7Ð7ÝÔØ °¸qÀ5¹yñ	
ô 	
ð 	
õ  Ô-Ð.AÑBÔBˆð
 ñ #	åÔ)Ð*=¸~ÑNÔNˆHÝÔ$ T¨8Ñ4Ô4ˆCÝÔ  TÑ*Ô*Ð*ÝÔ  SÑ)Ô)Ð)ÝÔ  UÑ+Ô+Ð+õ Ô ¨.Ñ9Ô9Ð9ÝÔ)¨$°Ñ9Ô9ˆHÝÔ ¨$Ñ/Ô/Ð/ÝÔ ¨#Ñ.Ô.Ð.ÝÔ ¨%Ñ0Ô0Ð0ð å#(Ô#5°eÐ=PÑ#QÔ#QÐ åÔÐ 4°cÑ:Ô:Ð:ÝÔÐ 4Ñ5Ô5Ð5ÝÔ Ð!5Ñ6Ô6Ð6Ð6ð$ð $ð $ð $ØCVð$ñ $ô $Ð ðð ð ð ð à/ðñ ô ˆCðð ð ð ð ð 0ðñ ô ˆHõ 	ÔÐ/°Ñ5Ô5Ð5åÔ˜OÐ-AÑBÔBÐBÝÔ˜O¨SÑ1Ô1Ð1ð !àñ >	åÔ  SÑ)Ô)Ð)ÝÔ  RÑ(Ô(Ð(åÔ&Ð':¸CÑ@Ô@ˆEÝÔ Ñ&Ô&Ð&ÝÔ  UÑ+Ô+Ð+à!ˆOàõ Ô&Ð':¸HÑEÔEˆEÝÔ ¨"Ñ-Ô-Ð-õ Ô  sÑ+Ô+Ð+ÝÔ ¨%Ñ0Ô0Ð0à'Ðàõ Ô*¨?¸MÑJÔJˆIÝÔ# IÐ/?ÐAQÑRÔRÐRõ Ô# N°I¸ÑOÔOÐOÑOå1ðð ð ð å*-Ð.AÀ3Ñ*GÔ*Gðñ ô ñô ˆOõ  3ðð ð ð õ 03Ð3FÈÑ/QÔ/Qðñ ô ñ
 ô 
 Ðõ Ô#ØØØØñ	ô ð õ Ô#ØØ ØØ ñ	ô ð ñ ðMkð kr7   )Úsingle_tensor_fnFr   c                óÈ  — t          d„ |D ¦   «         ¦  «        st          d¦  «        ‚t          d„ |D ¦   «         ¦  «        st          d¦  «        ‚|€t          | |	d¬¦  «        \  }}|r-t          j                             ¦   «         rt          d¦  «        ‚|r&t          j                             ¦   «         st          }nt          } || |||||||||||||||	|
¬	¦  «         dS )
zpFunctional API that performs NAdam algorithm computation.

    See :class:`~torch.optim.NAdam` for details.
    c              3   óJ   K  — | ]}t          |t          j        ¦  «        V — Œd S rl   ©r-   rE   r   ©r•   Úts     r6   r—   znadam.<locals>.<genexpr>q  ó.   è è € Ð@Ð@¨q�z˜!�Uœ\Ñ*Ô*Ð@Ð@Ð@Ð@Ð@Ð@r7   zPAPI has changed, `state_steps` argument must contain a list of singleton tensorsc              3   óJ   K  — | ]}t          |t          j        ¦  «        V — Œd S rl   rÒ   rÓ   s     r6   r—   znadam.<locals>.<genexpr>v  rÕ   r7   zPAPI has changed, `mu_products` argument must contain a list of singleton tensorsNF)Ú	use_fusedz6torch.jit.script not supported with foreach optimizers)re   rf   r$   r'   r(   r    r)   r&   r!   r"   rb   )r£   rW   r	   rE   ry   rz   rÎ   r‘   )r#   r]   r^   r_   r`   ra   r)   r   r!   r"   rb   r    re   rf   r$   r'   r(   r&   r¿   Úfuncs                       r6   r   r   U  sC  € õ8 Ð@Ð@°KÐ@Ñ@Ô@Ñ@Ô@ð 
ÝØ^ñ
ô 
ð 	
õ Ð@Ð@°KÐ@Ñ@Ô@Ñ@Ô@ð 
ÝØ^ñ
ô 
ð 	
ð €Ý1Ø�N¨eð
ñ 
ô 
‰
ˆˆ7ð ð U•5”9×)Ò)Ñ+Ô+ð UÝÐSÑTÔTÐTàð $•u”y×-Ò-Ñ/Ô/ð $Ý"ˆˆå#ˆà€DØØØØØØØØØØ!Ø%ØØ5ØØØ%Øð#ñ ô ð ð ð r7   )FNFFFF) Ú__doc__Útypingr   r   r   rE   r   Ú	optimizerr   r	   r
   r   r   r   r   r   r   r   r   r   r   r   r   r   Ú__all__r   r§   rG   rp   r‘   rÎ   r   rO   r7   r6   ú<module>rÝ      s§  ðà .Ð .à (Ð (Ð (Ð (Ð (Ð (Ð (Ð (Ð (Ð (à €€€Ø Ð Ð Ð Ð Ð ðð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð( �GÐ
€ðsð sð sð sð sˆIñ sô sð sðn&ðNà	ðð ð 
ðð ð 
ðð ð 
ðð ð  
ð!ð ð ñO>ð „ðF^Ø�ŒLð^à�Œ<ð^ð �6Œlð^ð �f”ð	^ð
 �f”ð^ð �f”ð^ð ð^ð ð^ð 	ð^ð ð^ð ð^ð 
ð^ð !ð^ð ð^ð  ð!^ð" ð#^ð$ ð%^ð ^ð ^ð ^ðBXØ�ŒLðXà�Œ<ðXð �6ŒlðXð �f”ð	Xð
 �f”ðXð �f”ðXð ðXð ðXð 	ðXð ðXð ðXð 
ðXð !ðXð ðXð  ð!Xð" ð#Xð$ ð%Xð Xð Xð Xðv  ÐÐ1EÐFÑFÔFð $)Ø"ØØ ØØðDð DØ�ŒLðDà�Œ<ðDð �6ŒlðDð �f”ð	Dð
 �f”ðDð �f”ðDð !ðDð �dŒ^ðDð ðDð ðDð ðDð ðDð  ð!Dð" ð#Dð$ 	ð%Dð& ð'Dð( ð)Dð* 
ð+Dð Dð Dñ GÔFðDð Dð Dr7   