MjUddlZddlmcmZddlmZddlmZgZe e e d<ejjGddZy)N)Tensor)2_scripted_functional_optimizer_deprecation_warning__all__cpeZdZ d deededeeefdeeefdededefdZd eed zfd Z y ) _FunctionalRpropparamslretas step_sizesforeachmaximize_allow_empty_param_listc\tdd|i|_||_||_||_||_t |dk(r |s tdd|i|_tjjttjtttjffi|_y)N) stacklevelr rz%optimizer got an empty parameter listr)rdefaultsr r r r len ValueError param_grouptorchjitannotatedictrstrstate)selfrr r r r r rs g/var/www/html/ia-enhanced/venv/lib/python3.12/site-packages/torch/distributed/optim/functional_rprop.py__init__z_FunctionalRprop.__init__s ;aH "   $   v;! $;DE E%f-YY''U\\4U\\@Q;R-R(SUWX gradientsNc|jd}g}g}g}g}g}|jd}|j\} } |j\} } t |t |k7r*t ddt |dzdt |zd} t ||D]\}}| | tj|z} |j||j|||jvrwi|j|<|j|}tjd|d <tj|tj |d <tj|||d <|j|}|j|d |j|d |j|d tj5t!j"|||||| | | | |j$|j&| dddy#1swYyxYw)Nrr zEthe gradients passed in does not equal to the size of the parameters!zParams length: z. zGradients length: Fgstep) memory_formatprev step_size) step_size_min step_size_maxetaminusetaplusr r has_complex)rrr r rrzipr is_complexappendrtensor zeros_likepreserve_format full_likeno_gradFrpropr r )rr rparams_with_gradgradsprevsr state_stepsr r(r)r&r'r*paramgradientrs rr"z_FunctionalRprop.step5s!!(+  ]]4  II''+$ } v;#i. (W#CK=34&s9~&678   "695 2OE8#u//66  ''. X& *(*DJJu% JJu-E$)LL$5E&M$)$4$4U-B-B%E&M*/2)FE+& 5) U6]+!!% "45""5=1% 2(]]_  GG ++! '    s 6HH )g{Gz?)g?g333333?)gư>2FFF) __name__ __module__ __qualname__listrfloattupleboolrr"rrrrs $.*4(-YV Y YE5L! Y %,' Y  YY"&Y85d6D=15rr)rtorch.optim._functionaloptim _functionalr3r,torch.distributed.optim._deprecation_warningrrr?r__annotations__rscriptrrCrrrJsN ## cRRRr