
    Phth                        d dl Z d dlZd dlZd dlZd dlZd dlZd dlZd dlZd dlZd dl	Z	d dl
mZmZ d dlZd dlZd dlmZ ddlmZ ddlmZ ddlmZ ddlmZmZ dd	lmZmZmZmZ dd
l m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z) ddl*m+Z+m,Z, ddl-m.Z.m/Z/ ddl0m1Z1  G d de,      Z2 G d de2      Z3 G d de2      Z4 G d de4      Z5 G d de,      Z6y)    N)DictList)TracingContext   )	variables)
is_allowed)unimplemented)GuardBuilderinstall_guard)
AttrSourceGetItemSourceODictGetItemSourceRandomValueSource)	all_hook_namesbuild_checkpoint_variablecheck_constant_argsget_custom_getattris_namedtuple_clsis_utils_checkpointistypenamedtuple_fieldsobject_has_getattribute   )MutableLocalVariableTracker)GenericContextWrappingVariableNullContextVariable)ConstDictVariablec                       e Zd Zy)UserDefinedVariableN__name__
__module____qualname__     oC:\Users\daisl\Desktop\realtime-object-detection\venv\Lib\site-packages\torch/_dynamo/variables/user_defined.pyr    r    '   s    r&   r    c                   x     e Zd Z fdZd Zd Zdeddf fdZ	 	 	 	 	 	 d fdZ	 	 	 	 	 	 d fd	Z	 fd
Z
 xZS )UserDefinedClassVariablec                 2    t        |   di | || _        y Nr%   )super__init__value)selfr.   kwargs	__class__s      r'   r-   z!UserDefinedClassVariable.__init__,   s    "6"
r&   c                     | j                   S N)r.   r/   s    r'   as_python_constantz+UserDefinedClassVariable.as_python_constant0   s    zzr&   c                 ,    t        | j                        S r3   )typer.   r4   s    r'   python_typez$UserDefinedClassVariable.python_type3   s    DJJr&   namereturnr   c                    ddl m} ddlm} | j                  t        | j                  |      nd }	 t        j                  | j                  |      }t        |t              r0t        j                  |j                  | j                        |      S t        |t              r"t        j                   |j"                  | |      S |r=t        j$                  |      r(  |||      |j                  | j                              S |t'        | j                  di       v s|j)                  |      r3|r  |||      |      S |j)                  |      r|j+                  |      S t,        | ]  ||      S # t        $ r d }Y ,w xY w)Nr   ConstantVariableVariableBuildersource__dict__) r=   builderr?   rA   r   inspectgetattr_staticr.   AttributeError
isinstancestaticmethodr   UserFunctionVariable__get__classmethodUserMethodVariable__func__ismemberdescriptorgetattr
is_literalcreater,   var_getattr)r/   txr9   r=   r?   rA   objr1   s          r'   rS   z$UserDefinedClassVariable.var_getattr6   sL   &,26++2IDKK.t	((T:C c<(11DJJ'  [)//d6RR2237.?2v.s{{4::/FGG74::z266:J:U:U;
 2r62377!,,S1'..s33w"2t,,'  	C	s    E1 1F ?F c                    |dk(  rt        |      dk(  r|sd| j                  j                  vrdt               i}t	               }| j                  j                         D ]W  }t        |j                  |j                        |j                        }|j                  t        j                  ||             Y t        j                  |fi |S t        	| =  ||||      S )N__subclasses__r   mutable_localr@   )lenr.   rB   r   listrW   r   import_sourcer#   r"   appendr   r)   ListVariabler,   call_method)
r/   rT   r9   argsr0   optionssubs_as_varssubrA   r1   s
            r'   r^   z$UserDefinedClassVariable.call_methodR   s     $$D	Q 

