Ë
    P/æjc  ã                   óæ   — d dl Z d dlZd dlmZ d dlZd dlmc mZ 	 d dl	Z
defd„Zd„ Zd„ Zd„ Z G d„ d	«      Zd
„ Zdej&                  fd„Z G d„ d«      Zdefd„Zdedefd„Zy# e$ r dZ
Y ŒNw xY w)é    N)ÚListÚreturnc                  ó>   — t         j                  j                  «       S )a£  
    Indicates whether we are tracing/compiling with torch.compile() or torch.export().

    If need to check specifically that TorchDynamo is used, then use
    torch.compiler.is_dynamo_compiling().

    TODO(khabinov): we should deprecate this function and use one of these two:
    * torch.compiler.is_compiling(),
    * torch.compiler.is_dynamo_compiling().
    It will depend on the context where to use what.
    )ÚtorchÚcompilerÚis_compiling© ó    ú_/root/repos/pocs/video-evals/.venv/lib/python3.12/site-packages/torch/_dynamo/external_utils.pyr   r      s   € ô �>‰>×&Ñ&Ó(Ð(r
   c                 óB   ‡ — t        j                  ‰ «      ˆ fd„«       }|S )zB
    Create an extra frame around fn that is not in skipfiles
    c                  ó   •—  ‰| i |¤ŽS ©Nr	   )ÚargsÚkwargsÚfns     €r   Úinnerzwrap_inline.<locals>.inner&   s   ø€ á�4Ð"˜6Ñ"Ð"r
   )Ú	functoolsÚwraps)r   r   s   ` r   Úwrap_inliner   !   s'   ø€ ô
 ‡_�_�RÓó#ó ð#ð €Lr
   c                 óD   —  | |Ž }|€|d   S |d   dk(  rt        d«      ‚|S )zA
    Used by compiled autograd to handle hook returning None
    r   Ú	hook_typeÚpost_acc_grad_hookz5Tensor post accumulate grad hooks should return None.)ÚRuntimeError)Úhookr   r   Úresults       r   Ú	call_hookr   -   s:   € ñ �4ˆ[€FØ€~Ø�A‰wˆØ	�Ñ	Ð 4Ò	4ÜÐRÓSÐSØ€Mr
   c                 óR   ‡ — t         s‰ S t        j                  ‰ «      ˆ fd„«       }|S )z�Decorator that turns a function from ``np.ndarray``s to ``np.ndarray``s into a function
    from ``torch.Tensor``s to ``torch.Tensor``s.
    c                  ó¶   •— t        j                  t        j                  d„ | |f«      \  } } ‰| i |¤Ž}t        j                  t        j
                  d„ |«      S )Nc                 ó"   — | j                  «       S r   )Únumpy©Úxs    r   ú<lambda>z*wrap_numpy.<locals>.wrap.<locals>.<lambda>C   s   €  A§G¡G£I€ r
   c                 ó,   — t        j                  | «      S r   )r   Ú	as_tensorr!   s    r   r#   z*wrap_numpy.<locals>.wrap.<locals>.<lambda>F   s   € ¼%¿/¹/È!Ó:L€ r
   )ÚpytreeÚtree_map_onlyr   ÚTensorÚnpÚndarray)r   r   ÚoutÚfs      €r   Úwrapzwrap_numpy.<locals>.wrap@   sR   ø€ ä×+Ñ+Ü�L‰LÑ-°°f¨~ó
‰ˆˆfñ �Ð ˜Ñ ˆÜ×#Ñ#¤B§J¡JÑ0LÈcÓRÐRr
   )r)   r   r   )r,   r-   s   ` r   Ú
wrap_numpyr.   9   s1   ø€ õ Øˆä‡_�_�QÓóSó ðSð €Kr
   c                   ó|   — e Zd Zdej                  j
                  j                  deej                     ddfd„Z	d„ Z
