
    [Thr                        % S SK r S SKJrJr  S SKJr   " S S5      r\ R                  " 5       q\ R                  \R                  \ R                  \   4   \S'   S\R                  S\SS4S	 jrS\R                  S\\   4S
 jrg)    N)castOptionalc                       \ rS rSrSrg)_State    N)__name__
__module____qualname____firstlineno____static_attributes__r       [/var/www/auris/envauris/lib/python3.13/site-packages/torch/distributed/_composable_state.pyr   r      s    r   r   _module_state_mappingmodulestatereturnc                 h    U [         ;  d   SU  S35       e[        R                  " U5      [         U '   g )Nz
Inserting z more than once.)r   weakrefref)r   r   s     r   _insert_module_stater      s4    ..U*VHDT0UU.$+KK$6&!r   c                     [        U [        5      (       a  [        [        U 5      $ U [        ;   a   [        U    nU" 5       nUc  [	        S5      eU$ g)aX  
Return the ``_State`` in ``model``.

Given a ``module``, this API finds out if the module is also a ``_State``
instance or if the module is managed by a composable API. If the module
is also a ``_State``, ``module`` will be casted to ``_State` and returned.
If it is managed by a composable API, the corresponding ``_State`` will
be returned.
Nz(State has already been garbage collected)
isinstancer   r   r   AssertionError)r   	state_refr   s      r   _get_module_stater      sS     &&!!FF## **-f5IKE}$%OPPLr   )r   typingr   r   torch.nnnnr   WeakKeyDictionaryr   ModuleReferenceType__annotations__r   r   r   r   r   <module>r$      s     ! 	 	  w00IIw$$V,,  
7 76 7d 7bii HV,< r   