(;(;;&7G26&Lzz002#B$4$4S^^$DcllS##66s6J 3 )),B'BBw"2tT6::r&   c                 	   ddl m} ddlm} | j                  t
        j                  u r
t               S t        t        | j                        t              r}t        | j                  d      rgt        | j                  d      rQt        ||      rEt        |      dk(  r7|D cg c]  }|j                          }}t        | | j                  |       S t        | j                        rt!        | j                        }| j                  j"                  }	t%        |      }
|
j'                  d gt        |      t        |
      z
  z         i }t)        ||
      D ].  \  }}|	||v r||   }n||	v sJ   |       ||	|         }|||<   0 |j+                         D ]  \  }}||v sJ ||
|j-                  |      <   ! t/        d	 |
D              sJ t1        j2                  |
| j                        S t5        j6                  | j                  d
d       t8        j:                  fv rH|j=                  | j                        r,| j>                  r|j@                  j                   jC                  | j>                  | j                  t        | j                  tD        jF                  jH                        rt0        jJ                  ntL        i       }t5        j6                  | j                  dd       tD        jF                  jH                  jN                  u rF|j@                  j                   jQ                  |dt0        jR                  jU                  d             |S |jW                  |d||       |S t0        jX                  j[                  | j                        r8dt]               i}t0        jX                  jU                  | j                  |||      S t0        j^                  j[                  | j                        r8dt]               i}t0        j^                  jU                  | j                  |||      S t0        j`                  j[                  | j                        rs| j>                  rgt1        j`                  t1        jb                  t$              je                  |||      j*                  | j                  | j>                  t]                     S tf        |   |||      S c c}w )Nr   )SideEffectsr   )SourcelessBuilder	__enter____exit__r   )cm_objc              3   $   K   | ]  }|d u 
 y wr3   r%   ).0xs     r'   	<genexpr>z9UserDefinedClassVariable.call_function.<locals>.<genexpr>   s     4eq}es   __new__r-   __call_nn_module_initTrX   )user_clsuser_cls_sourcerX   )4side_effectsrd   rD   re   r.   