y)ÚFakeBackwardCFunctionÚrealÚsaved_tensorsr   Nc                 ó    — || _         || _        y r   )r1   r2   )Úselfr1   r2   s      r   Ú__init__zFakeBackwardCFunction.__init__L   s   € ð
 ˆŒ	Ø*ˆÕr
   c                 ó„   — |dk(  r&t        j                  dt        «       | j                  S t	        | j
                  |«      S )NÚsaved_variablesz4'saved_variables' is deprecated; use 'saved_tensors')ÚwarningsÚwarnÚDeprecationWarningr2   Úgetattrr1   )r4   Únames     r   Ú__getattr__z!FakeBackwardCFunction.__getattr__T   s>   € ØÐ$Ò$Ü�M‰MØFÜ"ôð ×%Ñ%Ð%ô �t—y‘y $Ó'Ð'r
   )Ú__name__Ú
__module__Ú__qualname__r   ÚautogradÚfunctionÚBackwardCFunctionr   r(   r5   r=   r	   r
   r   r0   r0   K   sA   „ ð+à�n‰n×%Ñ%×7Ñ7ð+ð ˜EŸL™LÑ)ð+ð 
ó	+ó	(r
   r0   c                 ó€   — t        | |«      } |j                  j                  |g|¢­Ž }t        |«      t        ur|f}|S r   )r0   Ú_forward_clsÚbackwardÚtypeÚtuple)Úbackward_c_functionr2   r   ÚfakeÚgradss        r   Úcall_backwardrL   a   sF   € Ü Ð!4°mÓD€DØ&ˆD×Ñ×&Ñ& tÐ3¨dÒ3€Eô ˆEƒ{œ%ÑØ�ˆà€Lr
   r"   c                 ó>   — | j                  «       j                  «       S r   )Úuntyped_storageÚsizer!   s    r   Úuntyped_storage_sizerP   l   s   € Ø×ÑÓ×#Ñ#Ó%Ð%r
   c                   ó<   — e Zd Zed„ «       Zed„ «       Zed„ «       Zy)ÚFakeCompiledAutogradEnginec                 ó&   — | j                  |«       y r   )Úappend)Úfinal_callbacksÚcbs     r   Úqueue_callbackz)FakeCompiledAutogradEngine.queue_callbackq   s   € à×Ñ˜rÕ"r
   c                 ó„   — d}|t        | «      k  r | |   } |«        |dz  }|t        | «      k  rŒ | j                  «        y )Nr   é   )ÚlenÚclear)rU   ÚirV   s      r   Úexec_final_callbacksz/FakeCompiledAutogradEngine.exec_final_callbacksu   sI   € àˆØ”#�oÓ&Ò&Ø  Ñ#ˆBÙŒDØ�‰FˆAð ”#�oÓ&Ó&ð 	×ÑÕr
   c                   ó   — y r   r	   r	   r
   r   Ú_exec_final_callbacks_stubz5FakeCompiledAutogradEngine._exec_final_callbacks_stub~   s   € àr
   N)r>   r?   r@   ÚstaticmethodrW   r]   r_   r	   r
   r   rR   rR   p   s<   „ Øñ#ó ð#ð ñ ó ð ð ñó ñr
   rR   Ú	hook_namec                 ó&   —  t        | |«      |i |¤ŽS r   ©r;   )Úbw_statera   r   r   s       r   Úcall_hook_from_backward_statere   ƒ   s   € Ø'Œ7�8˜YÓ'¨Ð8°Ñ8Ð8r
   Ú
hooks_nameÚmodule_namec                ób   — t        ||«      }t        ||«      }|D ]  } |||g|¢­Ž }	|	€Œ|	}Œ |S r   rc   )
Ú_r   rd   rf   rg   r   ÚmoduleÚhooksr   Ú
new_results
             r   Ú%call_module_hooks_from_backward_staterm   ‡   sN   € ô �X˜{Ó+€FÜ�H˜jÓ)€EØò  ˆÙ˜& &Ð0¨4Ò0ˆ
ØÑ!Ø‰Fð ð €Mr
   )r   r8   Útypingr   r   Útorch.utils._pytreeÚutilsÚ_pytreer&   r    r)   ÚModuleNotFoundErrorÚboolr   r   r   r.   r0   rL   r(   rP   rR   Ústrre   rm   r	   r
   r   ú<module>ru      s�   ðó Û Ý ã ß $Ð $ðÛð
)�dó )ò	ò	ò÷$(ñ (ò,ð&˜EŸL™Ló &÷ñ ð&9¸có 9ð	Ø,/ð	Ø>Aô	øðs ò Ø	‚Bðús   �A& Á&A0Á/A0