
    [The"                     8   S SK r S SK Jr  S SKJrJrJrJr  S SKJrJ	r	J
r
Jr  S SKJr  SSKJr   " S S	\5      r\ R                   " S
 S5      5       rS\	R$                  S\S\4S jrS\S\4S jrS rS r\R.                  r\R2                  r\R6                  rg)    N)	dataclass)AnyCallableOptionalProtocol)_C_opsautogradTensor)_pytree   )utilsc                   6    \ rS rSr% \\   \S'   \\   \S'   Srg)InfoProtocol   _backward_fn_setup_context_fn N__name__
__module____qualname____firstlineno__r   r   __annotations____static_attributes__r       O/var/www/auris/envauris/lib/python3.13/site-packages/torch/_library/autograd.pyr   r      s    8$$))r   r   c                   6    \ rS rSr% \\   \S'   \\   \S'   Srg)Info   r   r   r   Nr   r   r   r   r   r      s    8$$))r   r   opinforeturnc                   ^ ^^^^	^
 ST R                    ST R                   ST R                   3n[        R                  " T R
                  5      m
[         " S S5      5       mU 4S jm	U
UU 4S jnUU 4S jn[        U[        R                  4[        U5      [        U5      S.5      mT R
                  n[        S	 / UR                  QUR                  Q7 5       5      (       a  [        T5      mUUU	4S
 jnU$ )NGeneratedBackwardFor__c                   H    \ rS rSr% \R
                  \S'   \\\	4   \S'   Sr
g)$make_autograd_impl.<locals>.Metadata   keysetkeyword_only_argsr   N)r   r   r   r   r   DispatchKeySetr   dictstrr   r   r   r   r   Metadatar(      s    !!!S>)r   r/   c                     > U S   nU S S n [         R                  " 5          UR                  nUR                  nTR                  " U[         R
                  -  /U Q70 UD6nUsS S S 5        $ ! , (       d  f       g = f)N)r   _AutoDispatchBelowAutogradr*   r+   
redispatch_after_autograd_keyset)argsmetadatar*   kwargsresultr!   s        r   forward_no_grad+make_autograd_impl.<locals>.forward_no_grad!   sj    8CRy**,__F//F]]6B,E,E#EWWPVWF	 -,,s   AA,,
A:c                   > US   nUS S n[         R                  " 5          UR                  nUR                  nTR                  " U[         R
                  -  /UQ70 UD6nTR                  (       aK  [        R                  " TR                  X5      u  pT(       a  TR                  XXES9  OTR                  XUS9  UsS S S 5        $ ! , (       d  f       g = f)Nr1   )ctxinputskeyword_only_inputsoutput)r<   r=   r?   )
r   r2   r*   r+   r3   r4   r   r   fill_defaults_schema)	r<   r5   r6   r*   r7   r8   has_kwarg_only_argsr"   r!   s	         r   forward#make_autograd_impl.<locals>.forward+   s    8CRy**,__F//F]]6B,E,E#EWWPVWF%%  %222::tL&**& +  **s*O/ -,,s   BC
Cc                   > TR                   (       aZ   U R                  nU R                  S S U l        TR                   " U /UQ76 nX l        [        U[        5      (       a  / UQS P7$ US 4$ [	        ST S35      e! WU l        f = f)Nr1   zTrying to backward through zQ but no autograd formula was registered. Please use register_autograd to add one.)r   needs_input_grad
isinstancetupleRuntimeError)r<   gradsprev_needs_input_gradr8   r"   r!   s       r   backward$make_autograd_impl.<locals>.backwardH   s    =(+(<(<%'*';';CR'@$**377'<$&%((&&4<)" .7 8
 	
	 (=$s   3A< <	B)rC   rL   c              3   b   #    U  H%  n[         R                  " UR                  5      v   M'     g 7fN)r   is_tensorlist_like_typetype).0as     r   	<genexpr>%make_autograd_impl.<locals>.<genexpr>c   s)      5A 	%%aff--5s   -/c                    > [         R                  " 5       (       aC  [        R                  " [        S U[
        5      (       a  TR                  " / UQT" X5      P76 nU$ T" / UQT" X5      P76 nU$ )Nc                     U R                   $ rO   )requires_grad)xs    r   <lambda>;make_autograd_impl.<locals>.autograd_impl.<locals>.<lambda>m   s    aoor   )r   is_grad_enabledr   tree_any_onlyr   not_list_of_tensorapply)r*   r5   r+   r8   	Generatedr/   r9   s       r   autograd_impl)make_autograd_impl.<locals>.autograd_implk   sq    G$9$9-t5G%
 %
 __PdPHV,OPF  %PdPHV,OPFr   )
_namespace_opname_overloadnamer   rB   rA   r   rQ   r
   Functionstaticmethodany	argumentsreturnssupports_tensorlist)r!   r"   namerC   rL   schemara   r`   r/   r9   rB   s   ``     @@@@r   make_autograd_implrn      s    'a

|1REUEUDVWD33BJJ?* * *:
" 			#G,$X.	
I ZZF
 56##5fnn5   (	2	 r   clsc                    ^^^^ U R                   mU R                  mU R                  m[         " S S5      5       mUU4S jnU4S jnUU4S jnXl         X l        X0l        U $ )a$  Allows a given autograd.Function class to support List[Tensor] inputs/outputs.

Regular autograd.Function has a constraint that it only directly supports autograd for
Tensors. Applying @supports_tensorlist enables an autograd.Function to support
autograd for List[Tensor] inputs and outputs.
c                   H    \ rS rSr% \\S'   Sr\\   \S'   Sr\\	   \S'   Sr
g)%supports_tensorlist.<locals>.Metadata   
input_specNoutput_specresult_is_tupler   )r   r   r   r   spec_tr   ru   r   rv   boolr   r   r   r   r/   rr      s%    (,Xf%,*.$.r   r/   c                 |  > US   nUS S n[        UT5      (       d  [        S5      e[        [        U5      UR                  5      nT" U /UQ76 n[        U[
        5      Ul        UR                  (       d  U4n[        U[        5      u  pEXRl	        [        U S5      (       a  [        S5      eX l        [        U5      $ )Nr1   zNYI: calling supports_tensorlist autograd.Function.forward directly. You should probably be calling .apply instead. Please file an issue if not._pt_metadataz@Please don't set ctx._pt_metadata; PyTorch uses it to store info)rG   NotImplementedError	unflattenlistrt   rH   rv   flattenr^   ru   hasattrrI   rz   )r<   r5   r6   r8   flat_resultru   r/   orig_forwards         r   new_forward(supports_tensorlist.<locals>.new_forward   s    8CRy(H--%/ 
 dX%8%89c)D)#-fe#< ''YF#*63E#F *3''R  $[!!r   c                 
  > [        U S5      (       d  [        S5      eU R                  n[        [	        U5      UR
                  5      nU R                  n [        [	        U R                  S S 5      UR                  5      U l        T" U /UQ76 nX0l        [        U[        5      (       d  U4n[        U[        5      u  pVXbR                  :w  a  [        SU SUR                   S35      e[        US /-   5      $ ! X0l        f = f)Nrz   zNYI: calling supports_tensorlist autograd.Function.backward directly. This will automatically get called by PyTorch autograd. Please file an issue if you need this.r1   zRExpected the return from backward to be of the same structure as the inputs. Got: z (return from backward), z	 (inputs))r   r{   rz   r|   r}   ru   rF   rt   rG   rH   r~   not_list_of_optional_tensorrI   )r<   rJ   r6   rK   grad_inputsflat_grad_inputsgrad_inputs_specorig_backwards          r   new_backward)supports_tensorlist.<locals>.new_backward   s   sN++%9  ##$u+x';';< !$ 4 4	9#,S))#2./1D1D$C  (4e4K#8 +u--&.K .54.
* 222''7&88Q&&'y2 
 %.//% $9 s   :C: :Dc                  $  > [        U [        S9u  pT" U5      nT" / UQUP76 nUR                  c   e[        [	        U5      UR                  5      nUR
                  (       d-  [        U[        5      (       d   e[        U5      S:X  d   eUS   $ U$ )N)is_leafr   r   )	r~   r^   ru   r|   r}   rv   rG   rH   len)r5   	flat_argsrt   r6   r8   r/   
orig_applys        r   	new_apply&supports_tensorlist.<locals>.new_apply   s     '6H I	J'1Y11##///4<)=)=>''fe,,,,v;!###!9r   )rC   rL   r_   r   )ro   r   r   r   r/   r   r   r   s       @@@@r   rk   rk   w   s]     ;;LLLMJ/ / /
"2*0X
 KLIJr   c                 ~    [        U [        5      (       a  g[        U [        5      (       a  [        S U  5       5      $ g)NFc              3   L   #    U  H  n[        U[        5      (       + v   M     g 7frO   rG   r   rR   ls     r   rT   %not_list_of_tensor.<locals>.<genexpr>   s     ;dz!V,,,ds   "$TrG   rH   r}   rh   trees    r   r^   r^      s3    $$;d;;;r   c                 ~    [        U [        5      (       a  g[        U [        5      (       a  [        S U  5       5      $ g)NFc              3   `   #    U  H$  oS L=(       a    [        U[        5      (       + v   M&     g 7frO   r   r   s     r   rT   .not_list_of_optional_tensor.<locals>.<genexpr>   s$     M1D=>Av)>%>>s   ,.Tr   r   s    r   r   r      s3    $$MMMMr   )dataclassesr   typingr   r   r   r   torchr   r	   r
   r   torch.utilsr    r   r   r   
OpOverloadrn   rk   r^   r   tree_flattenr~   tree_unflattenr|   TreeSpecrw   r   r   r   <module>r      s     ! 4 4 , ,  *8 *
 * * *
]4?? ], ]8 ]@eS eS eP 

""				r   