contextlibnullcontextr   
issubclassr7   hasattrr   rY   r5   r   r   r   _field_defaultsrZ   extendzipitemsindexallr   NamedTupleVariablerE   rF   objectrm   "cls_supports_mutation_side_effectsrA   outputtrack_object_newtorchnnModuleUnspecializedNNModuleVariableUserDefinedObjectVariabler-   
store_attrr=   rR   r^   CustomizedDictVariableis_matching_clsr   DataClassVariableRestrictedListSubclassVariableBuiltinVariablecall_functionr,   )r/   rT   r_   r0   rd   re   rk   unwrapped_argsfieldsfield_defaultsry   var_tracker_kwargs
field_namevar_tracker	field_varr9   r.   varr`   r1   s                      r'   r   z&UserDefinedClassVariable.call_functionk   s    	/.::///&((tDJJ'.

K0

J/#D&1Fq >BCda224dNC1!tzz>2  tzz*&tzz2F!ZZ77NJELL$3v;U#;<=!#+.vu+='
K&!V+$*:$6	)^;;;$7$5$7N:<V$W	5>&z2 ,>  2779ev~%~,1fll4()  : 4e4444//tzzBB""4::y$?FNNCTT>>tzzJ))((99

djj%((//: 77.C &&tzz:tD88??++, 		&&11+..55d;
 
Jf=
--==djjI&7G33::

D&'  ((88D&7G..55djj$PWXX44DDTZZP;;))$/==b$OUU $*n	  w$Rv66S Ds   !Sc                 Z    |dk(  r| j                   j                  S t        |   ||      S )Nr"   )r.   r"   r,   const_getattrr/   rT   r9   r1   s      r'   r   z&UserDefinedClassVariable.const_getattr   s-    :::&&&w$R..r&   r_   zList[VariableTracker]r0   zDict[str, VariableTracker]r:   r   )r"   r#   r$   r-   r5   r8   strrS   r^   r   r   __classcell__r1   s   @r'   r)   r)   +   su     -C -,= -8; &	;
 -; 
;2X7/X79UX7	X7t/ /r&   r)   c                        e Zd ZdZddhej
                  Zd fd	Zd Zd Ze	 e
j                  d      d               Zd	 Z	 	 	 	 	 	 d fdZ fdZd Z	 	 	 	 	 	 d fdZd Zd Zd Zd Zded
dfdZd Z xZS )r   z_
    Mostly objects of defined type.  Catch-all for something where we only know the type.
    r.   
value_typeNc                     t        |   di | || _        |xs t        |      | _        t        |      | j                  u sJ y r+   )r,   r-   r.   r7   r   )r/   r.   r   r0   r1   s       r'   r-   z"UserDefinedObjectVariable.__init__   s@    "6"
$3UE{doo---r&   c                     | j                   j                  }|dv r t        t        | j                  dd             }| j
                  j                   d| dS )N)builtin_function_or_methodgetset_descriptormethod_descriptormethodr"   ())r   r"   r   rP   r.   r1   )r/   inners     r'   __str__z!UserDefinedObjectVariable.__str__   sU    (( 
 
 

J=>E..))*!E7!44r&   c                     | j                   S r3   )r   r4   s    r'   r8   z%UserDefinedObjectVariable.python_type   s    r&   c                      t         j                   t         j                  t         j                  t         j                  h} | S r3   )randomrandint	randrangeuniform)fnss    r'   _supported_random_functionsz5UserDefinedObjectVariable._supported_random_functions   s0     MMNNNN	
 
r&   c                     |t        | j                  di       vr*	 t        j                  t	        | j                        |      S y # t
        $ r Y y w xY w)NrB   )rP   r.   rE   rF   r7   rG   )r/   r9   s     r'   _maybe_get_baseclass_methodz5UserDefinedObjectVariable._maybe_get_baseclass_method   sT    wtzz:r::--d4::.>EE  " s   (A 	AAr:   r   c           
         ddl m}m}m}m} | j                  |      }	|	|	t        j                  u r|j                  d       S |	t        j                  j                  u r|s|rJ | j                  sJ t        | j                  j                               }
t        t!        |j"                  |
            sJ t%        | j                  j'                  t(        j*                                ||
D cg c]  }|j                  |       c}      S |	t        j                  j,                  t.        j,                  fv rt1        |      dk(  rt3        |d   ||f      rt5        j6                  t9        | j                        d      t        j                  j                  t.        j                  fv rr|rJ | j                  sJ t%        | j                  j'                  t(        j*                               |j                  |d   j;                         | j                  v       S |	t        j                  j<                  u rt3        | j                  t        j                        rq| j                  sJ |s|rJ g }| j?                  |dg i       }
|
jA                  |      D ]+  }|jC                   ||| jE                  ||      g             -  ||      S |	t        j                  jF                  u r5t1        |      dk(  r'|rJ | j                  sJ | jE                  ||d         S t3        |	tH        jJ                        rJ| j                  d ntM        tM        | j                  d      |      } ||	| |      jO                  |||      S |	t        jP                  u r]| j                  rQ|sO|sMt%        | j                  j'                  t(        jR                                |t1        | j                              S tT        | }  ||||      S c c}w )Nr   )r   r=   TupleVariablerM   r   keysr1   r@   )+rC   r   r=   r   rM   r   r}   r-   rR   collectionsOrderedDictr   rA   rZ   r.   r{   maprQ   r   
make_guardr
   
ODICT_KEYS__contains__dictrY   rH   rE   rF   r7   r5   ry   r^   unpack_var_sequencer\   odict_getitem__getitem__typesFunctionTyper   r   __len__LIST_LENGTHr,   )r/   rT   r9   r_   r0   r   r=   r   rM   r   r   kry   keyrA   r1   s                  r'   r^   z%UserDefinedObjectVariable.call_method   ss   	
 	
 11$7('..t44 00555 F++{{"{DJJOO-.3/::DABBBdkk44\5L5LMN$$%O$Q&6&=&=a&@$%OPP ;22??ARARSSINtAw)9?(KL**4

+;VD++00$))<= "!z{{"{dkk44\5L5LMN'..G..0DJJ>  00666:

K33< {{"{ F++''FB;33B7CLL% $"4"4R"=> 8 %U++00<<<Ta!!z{{"{))"d1g66 &%"4"45 {{* #Jt{{K$H$O  *&$vFTTf  %$++tvdkk44\5M5MNO'DJJ88w"2tT6::i &Ps   1Pc           
      H   | j                   r| j                  d      t        j                  u r| j                  d      t        j                  u r| j                  d      t        j
                  u rt        | j                   j                  t        j                               t        t        | j                              D cg c]D  }t        j                  j                  | j                  |   t!        | j                   |            F c}S t"        | I  |      S c c}w )N__iter__r   r   r@   )rA   r   rZ   r   r   r   r   r   r
   r   rangerY   r.   r   LazyVariableTrackerrR   r   r,   r   )r/   rT   r   r1   s      r'   r   z-UserDefinedObjectVariable.unpack_var_sequenceK  s    KK00<M00;t||K00?4CSCSS$++001I1IJK s4::/
 0A	 --44JJqM(a8 5  0  w*2..s   A	Dc                 Z    	 | j                   | j                         v S # t        $ r Y yw xY w)NF)r.   r   	TypeErrorr4   s    r'   is_supported_randomz-UserDefinedObjectVariable.is_supported_random\  s1    	::!A!A!CCC 		s    	**c                 R   ddl m} ddlm} | j	                         rt        d |D              rt        d |j                         D              r|D cg c]  }|j                          }}|j                         D ci c]  \  }}||j                          }}}t        |j                        }	 | j                  |i |}
t        |	      }|j                  j                  | j                  ||f        |||      j                  |
      S t        | j                  t         j"                        rn| j                  j$                  }| j                  j&                  }|t(        j*                  j,                  j.                  j0                  u r`|j3                  |j4                        t6        j8                  k(  r4|s2|s0t7        j8                  |j4                        j;                  |||      S |t(        j<                  j>                  j@                  j0                  u r`|j4                  t(        j<                  j>                  j@                  u r/t7        jB                  |jD                        }t7        j8                  |j4                        j;                  ||g|      S t        | j                  tF        jH                        rStK        | j                  jL                        r3t        d tO        jP                  | j                  jR                  | j                  jT                  j                               D              r| jV                  rtY        t[        | jV                  d      j]                  t^        j`                        t[        | jV                  d	      j]                  t^        jb                        t[        | jV                  d
      j]                  t^        jb                               | j                  jR                  D cg c]!  }t6        jB                  je                  |      # }}|jg                  |       | j                  jT                  j                         D ci c]%  \  }}|t6        jB                  je                  |      ' }}}|ji                  |       tk        | j                  jL                        rtm               j;                  |||      S t7        jn                  | j                  jL                        j;                  |||      S tq        | j                        rR| jV                  r2tY        | jV                  j]                  t^        jr                               | ju                  |d||      S tv        | u  |||      S c c}w c c}}w c c}w c c}}w )Nr   )trace_rulesr   r>   c              3   <   K   | ]  }|j                           y wr3   is_python_constant)rj   r   s     r'   rl   z:UserDefinedObjectVariable.call_function.<locals>.<genexpr>k  s     9DqA((*D   c              3   <   K   | ]  }|j                           y wr3   r   rj   vs     r'   rl   z:UserDefinedObjectVariable.call_function.<locals>.<genexpr>l  s     DOqA((*Or   c              3   Z   K   | ]#  }t         j                  j                  |       % y wr3   )r   r=   rQ   r   s     r'   rl   z:UserDefinedObjectVariable.call_function.<locals>.<genexpr>  s)      WA **55a8Ws   )+funcr_   keywords__call__)<rC   r   rD   r?   r   r{   valuesr5   ry   rY   random_callsr.   r   r\   wrap_unspecialized_primitiver   r   
