
    [Th+                         % S SK Jr  S SKrS SKJs  Jr  S SKJr  S SKJ	r	  / r
\\   \S'   \R                  R                   " S S5      5       rg)    )OptionalN)Tensor)2_scripted_functional_optimizer_deprecation_warning__all__c                   z    \ rS rSr       SS\\   S\S\\\4   S\S\S\S\S	\4S
 jjr	S\\
\      4S jrSrg)_FunctionalAdamax   paramslrbetasepsweight_decayforeachmaximize_allow_empty_param_listc	                 Z   [        SS9  SU::  d  [        SU 35      eSU::  d  [        SU 35      eSUS   s=::  a  S:  d  O  [        SUS    35      eSUS	   s=::  a  S:  d  O  [        S
US	    35      eSU::  d  [        SU 35      eUUUS   US	   US.U l        X`l        Xpl        [
        R                  R                  [        [
        R                  [        [        [
        R                  4   4   0 5      U l        [        U5      S:X  a  U(       d  [        S5      eSU0U l        g )N   )
stacklevel        zInvalid learning rate: zInvalid epsilon value: r   g      ?z#Invalid beta parameter at index 0:    z#Invalid beta parameter at index 1: zInvalid weight_decay value: )r   r   beta1beta2r   z%optimizer got an empty parameter listr
   )r   
ValueErrordefaultsr   r   torchjitannotatedictr   strstatelenparam_group)	selfr
   r   r   r   r   r   r   r   s	            a/var/www/auris/envauris/lib/python3.13/site-packages/torch/distributed/optim/functional_adamax.py__init___FunctionalAdamax.__init__   s8    	;aHby6rd;<<cz6se<==eAh$$B58*MNNeAh$$B58*MNNl";L>JKK 1X1X(
  YY''U\\4U\\@Q;R-R(SUWX
v;!$;DEE %f-    	gradientsc                    U R                   S   n/ n/ n/ n/ n/ n[        U5      [        U5      :w  a*  [        SS[        U5       S3-   S[        U5       3-   5      eSn[        U R                   S   U5       GH#  u  pU
c  M  U[        R
                  " U	5      -  nUR                  U	5        UR                  U
5        XR                  ;  a  0 U R                  U	'   U R                  U	   n[        R                  " S5      US'   [        R                  " U	[        R                  S	9US
'   [        R                  " U	[        R                  S	9US'   U R                  U	   nUR                  US
   5        UR                  US   5        UR                  US   5        GM&     [        R                  " 5          [        R                  " UUUUUU R                  S   U R                  S   U R                  S   U R                  S   U R                  S   U R                  U R                   US9  S S S 5        g ! , (       d  f       g = f)Nr
   zEthe gradients passed in does not equal to the size of the parameters!zParams length: z. zGradients length: Fr   step)memory_formatexp_avgexp_infr   r   r   r   r   )r   r   r   r   r   r   r   has_complex)r"   r!   r   zipr   
is_complexappendr    tensor
zeros_likepreserve_formatno_gradFadamaxr   r   r   )r#   r(   r
   params_with_gradgradsexp_avgsexp_infsstate_stepsr.   paramgradientr    s               r$   r*   _FunctionalAdamax.stepC   s   !!(+$&v;#i.(W#CK=34&s9~&678  "4#3#3H#=yIOE#u//66 ''.X&

*(*DJJu% JJu-E$)LL$5E&M','7'7U-B-B(E)$ (-'7'7U-B-B(E)$ 

5)i 01i 01""5=1/  J2 ]]_HH MM%(mmG,mmG,==&!]]>:' __s   4A6H33
I)r   r   r   r"   r    N)gMbP?)g?g+?g:0yE>r   FFF)__name__
__module____qualname____firstlineno__listr   floattupleboolr%   r   r*   __static_attributes__ r'   r$   r   r      s    
 %1!(-'.V'. '. UE\"	'.
 '. '. '. '. "&'.R8d8F#34 8r'   r   )typingr   r   torch.optim._functionaloptim_functionalr6   r   ,torch.distributed.optim._deprecation_warningr   r   rD   r   __annotations__r   scriptr   rI   r'   r$   <module>rQ      sO      # # 
 c  b b br'   