ó
    "Eñiv@  ã            $       ó(  • S SK Jr  S SKrS SKJr  SSKJrJrJrJrJ	r	J
r
JrJrJrJrJrJrJrJrJr  SS/r " S S\5      rS	\ S
\	 S\ S\ S\ S3\l        S\\   S\\   S\\   S\\   S\\   S\\   S\S\S\S\S\S\S\S\S\SS4 S jrS\\   S\\   S\\   S\\   S\\   S\\   S\S\S\S\S\S\S\S\S\SS4 S jr\" \S9     S"S\\   S\\   S\\   S\\   S\\   S\\   S \S-  S\S\S\S\S\S\S\S\S\SS4"S! jj5       rg)#é    )ÚcastN)Ú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Ú
_to_scalarÚ_use_grad_for_differentiableÚ_view_as_realÚ	OptimizerÚParamsTÚASGDÚasgdc                   ó    ^ • \ rS rSr         SS\S\\-  S\S\S\S\S	\S-  S
\S\S\SS4U 4S jjjrU 4S jr	S r
\SS j5       rSrU =r$ )r   é   NÚparamsÚlrÚlambdÚalphaÚt0Úweight_decayÚforeachÚmaximizeÚdifferentiableÚ
capturableÚreturnc           
      óô   >• [        U[        5      (       a  UR                  5       S:w  a  [        S5      eSU::  d  [        SU 35      eSU::  d  [        SU 35      eUUUUUUUU	U
S.	n[        TU ]  X5        g )Nr   zTensor lr must be 1-elementg        zInvalid learning rate: zInvalid weight_decay value: )	r   r   r   r   r   r   r    r!   r"   )Ú
isinstancer   ÚnumelÚ
ValueErrorÚsuperÚ__init__)Úselfr   r   r   r   r   r   r   r    r!   r"   ÚdefaultsÚ	__class__s               €ÚM/home/mande/repo/quber/.venv/lib/python3.13/site-packages/torch/optim/asgd.pyr)   ÚASGD.__init__   s’   ø€ ô �bœ&×!Ñ! b§h¡h£j°A£oÜÐ:Ó;Ð;Ø�b‹yÜÐ6°r°dÐ;Ó<Ð<Ø�lÓ"ÜÐ;¸L¸>ÐJÓKÐKð ØØØØ(ØØ Ø,Ø$ñ

ˆô 	‰Ñ˜Õ*ó    c                 ó6  >• [         TU ]  U5        U R                   GHx  nUR                  SS 5        UR                  SS5        UR                  SS5        UR                  SS5        US    GH"  nU R                  R                  U/ 5      n[        U5      S:w  d  M1  [        R                  " US   5      (       d9  [        US   5      n[        R                  " U[        5       UR                  S	9US'   [        R                  " US
   5      (       d.  [        R                  " US
   [        5       UR                  S	9US
'   [        R                  " US   5      (       a  Mô  [        R                  " US   [        5       UR                  S	9US'   GM%     GM{     g )Nr   r    Fr!   r"   r   r   Ústep)ÚdtypeÚdeviceÚetaÚmu)r(   Ú__setstate__Úparam_groupsÚ
setdefaultÚstateÚgetÚlenÚtorchÚ	is_tensorÚfloatÚtensorr   r3   )r*   r9   ÚgroupÚpÚp_stateÚstep_valr,   s         €r-   r6   ÚASGD.__setstate__?   sE  ø€ Ü‰Ñ˜UÔ#Ø×&Õ&ˆEØ×Ñ˜Y¨Ô-Ø×Ñ˜Z¨Ô/Ø×ÑÐ-¨uÔ5Ø×Ñ˜\¨5Ô1Ø˜8•_�ØŸ*™*Ÿ.™.¨¨BÓ/�Ü�w“< 1Õ$Ü Ÿ?š?¨7°6©?×;Ñ;Ü#(¨°©Ó#9˜Ü*/¯,ª,Ø$Ô,=Ó,?ÈÏÉñ+˜ ™ô !Ÿ?š?¨7°5©>×:Ñ:Ü).¯ªØ# E™NÔ2CÓ2EÈaÏhÉhñ*˜ ™ô !Ÿ?š?¨7°4©=×9Ó9Ü(-¯ªØ# D™MÔ1BÓ1DÈQÏXÉXñ)˜ œô %ò 'r/   c                 ó”  • SnUS    GH»  n	U	R                   c  M  U[        R                  " U	5      -  nUR                  U	5        U	R                   R                  (       a  [        S5      eUR                  U	R                   5        U R                  U	   n
[        U
5      S:X  aÏ  [        R                  " SU	R                  [        5       S9U
S'   [        R                  " [        US   5      U	R                  [        5       S9R                  5       R                  5       U
S	'   [        R                  " SU	R                  [        5       S9U
S
'   [        R                   " U	[        R"                  S9U
S'   UR                  U
S
   5        UR                  U
S   5        UR                  U
S	   5        UR                  U
S   5        GM¾     U$ )NFr   z&ASGD does not support sparse gradientsr   © )r3   r2   r1   r   r4   r5   )Úmemory_formatÚax)Úgradr<   Ú
is_complexÚappendÚ	is_sparseÚRuntimeErrorr9   r;   Úzerosr3   r   Ú	as_tensorr   ÚcloneÚdetachÚonesÚ
zeros_likeÚpreserve_format)r*   r@   Úparams_with_gradÚgradsÚmusÚaxsÚetasÚstate_stepsÚhas_complexrA   r9   s              r-   Ú_init_groupÚASGD._init_groupW   sv  € ØˆØ�x•ˆAØ�v‰vÓ!Øœu×/Ò/°Ó2Ñ2�Ø ×'Ñ'¨Ô*Ø—6‘6×#×#Ü&Ð'OÓPÐPØ—‘˜QŸV™VÔ$àŸ
™
 1™�ä�u“: “?Ü$)§K¢KØ 1§8¡8Ô3DÓ3Fñ%�E˜&‘Mô ŸšÜ& u¨T¡{Ó3Ø#$§8¡8Ü"3Ó"5ñ÷
 ™›ß™›ð ˜%‘Lô #(§*¢*Ø 1§8¡8Ô3DÓ3Fñ#�E˜$‘Kô #(×"2Ò"2Ø¬×)>Ñ)>ñ#�E˜$‘Kð —
‘
˜5 ™;Ô'Ø—
‘
˜5 ™;Ô'Ø—‘˜E %™LÔ)Ø×"Ñ" 5¨¡=×1ñC !ðD Ðr/   c                 óh  • U R                  5         SnUb%  [        R                  " 5          U" 5       nSSS5        U R                   HV  n/ n/ n/ n/ n/ n/ n	U R	                  X4XVXxU	5      n
[        UUUUUU	US   US   US   US   US   US   US   US	   US
   U
S9  MX     U$ ! , (       d  f       Nv= f)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   r   r   r   r   r   r    r!   r"   r[   )Ú'_accelerator_graph_capture_health_checkr<   Úenable_gradr7   r\   r   )r*   ÚclosureÚlossr@   rU   rV   rW   rX   rY   rZ   r[   s              r-   r1   Ú	ASGD.step}   sò   € ð 	×4Ñ4Ô6àˆØÑÜ×"Ò"Õ$Ù“y�÷ %ð ×&Ô&ˆEØ-/ÐØ"$ˆEØ "ˆCØ "ˆCØ!#ˆDØ(*ˆKà×*Ñ*Ø¨°SÀóˆKô Ø ØØØØØØ˜G‘nØ˜‘;Ø˜‘;Ø˜G‘nØ" >Ñ2Ø˜iÑ(Ø˜zÑ*Ø$Ð%5Ñ6Ø  Ñ.Ø'ô!ñ 'ð> ˆ÷E %Õ$ús   «B#Â#
B1rF   )	g{®Gáz„?g-Cëâ6?g      è?g    €„.Ar   NFFF©N)Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__r   r>   r   Úboolr)   r6   r\   r   r1   Ú__static_attributes__Ú__classcell__)r,   s   @r-   r   r      s¾   ø† ð "ØØØØØ#ØØ$Ø ñ+àð+ð �F‰Nð+ð ð	+ð
 ð+ð ð+ð ð+ð ˜‘ð+ð ð+ð ð+ð ð+ð 
÷+ð +õBò0$ðL "ó-ó "ö-r/   zŸImplements Averaged Stochastic Gradient Descent.

    It has been proposed in `Acceleration of stochastic approximation by
    averaging`_.

    Args:
        am  
        lr (float, Tensor, optional): learning rate (default: 1e-2)
        lambd (float, optional): decay term (default: 1e-4)
        alpha (float, optional): power for eta update (default: 0.75)
        t0 (float, optional): point at which to start averaging (default: 1e6)
        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
        z	
        z¿

    .. _Acceleration of stochastic approximation by averaging:
        https://meyn.ece.ufl.edu/wp-content/uploads/sites/77/archive/spm_files/Courses/ECE555-2011/555media/poljud92.pdf

    r   rV   rX   rW   rY   rZ   r   r   r   r   r   r    r!   r"   r[   r#   c       	   
      ó„  • [         R                  R                  5       (       d  [        U5      n[	        U 5       GH  u  nnX   nU(       d  UOU* nX?   nX/   nXO   nX_   n[         R
                  R                  5       (       d¡  U(       aš  [        5       nUR                  R                  UR                  R                  s=:X  a5  UR                  R                  s=:X  a  UR                  R                  :X  a  O  OUR                  R                  U;   d  [        SU S35      e[         R                  " U5      (       aB  [         R                  " U5      n[         R                  " U5      n[         R                  " U5      nUS-  nU
S:w  a  UR                  UU
S9nU(       a)  UR                  SUU-  -
  5        UR                  UUSS9  O3[!        U5      nUR                  SUU-  -
  5        UR#                  UU* S9  U(       d  UR%                  5       S:w  a0  UR#                  UR'                  U5      R                  U5      5        OUR)                  U5        U(       ab  UR)                  USXg-  U-  -   U	-  -  5        UR)                  S[         R*                  " UU-
  [         R,                  " U5      5      -  5        GM‹  [!        U5      n[         R.                  " USXg-  U-  -   U	-  -  5      nUR)                  U5        [         R.                  " S[1        SUU-
  5      -  5      nUR)                  U5        GM     g )NúUIf capturable=True, params, mus, etas, and state_steps must be on supported devices: Ú.r   r   ©r   éÿÿÿÿ©Úvalue)r<   ÚjitÚis_scriptingr   Ú	enumerateÚcompilerÚis_compilingr   r3   ÚtypeÚAssertionErrorrJ   Úview_as_realÚaddÚmul_Úaddcmul_r   Úadd_ÚitemÚsubÚcopy_ÚmaximumÚ	ones_likerO   Úmax)r   rV   rX   rW   rY   rZ   r   r   r   r   r   r    r!   r"   r[   ÚiÚparamrI   r5   rH   r4   Ústep_tÚcapturable_supported_devicesÚ	eta_valuer1   Únew_etaÚnew_mus                              r-   Ú_single_tensor_asgdrŒ   Å   s§  € ô$ �9‰9×!Ñ!×#Ñ#Ü˜‹^ˆä˜f×%‰ˆˆ5Ø‰xˆÞ#‰t¨$¨ˆØ‰VˆØ‰VˆØ‰gˆØ‘ˆô �~‰~×*Ñ*×,Ñ,¶Ü+LÓ+NÐ(à—‘×!Ñ!Ø—9‘9—>‘>õ&à—:‘:—?‘?õ&ð —=‘=×%Ñ%ö&ð —L‘L×%Ñ%Ð)EÓEä$ð-Ø-IÐ,JÈ!ðMóð ô
 ×Ò˜E×"Ñ"Ü×%Ò% dÓ+ˆDÜ×&Ò& uÓ-ˆEÜ×#Ò# BÓ'ˆBð 	�!‰ˆà˜1ÓØ—8‘8˜E¨�8Ð6ˆDæØ�J‰J�q˜5 3™;‘Ô'Ø�N‰N˜4 ¨BˆNÒ/ä" 3›ˆIØ�J‰J�q˜5 9Ñ,Ñ,Ô-Ø�J‰J�t I :ˆJÑ.ö ˜Ÿ™› a›Ø�G‰G�E—I‘I˜b“M×&Ñ& rÓ*Õ+à�H‰H�UŒOæà�I‰I�b˜Q ¡¨fÑ!4Ñ4¸Ñ>Ñ?Ô@Ø�H‰H�QœŸš v°¡{´E·O²OÀFÓ4KÓLÑL×Mä˜fÓ%ˆDÜ—o’o b¨Q°±¸dÑ1BÑ-BÀuÑ,LÑ&MÓNˆGØ�I‰I�gÔÜ—_’_ Q¬¨Q°°r±	Ó):Ñ%:Ó;ˆFØ�H‰H�V×òs &r/   c       	         óÔ  ^"• [        U 5      S:X  a  g U(       a  [        S5      e[        R                  R	                  5       (       dB  U(       a;  [        SS9m"[        U"4S j[        XXESS9 5       5      (       d  [        ST" S	35      e[        U5      n[        R                  " XX#XE/5      nUR                  5        GH�  u  u  nnu  u  nnnnnnn[        [        [           U5      n[        [        [           U5      n[        [        [           U5      n[        [        [           U5      n[        [        [           U5      n[        [        [           U5      nU(       a  [        UUU5        U(       a  [        R                   " U5      n[        R                  R	                  5       (       d>  US   R"                  (       a*  [        R$                  " U[        R&                  " S
SS9S
S9  O[        R$                  " US5        U
S:w  aM  U(       a  [        R$                  " UUU
S9  UnO[        R(                  " UUU
S9n[        R$                  " UUUS9  O[        R(                  " UUUS9n[        R*                  " UUUSS9  A[        R,                  " UU5      n[        R*                  " UUU5        AU(       aÿ  [        R,                  " UU5      n[        R.                  " US
5        [        R0                  " U5        [        R2                  " UU5        A[        R4                  " UU5      n [        R6                  " U U5        [        R$                  " U S5        [        R8                  " U U	5        [        R0                  " U 5        [        R6                  " U U5        [        R2                  " UU 5        GMþ  U V!s/ s H&  n![        R:                  " USXg-  U!-  -   U	-  -  US9PM(     n n!U V!s/ s H1  n![        R:                  " S[=        S[?        U!5      U-
  5      -  US9PM3     nn![        R2                  " UU 5        [        R2                  " UU5        GM      g s  sn!f s  sn!f )Nr   z#_foreach ops don't support autogradF)Úsupports_xlac              3   óN  >#   • U  Hš  u  pp4UR                   R                  UR                   R                  s=:H  =(       a:    UR                   R                  s=:H  =(       a    UR                   R                  :H  Os  =(       a    UR                   R                  T;   v •  Mœ     g 7frd   )r3   rx   )Ú.0rA   r5   r4   r1   rˆ   s        €r-   Ú	<genexpr>Ú%_multi_tensor_asgd.<locals>.<genexpr>3  sr   øé € ð 
ò %UÑ ��sð �H‰H�M‰M˜RŸY™YŸ^™^×RÓR¨s¯z©z¯©×RÓRÀ$Ç+Á+×BRÑBRÔR÷ >Ø—‘—‘Ð!=Ñ=ô>â$Tùs   ƒB"B%T)Ústrictrm   rn   g      ð?Úcpu)r3   ro   r   rp   rq   ) r;   ry   r<   rv   rw   r   ÚallÚzipr   r   Ú"_group_tensors_by_device_and_dtypeÚitemsr   Úlistr   r   Ú_foreach_negÚis_cpuÚ_foreach_add_r?   Ú_foreach_addÚ_foreach_addcmul_Ú_foreach_subÚ_foreach_maximum_Ú_foreach_reciprocal_Ú_foreach_copy_Ú_foreach_mulÚ_foreach_mul_Ú_foreach_pow_rO   r„   r   )#r   rV   rX   rW   rY   rZ   r   r   r   r   r   r    r!   r"   r[   Úgrouped_tensorsr3   Ú_Úgrouped_params_Úgrouped_grads_Úgrouped_axs_Úgrouped_mus_Úgrouped_etas_Úgrouped_state_steps_Úgrouped_paramsÚgrouped_gradsÚgrouped_axsÚgrouped_musÚgrouped_etasÚgrouped_state_stepsÚintermediateÚnew_musÚnew_etasr1   rˆ   s#                                     @r-   Ú_multi_tensor_asgdr·     sç  ø€ ô$ ˆ6ƒ{�aÓØæÜÐBÓCÐCô �>‰>×&Ñ&×(Ñ(®ZÜ'HØñ(
Ð$ô ô 
ô %(¨°TÈtÒ$Tó
÷ 
ñ 
ô
 !ð&Ø&BÐ%CÀ1ðFóð ô
 
�B‹€Bä×BÒBØ	˜ $Ð4ó€Oð 
×	Ñ	×	 ñ
	‰ˆ�ñ 
ñ	
ØØØØØØ à	äœd¤6™l¨OÓ<ˆÜœT¤&™\¨>Ó:ˆÜœ4¤™<¨Ó6ˆÜœ4¤™<¨Ó6ˆÜœD¤™L¨-Ó8ˆÜ"¤4¬¡<Ð1EÓFÐæÜ˜.¨-¸ÔEæÜ!×.Ò.¨}Ó=ˆMô �~‰~×*Ñ*×,Ñ,Ð1DÀQÑ1G×1N×1NÜ×ÒØ#¤U§\¢\°#¸eÑ%DÈCóô ×ÒÐ 3°QÔ7ð ˜1ÓÞÜ×#Ò# M°>ÈÒVØ,‘ä$×1Ò1Ø! >¸ñ �ô ×Ò ¨nÀEÓJä ×-Ò-Ø˜~°UñˆLô 	×Ò °¸lÐRTÒUØô ×)Ò)¨.¸+ÓFˆÜ×Ò ¨\¸;ÔGØö ä×(Ò(Ð)<¸bÓAˆGÜ×#Ò# G¨SÔ1Ü×&Ò& wÔ/Ü× Ò  ¨gÔ6Øô ×)Ò)Ð*=¸uÓEˆHÜ×Ò ¨"Ô-Ü×Ò ¨!Ô,Ü×Ò ¨%Ô0Ü×&Ò& xÔ0Ü×Ò ¨"Ô-Ü× Ò  ¨x×8ñ 0óâ/�Dô —’  q¨5©:¸Ñ+<Ñ'<ÀÑ&FÑ GÐPVÔWÙ/ð ð ñ 0óâ/�Dô —’ ¤C¨¬:°dÓ+;¸bÑ+@Ó$AÑ AÈ&ÔQÙ/ð ð ô × Ò  ¨xÔ8Ü× Ò  ¨g×6òs 
!ùò`ùòs   Î?-Q Ï28Q%)Úsingle_tensor_fnr   c                ó$  • Uc  [        XSS9u  nnU(       a.  [        R                  R                  5       (       a  [	        S5      eU(       a*  [        R                  R                  5       (       d  [
        nO[        nU" U UUUUUUUUUUUUU	U
S9  g)zfFunctional API that performs asgd algorithm computation.

See :class:`~torch.optim.ASGD` for details.
NF)Ú	use_fusedz6torch.jit.script not supported with foreach optimizers)	r   r   r   r   r   r    r!   r"   r[   )r   r<   rs   rt   rM   r·   rŒ   )r   rV   rX   rW   rY   rZ   r   r    r!   r"   r[   r   r   r   r   r   r§   Úfuncs                     r-   r   r   ¨  s“   € ð4 �Ü1Ø¨eñ
‰
ˆˆ7ö ”5—9‘9×)Ñ)×+Ñ+ÜÐSÓTÐTæ”u—y‘y×-Ñ-×/Ñ/Ü!‰ä"ˆáØØØØØØØØØØØ!ØØ%ØØór/   )NFFFF)Útypingr   r<   r   Ú	optimizerr   r   r   r	   r
   r   r   r   r   r   r   r   r   r   r   Ú__all__r   Ú__doc__r™   r>   ri   rŒ   r·   r   rF   r/   r-   Ú<module>rÀ      s  ðå ã Ý ÷÷ ÷ ÷ ñ ð& �6Ð
€ôNˆ9ô Nðb	ð 
ˆð 	ð 
ˆð 	Ø	ˆð 	Ø	Ðð 	Ø	Ðð ð€„ð.NØ�‰LðNà�‰<ðNð 
ˆf‰ðNð 
ˆf‰ð	Nð
 ˆv‰,ðNð �f‘ðNð ðNð 	ðNð 	ðNð ðNð ðNð ðNð ðNð ðNð  ð!Nð" 
ô#NðbO7Ø�‰LðO7à�‰<ðO7ð 
ˆf‰ðO7ð 
ˆf‰ð	O7ð
 ˆv‰,ðO7ð �f‘ðO7ð ðO7ð 	ðO7ð 	ðO7ð ðO7ð ðO7ð ðO7ð ðO7ð ðO7ð  ð!O7ð" 
ô#O7ñd  Ð1DÑEð  ØØ ØØñ6Ø�‰Lð6à�‰<ð6ð 
ˆf‰ð6ð 
ˆf‰ð	6ð
 ˆv‰,ð6ð �f‘ð6ð �D‰[ð6ð ð6ð ð6ð ð6ð ð6ð ð6ð  	ð!6ð" 	ð#6ð$ ð%6ð& ð'6ð( 
ô)6ó Fñ6r/   