MethodTyperN   __self__r   utils_contextlib_DecoratorContextManagerclonelookupr1   r   TorchCtxManagerClassVariabler   autograd	grad_modeinference_moder=   mode	functoolspartialr   r   	itertoolschainr_   r   rA   r   r   r   r
   ID_MATCHCONSTANT_MATCHrR   rw   updater   r   TorchVariablecallableFUNCTION_MATCHr^   r,   )r/   rT   r_   r0   r   r?   rk   r   r   random_call_indexexample_valuerA   r   rU   r   partial_argspartial_kwargsr1   s                    r'   r   z'UserDefinedObjectVariable.call_functionc  sS    	#, $$&9D99DFMMODD489DqA((*DD9<BLLNKNDAqa--//NFK #BOO 4&DJJ77M&'89FOO""DJJf#=>"2v.KK  DJJ 0 01::&&D**%%C//HHNNN&&s}}599: ==MM-D&12
 00??EEEMMU^^%=%=%L%LL  00: ==MM-SE623 4::y0014::??+ "$**:M:M:T:T:VW 
 {{t{{F3>>|?T?TUt{{F3>>$33 t{{J7BB$33 ?Cjjoo>M	**11!4o   % !JJ//5577DAq 9--44Q777   !!&)"4::??302@@n  **4::??;IIL.  djj!{{dkk44\5P5PQR##B
D&AAw$Rv66O :Kbs   V V	&V)*V#c                 F    t        | j                        rt        d       y y )Nz6UserDefinedObjectVariable with custom __getattribute__)r   r.   r	   r4   s    r'   _check_for_getattributez1UserDefinedObjectVariable._check_for_getattribute  s    "4::.RS /r&   c                 ,    t        | j                        S r3   )r   r.   r4   s    r'   _check_for_getattrz,UserDefinedObjectVariable._check_for_getattr  s    !$**--r&   c                 b   t        | j                  t        j                  j                        sHd| j                  j
                  j                  v s&t        | j                        t        j                  k(  rt        | j                  |      }|S t        j                  | j                  |      }|S )N	__slots__)rH   r.   r   r   r   r1   rB   r7   	threadinglocalrP   rE   rF   )r/   r9   subobjs      r'   _getattr_staticz)UserDefinedObjectVariable._getattr_static  s{    tzz588??3djj22;;;DJJ9??2 TZZ.F  ++DJJ=Fr&   c           	      B   ddl m} ddlm} | j                  }| j
                  rt        | j
                  |      nd }| j                          | j                         } G d d      }	 | j                  |      }	t        |	t&              r3t        j                  |	j(                  | |      j!                  |g i       S t        |	t*        j,                  j.                  j0                        rKt3        |	|      }
t        j                  |	j4                  j6                  |
|      j!                  || gi       S t        |	t8              r0t        j:                  |	j5                  | j                        |      S t        |	t<              r"t        j                  |	j6                  | |      S t        |	t        j                        sJt        |	t        j>                        rt        | j                  t*        j@                  jB                        r_tE        | j                  |      }||	u r;tG        |	d      r/|	jH                  }	|	}|rt        |d      nd }||	u rtG        |	d      r/t        |	t        j>                        r@|jJ                  | j                  urt%        d	       |	j6                  }|rt        |d
      nd }nt        |	t        j                        sJ |	}tM        jN                  |      rt        j                  || |      S tM        jP                  |      rPtS        |      rtU        |      S tW        |      rt        jX                  ||      S t        j:                  ||      S |tE        |di       v sE|j[                  |	      s4t        |	t*        j\                  t*        j@                  jB                  f      r3|r  |||      |	      S |j[                  |	      r|j#                  |	      S |tE        |di       vrt_        |      j`                  jc                  d      rdt_        |      j`                  vrte        |      s|stE        tg        jh                  t_        |      j`                        t_        |      jj                        t_        |      u sJ t        t        |jm                  t_        |      j`                        t_        |      jj                        |      }  |||      |	      S d|i}t        |	t*        j,                  jn                  jp                  t*        j,                  jn                  jr                  t*        j,                  jn                  jt                  f      rt3        |	fi |S t        | j                  t*        j@                  jB                        rK|tv        v rCt        |	tx        jz                        sJ |	s%t        j|                  |	tx        jz                  fi |S |dk(  rt        t_        | j                        fi |S t        j                  | |fi |S # t        $ rh |}	t        |t        j                        r;t        j                  || |      j!                  ||j#                  |      gi       cY S |t%        d       Y $w xY w)Nr   r<   r>   c                       e Zd Zy)=UserDefinedObjectVariable.var_getattr.<locals>.NO_SUCH_SUBOBJNr!   r%   r&   r'   NO_SUCH_SUBOBJr    s    r&   r  r@   z)UserDefined with non-function __getattr___torchdynamo_inlinez"__self__ mismatch for bound methodrN   rB   ztorch.ztorch.optimrA   r1   )ArC   r=   rD   r?   r.   rA   r   r   r   r   rG   rH   r   r   r   rM   r   rR   r	   propertyfgetr   distributionsr   lazy_propertyr   rK   rN   rI   rJ   rL   r   r   r   rP   ru   r  r   rE   ismethod
isfunctionr   r   r   r   rQ   Tensorr7   r#   
startswithr   	importlibimport_moduler"   r[   constraints	_Interval_Real
Constraintr   r   r   r   r)   GetAttrVariable)r/   rT   r9   r=   r?   r.   rA   
getattr_fnr  r   
subobj_vardynamic_subobjr   r`   s                 r'   rS   z%UserDefinedObjectVariable.var_getattr  s   &,

26++DKK.4$$&,,.
	 			K))$/F fh'//T&mBB'(  3 3 9 9 G GH26&IJ//''FmB+, -11tzz*6  ,//fUU 2 23vu//04::uxx7
 %TZZ6N F*wv?T/U33!'FLF,ABRV !F*wv?T/U
 &%"2"23!**$**<!"FG;AFJ7t!&%*<*<===/ 33D$vNN##N3&t,4FCC%$224GG 55d6JJ GE:r22**62LLHHOO 2r626::!,,V4'..v66 z266U&&11(;T%[%;%;;UO++DK,B,BCK(( %[! ! ! $((e)?)?@$u+BVBV 	 /?2v.v66V$##//99##//55##//::
 -V?w??

EHHOO49Ofk&=&=>>> 22K337>  ;+D,<HHH((t?w??_  	K#F*e&8&89 33V-%5%<%<T%B$CRHI 'IJ	Ks   (X- -AZZZr9   c                    |j                   j                  j                  |       ra	 |j                   j                  j                  | |d      }t        j
                  j                  t        |t        j                               S | j                  r<t        t        | j                  |      j                  t        j                               | j!                         s| j#                         rt%        d       	 | j'                  |       t        j
                  j                  d      S # t        $ r Y w xY w# t(        $ r" t        j
                  j                  d      cY S w xY w)NT)
deleted_okzhasattr with custom __getattr__F)r   rq   is_attribute_mutation	load_attrr   r=   rR   rH   DeletedVariableKeyErrorrA   r   r   r   r
   HASATTRr   r   r	   r   rG   )r/   rT   r9   results       r'   call_hasattrz&UserDefinedObjectVariable.call_hasattrI  s$   99!!77=//99$QU9V 1188"69+D+DEE 
 ;;4;;-889M9MN '')T-D-D-F;<	<  &--44T::    	<--44U;;	<s$   AD+ ;/D: +	D76D7:(E%$E%c                 :   ddl m} t        j                  |      r|j                  |j                  n|j                         }  ||t        | j                  |            t        j                  j                  | j                  |j                                     S )Nr   r>   )rD   r?   r   is_valid_keyrA   r5   r   r   r   r   r.   )r/   rT   r   r?   rz   s        r'   r   z'UserDefinedObjectVariable.odict_getitem_  s    , !--c2szz7M JJ'') 	
t{{E2
 
!
!
-
-djj#:P:P:R
SU 	Ur&   r3   r   )r"   r#   r$   __doc__r    _nonvar_fieldsr-   r   r8   rI   r   	lru_cacher   r   r^   r   r   r   r   r   r   rS   r   r  r   r   r   s   @r'   r   r      s     |Q.A.P.PQN.	5 Y  P; &	P;
 -P; 
P;d/"R7/R79UR7	R7hT.
}@~<S <-> <,Ur&   r   c                   8     e Zd Zed        Z fdZ fdZ xZS )KeyedJaggedTensorVariablec                 z    t         j                  j                  d      }|d uxr t        |       |j                  u S )Nztorchrec.sparse.jagged_tensor)sysmodulesgetr7   KeyedJaggedTensor)rU   mods     r'   is_matching_objectz,KeyedJaggedTensorVariable.is_matching_objecto  s3    kkoo=>$E490E0E#EEr&   c                 P    ddl m} t        |      |u sJ t        |   |fi | y )Nr   )r+  )torchrec.sparse.jagged_tensorr+  r7   r,   r-   )r/   r.   r0   r+  r1   s       r'   r-   z"KeyedJaggedTensorVariable.__init__t  s+    CE{////)&)r&   c                    t         j                  j                  j                  r@| j                  4|dv r0t        j                  d      5  t        | !  ||      cd d d        S t        | !  ||      S # 1 sw Y   xY w)N)_length_per_key_offset_per_keyT)#force_unspec_int_unbacked_size_like)	r   _dynamoconfig3force_unspec_int_unbacked_size_like_on_torchrec_kjtrA   r   patchr,   rS   r   s      r'   rS   z%KeyedJaggedTensorVariable.var_getattrz  sk    MM  TT'>>%%$Ow*2t4 POw"2t,, POs   A66A?)r"   r#   r$   rI   r-  r-   rS   r   r   s   @r'   r&  r&  n  s&    F F*- -r&   r&  c                   8     e Zd Z	 	 d fd	Z fdZ fdZ xZS )RemovableHandleVariablec                 @    t        |   di | || _        || _        y r+   )r,   r-   rX   idx)r/   rX   r;  r0   r1   s       r'   r-   z RemovableHandleVariable.__init__  s$     	"6"*r&   c                     |dk(  rN|j                   j                  j                  | j                         t        j
                  j                  d       S t        | !  ||||       y )Nremove)	r   rq   remove_hookr;  r   r=   rR   r,   r^   )r/   rT   method_namer_   r0   r1   s        r'   r^   z#RemovableHandleVariable.call_method  sQ    ("II""..txx8--44T::BT6:r&   c                 r    | j                   r|j                  | j                         gS t        |   |      S r3   )user_code_variable_namecreate_loadr,   reconstruct)r/   codegenr1   s     r'   rC  z#RemovableHandleVariable.reconstruct  s6    ''''(D(DEFFw"7++r&   )NN)r"   r#   r$   r-   r^   rC  r   r   s   @r'   r9  r9    s     		;, ,r&   r9  )7r   rr   r   r  rE   r   r   r(  r   r   typingr   r   torch._dynamo.configr   torch.nntorch._guardsr   rC   r   allowed_functionsr   excr	   guardsr
   r   rA   r   r   r   r   r   r   r   r   r   r   r   r   r   r   baser   r   ctx_managerr   r   dictsr   r    r)   r   r&  r9  r%   r&   r'   <module>rO     s           
      (  *  0 U U
 
 
 0 L $	/ 	]/2 ]/@`U 3 `UF- 9 -.,o ,r&   