
    ǄgWq                   	   d dl m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Zd dlmZ d dlmZ d dl
mZmZmZmZmZmZmZmZ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# d dl$Z$d dl$m%Z%m&Z&m'Z'm(Z( d d	l)m*Z*m+Z+ d d
l,m-Z- d dl.m/Z/m0Z0m1Z1m2Z2m3Z3 d dl4m5Z5 d dl6m7Z7 d dl8m9Z9 d dl:m;Z; d dl<m=Z= d dl>m?Z?m@Z@ d dlAmBZB d dlCmDZD d dlEmFZFmGZG d dlHmIZImJZJmKZKmLZL d dlMmNZN d dlOmPZP ddlQmRZRmSZSmTZT erd dlUmVZV d dlWmXZX d dlYmZZZ d dl[m\Z\m]Z]  ej                  e_      Z`	 e$j                  j                  e_d      Zc G d  d!      Zg eg       ZheZie$j                  j                  Zl ed"      Zme$j                  j                  j                  ZpdZqd ar G d# d$      Zse G d% d&et             Zue G d' d(et             Zve G d) d*et             Zwe G d+ d,et             Zxd]d-Zyej                  d^d.       Z{d_d/Z|d`d0Z}dad1Z~ ej                  d      dbd2       Z ej                  d      dcd3       Zddd4Zded5Z G d6 d7      Z ej                  d      dfd8       Zej                  	 	 	 	 dgd:       Zdcd;Z G d< d=      Z G d> d?      Z G d@ dAe(      Zee?dBdCf   ZeBe G dD dE                    ZdhdFZeBe G dG dH                    ZeB edIJ       G dK dL                    ZeB edIJ       G dM dNe                    ZeB edIJ       G dO dP                    Z G dQ d9eG      ZeZ	 	 	 	 	 	 	 	 	 	 	 	 didRZ G dS dTe=      ZdjdUZdkdVZdldWZe$j                  j6                  j8                  j:                  ee$j                  j                  j<                  j:                  dX e$j                  j                  j>                  j:                  dY e$j                  j                  j@                  j:                  dZ iZ eye$j                  j                  jD                  j:                  e$j                  j                  jF                  j:                  e$j                  j                  jH                  j:                        Zd d[lmZmZmZmZmZmZmZmZ ej^                  dfd\       Zy# ed$ rZed efee      v rd dlZcneeY dZe[edZe[eww xY w)m    )annotationsN)defaultdict)	dataclass)AnyCallablecastDict	GeneratorIterableListLiteralMappingOptionalSequenceSetTupleTypeTYPE_CHECKINGTypeVarUnion)Self	TypeGuard)ReferenceType)SymBoolSymFloatSymIntTensor)is_functorch_wrapped_tensoris_legacy_batchedtensor)suggest_memory_format)	assert_eqassert_metadata_eqis_sparse_anyis_sparse_compressedMetaConverter)render_call)immutable_dict)normalize_function)StorageWeakRef)TorchFunctionMode)IntLikeTypepy_sym_types)dataclass_slots)no_dispatch)is_traceable_wrapper_subclassTorchDispatchMode)PyTreetree_map	tree_map_TreeSpec)count)CapturedTraceback   )_CacheKeyState_PySymInputStub_SymIntOutputStub)TracebackType)Source)
OpOverload)ShapeEnvSymbolicContextnot_implementedz 'not_implemented' not registeredc                      e Zd Zy)_UnassignedN)__name__
__module____qualname__     e/home/mcse/projects/flask_80/flask-venv/lib/python3.12/site-packages/torch/_subclasses/fake_tensor.pyrB   rB   X   s    rG   rB   Tc                      e Zd ZddZddZy)IncrementRecursionCountc                    t         dz  a y Nr7   RECURSION_COUNTselfs    rH   __init__z IncrementRecursionCount.__init__o       1rG   c                    t         dz  a y rM   rN   rP   s    rH   __del__zIncrementRecursionCount.__del__s   rS   rG   NreturnNone)rC   rD   rE   rR   rU   rF   rG   rH   rK   rK   n   s    rG   rK   c                      e Zd ZU ded<   y)UnsupportedFakeTensorExceptionstrreasonNrC   rD   rE   __annotations__rF   rG   rH   rZ   rZ   x   s    KrG   rZ   c                      e Zd ZU ded<   y)DynamicOutputShapeExceptionr=   funcNr]   rF   rG   rH   r`   r`   }       
rG   r`   c                      e Zd ZU ded<   y)DataDependentOutputExceptionr=   ra   Nr]   rF   rG   rH   rd   rd      rb   rG   rd   c                      e Zd ZU ded<   y)UnsupportedOperatorExceptionr=   ra   Nr]   rF   rG   rH   rf   rf      rb   rG   rf   c                 .    t         j                  | d      S NT)dictfromkeys)itemss    rH   ordered_setrl      s    ==%%rG   c               #  *  K   t         j                  j                  t         j                  j                  j                        } 	 |  |  t         j                  j                  |        y y # |  t         j                  j                  |        w w xY wwN)torch_C_unset_dispatch_mode_TorchDispatchModeKeyFAKE_set_dispatch_mode)olds    rH   unset_fake_temporarilyrv      sl     
((
'
'(F(F(K(K
LC-	?HH'', 3?HH'', s   ABA, 	#B,$BBc                .   t        |       sJ g }| g}|r|j                         }t        |      s$t        |t              sJ |j	                  |       A|j                         \  }}t        |      D ]  }|j	                  t        ||              |r|S rn   )r/   pop
isinstancer   append__tensor_flatten__reversedgetattr)subclassplain_tensorstodocurr
inner_keys_keys          rH   get_plain_tensorsr      s    (222"$M:D
xxz,T2dF+++  &//1
AJ' 	,CKKc*+	,  rG   c                   t        | t              ryt        |       rht        |       j	                  |       \  }}|D cg c]  }t        | |       }}t        d |D              }t        d |D              }||k(  sJ d       |S t        | t              rht        j                  |       rSt        j                  j                         }t        j                  j                  j                  | |      }t        |      S t        | t              r?t!        |       r4t        j                  j                  j#                  |       }t        |      S yc c}w )NTc              3  2   K   | ]  }t        |        y wrn   is_fake.0xs     rH   	<genexpr>zis_fake.<locals>.<genexpr>        =awqz=   c              3  2   K   | ]  }t        |        y wrn   r   r   s     rH   r   zis_fake.<locals>.<genexpr>   r   r   z got mixed fake and real tensors!F)ry   
FakeTensorr/   typer{   r}   allanyr   ro   _is_functional_tensorrp   $_functionalization_reapply_views_tls
_functorch_unwrap_functional_tensorr   r   get_unwrapped)	r   attrsr   attrflattened_tensorsall_fakeany_fakereapply_views	unwrappeds	            rH   r   r      s   !Z $Q'7--a0q:?@$WQ-@@=+<===+<==8#G%GG#	Av	5#>#>q#AEEGHH''AA!]S	y!!	Av	#>q#AHH''55a8	y!! As   Ec           	        t        | t              r| j                  S t        |       rR| j	                         \  }}|D cg c]  }t        t        | |             }}|d   t        fd|D              sJ S t        | t              rht        j                  |       rSt        j                  j                         }t        j                  j                  j                  | |      }t        |      S t        | t              r?t        |       r4t        j                  j                  j!                  |       }t        |      S y c c}w )Nr   c              3  &   K   | ]  }|u  
 y wrn   rF   )r   r   ms     rH   r   z&maybe_get_fake_mode.<locals>.<genexpr>   s     )a16)s   )ry   r   	fake_moder/   r{   maybe_get_fake_moder}   r   r   ro   r   rp   r   r   r   r   r   )tinner_tensor_namesr   t_namemodesr   r   r   s          @rH   r   r      s   !Z {{$Q' ! 4 4 6ABT
8>6 23
 
 !H)5))))	Av	5#>#>q#AEEGHH''AA!]S	"9--	Av	#>q#AHH''55a8	"9--
s    Ec                T    t         j                  j                  | j                        S rn   )ro   rp   _SchemaInfo_schemara   s    rH   get_schema_infor      s    88--rG   c                    ddl m} t        j                  j                  }||    j
                  j                  d      xr ||    j                  t        |      v S )Nr   decomposition_tableztorch._decomp)	torch._decompr   ro   _decompdecompositionsrD   
startswithrC   dir)ra   r   r   s      rH   torch_decomp_decompositionsr      sV    1]]11N t$//:: D
d
#
,
,N0C
CDrG   c                p    t         j                  |      }|D cg c]  }t        ||       s| c}S c c}w rn   )pytreetree_leavesry   )tytree	flat_valselems       rH   tree_flatten_onlyr      s/    ""4(I&?T*T2*>D???s   33c                    t        |       t        u xrk | j                  t        j                  k(  xrL | j
                  xs= | j                  xs/ t        |       xs" t        |       xs t        j                  |        S rn   )
r   r   layoutro   strided	is_sparse	is_nestedr   r   r   r   s    rH   _is_plain_tensorr      sx    Q6 	
HH%	
 KK .{{.*1-. 'q). **1-

rG   c                      e Zd ZU e	 	 dd       Zded<   ded<   ded<   ddd	dd
ZddZddZddZ	ddZ
	 	 ddddd	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 ddZ	 	 	 	 	 	 	 	 ddZy)FakeTensorConverterc                .    | j                   j                  S rn   )meta_convertertensor_memorP   s    rH   r   zFakeTensorConverter.tensor_memo  s     ""...rG   r%   r   z)Dict[StorageWeakRef, List[ReferenceType]]constant_storage_mappingboolexportF	copy_datar   c               B    t        |      | _        || _        i | _        y )N)r   )r%   r   r   r   )rQ   r   r   s      rH   rR   zFakeTensorConverter.__init__  s     +i@ )+%rG   c                "   t        |t              r|j                  J t        |j                  j	                               }|| j
                  vrg | j
                  |<   | j
                  |   j                  t        j                  |             y rn   )	ry   r   constantr)   _typed_storager   rz   weakrefref)rQ   fake_tensorweak_sts      rH   add_constant_storage_mappingz0FakeTensorConverter.add_constant_storage_mapping  sz     +z2{7K7K7WWW !5!5!D!D!FG
 $77757D))'2%%g.55gkk+6NOrG   c                    t        |t              rJ t        |j                               }|| j                  vry | j                  |   D ]#  } |       }||j                          d |_        % | j                  |= y rn   )ry   r   r)   r   r   _fix_weakrefr   )rQ   tensorr   weak_tensor_reftens        rH   invalidate_constant_aliasesz/FakeTensorConverter.invalidate_constant_aliases%  s    fj111 !6!6!89$777#<<WE 	$O!#C  "#		$ ))'2rG   c                    | j                   j                  j                  j                  |      }|y | j                  j                  |      S rn   )r   	describerlookup_tensorgetr   )rQ   r   tids      rH   	_get_memozFakeTensorConverter._get_memo4  sC    !!++99==a@;##C((rG   c                    | j                   j                  j                  |      }|| j                   j                  |<   y rn   )r   r   get_tensor_idr   )rQ   r   vr   s       rH   set_tensor_memoz#FakeTensorConverter.set_tensor_memo:  s4    !!++99!</0'',rG   NT)sourcesymbolic_contexttracec                  |sk|si|rgt         j                  j                  j                         x}r=|j                  v r/|j                     }ddlm}	 t        ||	      sJ |j                  }| j                        }
|
|
S j                  j                  rt        d      t              t         j                  j                  u rrJ dfd}| j!                  |||||      }|t"        u rt        d      ddlm} d }| j(                  st+              rj-                         dk(  rwj                  j                  dk(  r]j.                  t         j0                  t         j2                  t         j4                  t         j6                  t         j8                  fv r| t        ||      s|dd	lm}m} dd
lm} tA               5  jC                         }d d d        tE        jF                  |      st        ||      r|jH                  }n ||      }|jK                  |||jL                        }j.                  t         j0                  k(  r|jO                  |||      |_(        n6j.                  t         j8                  k(  r|jS                  |||      |_(        r| jU                  |       |S # 1 sw Y   xY w)Nr   )StatefulSymbolicContextzquantized nyi in meta tensorsc                t    t               5  t         |        rnd       cd d d        S # 1 sw Y   y xY w)Nr   )r.   r   )make_meta_texisting_devicer   make_constantr   s    rH   mk_fake_tensorz<FakeTensorConverter.from_real_tensor.<locals>.mk_fake_tensore  s;      !M# #0QT  s   .7)	shape_envcallbackr   r   r   zmeta converter nyi)RandomValueSourcecpu)CallMethodItemSourceFloatTensorSource)
DimDynamic)r   dynamic_dim)hintr   )r   zCallable[[], object]rW   r   )+ro   _guardsTracingContexttry_gettensor_to_context%torch.fx.experimental.symbolic_shapesr   ry   tensor_sourcer   deviceis_quantizedrZ   r   nn	Parameterr   NotImplementedtorch._dynamo.sourcer   r   r   dimdtypeint64int32int16int8float64r   r   r   r.   itemmathisnanbasecreate_unspecified_symbolDYNAMICcreate_symintnode	item_memocreate_symfloatnoder   )rQ   r   r   r   r   r   r   r   tracing_contextr   
maybe_memor   outr   valuer   r   r   item_sourcesymbolr   s    ```                @rH   from_real_tensorz$FakeTensorConverter.from_real_tensorD  s     9"'--">">"F"F"HHH999'6'H'H'K$ &&68OPPP-;;F^^A&
!((>>01PQQ7ehh((($$$	 	$ !!#- " 
 . 01EFF: #1& U[[%++uzz5==QR"$ v'89 %TH !!::e$f&78"(++K"6v">K"<<& * 2 2 =  77ekk)$-$?$?"* %@ %CM
 WW-$-$A$A"* %B %CM
 --c2
=! !s   KKc                    |j                   j                  dk(  sJ d|j                   j                   d       | j                  |      }||S t        |||      }| j	                  ||       |S )Nmetaz$tensor's device must be `meta`, got z instead)r  r   r   r   r   )rQ   r   r   r  r  r  s         rH   from_meta_and_devicez(FakeTensorConverter.from_meta_and_device  sv     HHMMV#	J1!((--I	J# ^^A&
!Av.Q$
rG   )rW   zweakref.WeakValueDictionary)r   r   r   r   rW   rX   )r   r   rW   rX   )r   r   rW   rX   )r   r   rW   Optional[FakeTensor])r   r   r   r   rW   rX   )FN)r   FakeTensorModer   r   r   r   r   Optional[ShapeEnv]r   Optional[Source]r   Optional[SymbolicContext]r   r   rW   r   )r   r%  r   r   r  torch.devicerW   r   )rC   rD   rE   propertyr   r^   rR   r   r   r   r   r   r#  rF   rG   rH   r   r     s    /	$/ / "!GGL,1% +P3)1 $(,F $(6:F!F F 	F
 &F !F 4F F 
FR',2<H	rG   r   c                 <   t         j                  j                         rHt         j                  j                  t        j
                  dd      nt        j                  dd       t         j                  j                         rt        j
                  dd       y y )Nr7   cudar  xpu)ro   r,  is_availableversionhipemptyzerosr.  rF   rG   rH   init_gpu_contextr4    sj     zz  }}  ( KK&)Qv. yy	Qu	%  rG   r%  c              #    K   | j                   }t        j                  j                         }||k(  sJ | d|        t        j                  j	                         5  d| _         t        j                  j                         5  t        j                  j                  d       	 d  || _         	 d d d        d d d        y # || _         w xY w# 1 sw Y   xY w# 1 sw Y   y xY ww)N, T)in_kernel_invocationro   rp   _meta_in_tls_dispatch_include_DisableTorchDispatch_PreserveDispatchKeyGuard!_set_meta_in_tls_dispatch_include)r   prev_in_kernelmeta_in_tlss      rH   in_kernel_invocation_managerr>    s     
 33N((88:K.(L[MN;K*LL(		'	'	) 	@)-	& XX//1 	@HH66t<@1?	.	@		@ 	@ 2@	.	@ 	@		@ 	@sN   AC$&C C#C 'C/C7	C$ 	C		CC	CC!C$c                    t         j                  j                  | j                         j	                  d      d   j	                  d      d         S )Nz::.r   )ro   rp    _should_allow_numbers_as_tensorsnamesplitr   s    rH   should_allow_numbers_as_tensorsrE    sB    8844		$#))#.q1 rG   c                  J    e Zd Zej                  j                  dd      dk(  Zy)FakeTensorConfigTORCH_FAKE_TENSOR_DEBUG01N)rC   rD   rE   osenvironr   debugrF   rG   rH   rG  rG    s    JJNN4c:cAErG   rG  c                  n    e Zd ZU ded<   ded<   ddddZddZdd	Zdd
ZddZ	 d	 	 	 	 	 ddZ	ddZ
y)SymIntMemoDescriptorr[   _namer   _is_nested_intFis_nested_intc                   || _         y rn   )rQ  )rQ   rS  s     rH   rR   zSymIntMemoDescriptor.__init__!  s
    +rG   c                    || _         y rn   rP  )rQ   ownerrC  s      rH   __set_name__z!SymIntMemoDescriptor.__set_name__$  s	    
rG   c                     d| j                    S )Nr   rV  rQ   objs     rH   _memozSymIntMemoDescriptor._memo'  s    4::,rG   c                "    d| j                    dS )Nr   _vcrV  rZ  s     rH   _memo_vczSymIntMemoDescriptor._memo_vc*  s    4::,c""rG   c                "    d| j                    dS )Nr   _epochrV  rZ  s     rH   _memo_epochz SymIntMemoDescriptor._memo_epoch1  s    4::,f%%rG   Nc                `   t        || j                  |            x}y | j                  s(t        || j                  |            |j                  k7  s>| j                  sOt        || j                  |            |j                  j                  k7  rt        || j                  |      d        y |S rn   )	r}   r\  rQ  r_  _versionrb  r   epochsetattr)rQ   r[  objtypers       rH   __get__zSymIntMemoDescriptor.__get__4  s     djjo..A7 ##T]]35G(HCLL(X##T--c23s}}7J7JJCC$/rG   c                   |Ut        || j                  |      d        t        || j                  |      d        t        || j                  |      d        y |j	                         r| j
                  rt        || j                  |      |       | j
                  s&t        || j                  |      |j                         t        || j                  |      |j                  j                         y y rn   )	rf  r\  r_  rb  is_inferencerQ  rd  r   re  )rQ   r[  r  s      rH   __set__zSymIntMemoDescriptor.__set__E  s    =CC$/Cs+T2C))#.5!!#t':':CC%0&&T]]3/>C))#.0C0CD	 (;rG   )rS  r   rW   rX   )rW  r[   rC  r[   rW   rX   )r[  r   rW   r[   rn   )r[  r   rg  zOptional[Type[FakeTensor]]rW   Optional[torch.SymInt])r[  r   r  rm  rW   rX   )rC   rD   rE   r^   rR   rX  r\  r_  rb  ri  rl  rF   rG   rH   rO  rO    sU    J 05 , #& FJ(B	"	ErG   rO  c                      e Zd ZU dZded<   ded<   ded<   ded<    e       Z e       Z e       Z ed	
      Z	e
j                  j                  j                  Zedd       Zej"                  dd       Zedd       Zej"                  dd       Ze	 	 d	 	 	 	 	 	 	 	 	 	 	 dd       Zd fdZed d       Zeed e       f	 	 	 	 	 	 	 	 	 d!d              Ze	 	 	 	 	 	 d"d       Zdd	 	 	 d#dZd$dZ xZS )%r   al  
    Meta tensors give you the ability to run PyTorch code without having to
    actually do computation through tensors allocated on a `meta` device.
    Because the device is `meta`, meta tensors do not model device propagation.
    FakeTensor extends MetaTensors to also carry an additional `fake_device`
    which tracks devices that would have been used.
    r)  fake_devicer%  r   Optional[Tensor]r   real_tensorTrR  c                p    | j                   j                  rt        j                  d      S | j                  S Nr"  )r   r7  ro   r  ro  rP   s    rH   r  zFakeTensor.deviceo  s*    >>..<<''###rG   c                    t         rn   NotImplementedErrorrQ   r   s     rH   r  zFakeTensor.devicev      !!rG   c                    t        d      )Nz+torch.compile doesn't support named tensors)rZ   rP   s    rH   nameszFakeTensor.names  s    ,9
 	
rG   c                    t         rn   ru  rw  s     rH   rz  zFakeTensor.names  rx  rG   c                   t        j                  | ||j                  d|      }|j                  s t        j
                  j                  |       nt        j
                  j                  |       |j                  j                  dk(  sJ |j                  j                         t        |t        j                        r|nt	        j                  |      }|j                  s|j                  dk7  sJ |j                  dv r
t                |j                  dddt        j
                  j                         fv r|j                  t        t        |j                        j!                         rKt	        j                  |j                   dt        t        |j                        j#                                }n"t	        j                  |j                   d	      }||_        ||_        ||_        t        |t*              rJ ||_        d |_        d |_        d |_        d |_        t6        j8                  rt;        j<                         |_        |S )
NT)dispatch_devicedevice_for_backend_keysr"  )r,  r.  r,  hpur.  :z:0) r   _make_subclassrequires_grad_allow_unsafe_data_ptr_accessro   rp   _set_throw_on_mutable_data_ptr(_set_warn_deprecated_on_mutable_data_ptrr  r   ry   
allow_metar4  _get_privateuse1_backend_nameindexr}   is_initializedcurrent_devicero  r   r   r   rq  nonzero_memor  unique_memonested_int_memorG  rM  r6   extract_debug_trace)clsr   r   r  r   rq  rQ   s          rH   __new__zFakeTensor.__new__  s    $$ $*
 66HH33D9HH==dC{{6);4;;+;+;;)%fell;fAU ##;;&(((;;/) KKuehh&L&L&NOP$ufkk*99;{{m1WUFKK%@%O%O%Q$RS R&89!" k:666& #!! 1 9 9 ;DrG   c                "    t         |           y rn   )superrR   )rQ   argskwargs	__class__s      rH   rR   zFakeTensor.__init__  s    rG   c                $    |j                  |       S rn   )from_tensor)r   r   s     rH   r  zFakeTensor.from_tensor  s    $$Q''rG   rF   c                   |t         j                  j                  j                  j                  k(  r`t        |      dk(  rt        |d   t              sJ |d   j                  j                  rt        j                  d      S |d   j                  S t        j                  |      x}r ||      S |D cg c]  }t        |t              r|t        us|  }}|rt        j                   d|       t"        S d }t%        j&                  |i |D ]   }	t        |	t              s|	j                  } n |J t         j(                  j+                  t         j(                  j,                  j.                        }
|
rt        j                   d||
       t"        S |j                  rJ |5   ||i |cd d d        S c c}w # 1 sw Y   y xY w)Nr7   r   r"  z(FakeTensor unrecognized subclass(es): %sz(FakeTensor mode already active: %s in %s)ro   opsprimr  defaultlenry   r   r   r7  ro  _DISPATCH_META_HANDLERSr   
issubclassr   not_implemented_logrM  r  r   arg_tree_leavesrp   _get_dispatch_moderr   rs   )r  ra   typesr  r  handlerr   unrecognized_typesr   argmaybe_cur_fake_modes              rH   __torch_dispatch__zFakeTensor.__torch_dispatch__  s    599>>((000t9>ja*&EEEAw  55||F++Aw*** .11$77774=  
Jq*$=!6/A
 
 %%:<N "!	))4:6: 	C#z*MM		
 $$$ $hh99HH**//
 %%:#
 "!1111 	)((	) 	)K
J	) 	)s   5F8F8F8&F==Gc                     d d}d ddd fd}|D ]
  } ||        t               rd}t        j                  d      
J d         |fS )	NFc                ^    | j                   j                  dk(  xr | j                         dk(  S )Nr   r   )r  r   r
  r   s    rH   cpu_zero_dimz4FakeTensor._find_common_device.<locals>.cpu_zero_dimB  s$    88==E):aeegl:rG   c                    t        | t              sy | j                   |       y  |       }| j                  k(  rr|y |ry r| j                  |y t        d d d| j                         )Nz,Unhandled FakeTensor Device Propagation for z, found two different devices r6  )ry   r   r  RuntimeError)r   t_is_cpu_zero_dimcommon_devicer  ra   is_cpu_zero_dims     rH   merge_devicesz5FakeTensor._find_common_device.<locals>.merge_devicesE  s     a,$ !".q/ ,Qxx=("&7O !  !"3 >tfDbcpbqqstut|t|s}~ rG   Tr   z!Could not find common device for r   r   rW   r   )r   objectrW   rX   )rE  ro   r  )ra   	flat_argshas_scalar_only_inputsr  r  r  r  r  s   `    @@@rH   _find_common_devicezFakeTensor._find_common_device5  s     !&	; 	  	D  	C#	 +40]5J%)"!LL/M(T,MdV*TT(444rG   r7   )coeffc               z    | j                   !| j                  j                  d       | _         | j                   |z  S )Nnt_tensor_id)r  r   create_symbolic_nested_int)rQ   r  s     rH   get_nested_intzFakeTensor.get_nested_intv  sC    
 '#'>>#L#L! $M $D  ##e++rG   c                    | j                         dk(  r| j                         S | j                         dk(  r| D cg c]  }|j                          c}S | D cg c]  }|j                          c}S c c}w c c}w )Nr   r7   )r
  r  tolist)rQ   r   s     rH   r  zFakeTensor.tolist  s^    88:?99;XXZ1_,01DDIIK11.23dDKKM33 23s   A3A8)rW   r)  )r   r)  rW   rX   )rW   	List[str])r   r  rW   rX   )NN)r   r%  r   r   r  r)  r   rp  rq  rp  rW   r   )r  r  r  r  rW   rX   )r   r   r   r%  rW   r   
ra   r=   r  Sequence[Type]r  Sequence[object]r  Mapping[str, object]rW   r  )ra   r=   r  r  rW   zTuple[torch.device, bool])r  zUnion[int, torch.SymInt]rW   torch.SymInt)rW   r   )rC   rD   rE   __doc__r^   rO  r  r  r  r  ro   rp   rr   rs   	_mode_keyr*  r  setterrz  staticmethodr  rR   r  classmethodr5   r'   r  r  r  r  __classcell__r  s   @rH   r   r   Q  s    !!
 ()L$&I&(K
 +>O ..33I$ $ ]]" "& 
 

 \\" "  &*(,:!: : 	:
 #: &: 
: :\ ( ( 

 "$'5'7G)G) G) 	G)
 %G) 
G)  G)R >5>5%5>5	">5 >5F +,	, (	, 
		,4rG   r   r9   r:   c                      e Zd ZU dZded<   ded<   ded<   ded<   d	ed
<   ded<   ded<   ded<   ded<   ded<   ded<   ded<   ded<   ded<   ded<   ded<   ded<   	 	 	 	 	 	 	 	 ddZy) TensorMetadatazK
    The Tensor metadata relevant to hashing FakeTensors when caching.
    ztorch.dtyper  zTuple[_MetadataIntLike, ...]shapestrider)  r  ztorch.layoutr   zOptional[torch.memory_format]memory_format_MetadataIntLikestorage_offsetzOptional[_MetadataIntLike]storage_bytesr   r  r  is_conjis_negrk  r   Optional[bool]is_coalescedOptional[int]	dense_dim
sparse_dimc                >   t        j                  |       D ]  }t        | |j                        }t	        |t
        t        t        j                  f      r|j                  |||       Rt	        |t              r|j                  ||       u|j                  |        y rn   )dataclassesfieldsr}   rC  ry   tuplelistro   Size_prep_args_for_hashr   convert_sym_intrz   )rQ   resultmodestatefieldr  s         rH   _flatten_intozTensorMetadata._flatten_into  s}     !''- 		%ED%**-E%%uzz!:; ((>E6*%%fe4e$		%rG   N)r  List[object]r  r%  r  r8   rW   rX   )rC   rD   rE   r  r^   r  rF   rG   rH   r  r    s     ''((00$$--MLO  %% % 	%
 
%rG   r  c                   t        |       }| j                  st        |       s| j                  |      sd}| j	                         }t        | j                  | j                  | j                  t        j                  k(  r| j                         nd| j                  | j                  ||t        |       s| j                         j                         nd| j                  | j                   | j#                         | j%                         | j'                         | j(                  | j(                  r| j+                         ndt        |       r| j-                         ndt        |       r| j/                               S d      S )z1
    Extract the TensorMetadata of a tensor.
    )r  NrF   )r    _has_symbolic_sizes_stridesr#   is_contiguousr  r  r  r  r   ro   r   r  r  untyped_storagenbytesr  r  r  r  rk  r   r  r  r  )r   r  r  s      rH   extract_tensor_metadatar    s    4I3KM 	
%%];%%'N		hh%--/
R		,9!,<""$$						
		KKT&q)t'*% $ 15% rG   c                  F    e Zd ZU dZded<   ded<   ddZddZddZdd	Zy
)_DispatchCacheKeyz0
    Key for the FakeTensor dispatch cache.
    Tuple[object, ...]r   int	hashvaluec                2    || _         t        |      | _        y rn   )r   hashr  )rQ   tups     rH   rR   z_DispatchCacheKey.__init__  s    crG   c                X    t        |t              xr | j                  |j                  k(  S rn   )ry   r  r   )rQ   others     rH   __eq__z_DispatchCacheKey.__eq__  s"    %!23MEII8MMrG   c                    | j                   S rn   )r  rP   s    rH   __hash__z_DispatchCacheKey.__hash__  s    ~~rG   c                h    | j                   D ]#  }t        |t              s|j                          % y rn   )r   ry   r9   strip_shape_env)rQ   r   s     rH   r  z!_DispatchCacheKey.strip_shape_env  s-      	$A!_-!!#	$rG   N)r  r  rW   rX   )r  r  rW   r   )rW   r  rV   )	rC   rD   rE   r  r^   rR   r  r  r  rF   rG   rH   r  r    s)     
N#N$rG   r  T)frozenc                  0    e Zd ZU dZded<   ded<   ded<   y)_DispatchCacheEntryaD  
    Entry type for the FakeTensor dispatch cache. Accounts for two possibilities:
    1) The op is inplace, and a hit means we need to alias the argument at a
       given index.
    2) We need to synthesize a new FakeTensor given tensor metadata. For view
       ops, we further capture the index of the arg to alias.
    r  inplace_idxzOptional[TensorMetadata]metadataview_idxNrC   rD   rE   r  r^   rF   rG   rH   r   r     s     &&rG   r   c                      e Zd ZU dZded<   y)_BypassDispatchCachez<
    Signals cases that should skip FakeTensor caching.
    r[   r\   Nr  rF   rG   rH   r  r    s     KrG   r  c                  :    e Zd ZU dZded<   ded<   ded<   ded<   y)	DispatchCacheInfozG
    Information about the state of the FakeTensor dispatch cache.
    r  hitsmissesDict[str, int]bypassessizeNr  rF   rG   rH   r  r    s     IK
IrG   r  c                      e Zd ZU i Zded<   dZded<   dZded<    ee      Z	ded<   dZ
ded	<   d
Zded<   ded<   ded<   ded<   ded<   dZded<   dZded<   dd
ddd
d	 	 	 	 	 	 	 	 	 	 	 d@ fdZdAdZdBdZedCd       ZedDd       Zed e       f	 	 	 	 	 	 	 	 	 dEd       ZdF fd Z	 	 	 	 	 	 	 	 dG fd!ZedCd"       ZedHd#       ZedAd$       Z	 	 	 	 	 	 	 	 	 	 dEd%Z	 	 	 	 	 	 	 	 	 	 dId&Z	 	 	 	 	 	 	 	 dJd'Z	 	 	 	 	 	 	 	 dKd(Z 	 	 	 	 	 	 	 	 	 	 	 	 	 	 dLd)Z!	 	 	 	 	 	 	 	 	 	 	 	 dMd*Z"	 	 	 	 	 	 	 	 	 	 	 	 dNd+Z#d e       f	 	 	 	 	 	 	 	 	 dEd,Z$	 	 	 	 	 	 	 	 	 	 dOd-Z% e&d.d/d0d1d2d3d4d5      Z'dPd6Z(	 	 	 	 	 	 	 	 	 	 dQd7Z)	 	 	 	 	 	 	 	 	 	 dRd8Z*dd9	 	 	 dSd:Z+ e&e,jZ                  j\                  e,j^                  j`                  e,jb                  j`                  e,jd                  j`                  e,jf                  j`                  e,jh                  j`                  e,jj                  j`                  e,jl                  j`                  e,jn                  j`                  e,jp                  jr                  e,jt                  j`                        Z;dPd;Z< e&e,jz                  j`                  e,j|                  j`                        Z?dTd<Z@	 	 	 	 	 	 	 	 	 	 dUd=ZAddddd>	 	 	 	 	 	 	 	 	 	 	 dVd?ZB xZCS )Wr%  z,Dict[_DispatchCacheKey, _DispatchCacheEntry]cacher   r  
cache_hitscache_missesr  cache_bypassesre  Fr   r7  static_shapesr&  r   zOptional[str]_stackr  r@  nt_tensor_id_counternt_tensor_id_initial_countTN)allow_fallback_kernelsallow_non_fake_inputsr   r  r   c                  t         j                  dt        |              t        |           || _        dd l}dd l}|j                  j                  j                  | _        t        | j                  |      | _        ||| _        n	|d u | _        d| _        |j                  j                  j                   | _        |j                  j                  j$                  | _        |j(                  j                  j*                  xr | j                   | _        |j(                  j                  j.                  | _        || _        d| _        g | _        || _        t;        j<                         | _        d | _         |jB                  jD                  jF                  | _$        dd l%}|jL                  jN                  jP                  jR                  | _*        | jT                  | _+        y )Nzcreate_mode 0x%xr   r   F),logrM  idr  rR   r  torch._dynamo.configtorch._functorch.configr   config"fake_tensor_propagate_real_tensorspropagate_real_tensorsr   fake_tensor_converterr  allow_scalar_outputs(fake_tensor_allow_unsafe_data_ptr_accessr  fake_tensor_allow_metar  _dynamofake_tensor_cache_enabledcache_enabled$fake_tensor_cache_crosscheck_enabledcache_crosscheck_enabledr  r7  enter_stackr   	tracebackextract_stack_stack_tracer  rp   rr   rs   r  $torch.nested._internal.nested_tensornested	_internalnested_tensor_tensor_id_counterr  r  )rQ   r  r  r   r  r   ro   r  s          rH   rR   zFakeTensorMode.__init__E  s   $ 			$bh/&<##& ##FF 	# &911&
"
 $!.D!*d!2D %*! ##LL 	*  **11HHMM  :: 0/// 	
 MM  EE 	% &;" %*!  	 #%335 77<<3 LL""00CC 	' %)$C$C!rG   c                &    | j                   | _        y rn   )r  r  rP   s    rH   reset_nt_tensor_id_counterz)FakeTensorMode.reset_nt_tensor_id_counter  s    $($C$C!rG   c                B    t        |t              xr |j                  | u S rn   )ry   r   r   rQ   r   s     rH   is_our_fakezFakeTensorMode.is_our_fake  s    !Z(@Q[[D-@@rG   c                    t         j                  j                         r?t         j                  j                         rJ t         j                  j	                          S t         j                  j	                          S rn   )ro   r.  _is_compiledr,  r/  rP   s    rH   avoid_device_initz FakeTensorMode.avoid_device_init  sT    99!!#zz..000yy--///::**,,,rG   c                    | j                   3dj                  t        j                  | j                              | _         | j                   S )N )r  joinr+  format_listr-  rP   s    rH   stackzFakeTensorMode.stack  s7    ;;'')"7"78I8I"JKDK{{rG   rF   c                   t         j                  j                  t         j                  j                  j                        J |       	 | j                  ||||      S # t        $ r t        j                  d        w xY w)Nzfake tensor raised TypeError)	ro   rp   r  rr   rs   dispatch	TypeErrorr  	exception)rQ   ra   r  r  r  s        rH   r  z!FakeTensorMode.__torch_dispatch__  sr     HH''(F(F(K(KLT		T	==udF;; 	MM89	s   
A  A>c                   dd l }d }| j                  r5|j                  j                         }|j                  j	                  d       |j                  j                  | j                        }| |ur,| j                  j                  d||f       t        | )         S |j                  j                  |        | j                  j                  dd |f       | S )Nr   TF)r.  r:  rp   _only_lift_cpu_tensors_set_only_lift_cpu_tensorsrq   r  r*  rz   r  	__enter__rt   )rQ   ro   prev_only_lift_cpu_tensorsmaybe_prev_fake_moder  s       rH   rG  zFakeTensorMode.__enter__  s    3%)"!!).)H)H)J&HH//5$xx<<T^^L++##+-GH 7$&& HH''-##UD2L$MNrG   c                    | j                   j                         \  }}}|rUt        |   |||      }|t        j
                  j                  |       | t        j
                  j                  |       y y y rn   )r*  rx   r  __exit__ro   rp   rt   rF  )	rQ   abcliverI   maybe_prev_only_lift_cpu_tensorsr  r  s	           rH   rK  zFakeTensorMode.__exit__  sw       "		
 ,'"1a+C#/++,@A/;334TU < rG   c                     yrh   rF   r  s    rH   is_infra_modezFakeTensorMode.is_infra_mode  s    rG   c                    t        t        j                  t        j                  t	        t        j
                        t        t        j                              S )z8
        Query the state of the dispatch cache.
        )r  r%  r  r  ri   r  r  r  rR  s    rH   
cache_infozFakeTensorMode.cache_info  s?    
 !%%''../$$%	
 	
rG   c                    d| _         d| _        | j                  j                          | j                  j                          y)z+
        Clear the dispatch cache.
        r   N)r  r  r  clearr  rR  s    rH   cache_clearzFakeTensorMode.cache_clear  s3    
   "		rG   c                <   t         }	 t        | j                        }| j                  ||||      }|j	                         r%|j                  J |j                  j
                  }nt        j                  }|j                  |d      }	|	P| j                  ||	|||      }t        xj                  dz  c_
        | j                  r| j                  |||||       nk| j                  |||       | j                  ||||      }| j                  ||||||      }	|j!                          |	||<   t        xj"                  dz  c_        |t         u r| j                  ||||      }|S # t$        $ r/}
t        j&                  |
j(                  xx   dz  cc<   Y d}
~
Qd}
~
ww xY w)z
        Lookup a cache entry for the given arguments. If none exists, dispatch
        and cache the result (if the result is eligible for caching).
        Nr7   )_UNASSIGNEDr8   r   
_cache_keycache_on_shape_envfake_tensor_cacher%  r  r   _output_from_cache_entryr  r)  _crosscheck_cache_output_validate_cache_key_dispatch_impl_make_cache_entryr  r  r  r  r\   )rQ   ra   r  r  r  outputr  r   r  entryes              rH   _cached_dispatch_implz$FakeTensorMode._cached_dispatch_impl  s    %	9"4>>2E//%tV<C'')22299&,,IIc4(E 66ueS$PTU))Q.)00 11&$tVT((tV<,,T5$G..uc4vvV##%"c
++q0+ [ ((udFCF $ 	9))!((3q833	9s   D=E# #	F,%FFc                P   |t        j                         t         j                  j                         t        j                         | j
                  r| j
                  j                  ndg}|r| j                  |||       |r| j                  |||       t        t        |            S )z
        Create a cache key given the dispatch args. Raises _BypassDispatchCache
        for any situation that precludes caching.
        N)
ro   get_default_dtyperp   _get_default_deviceis_inference_mode_enabledr   settingsr  r  r  )rQ   r  ra   r  r  
key_valuess         rH   r[  zFakeTensorMode._cache_keyH  s      ##% HH((* ++- (,~~DNN##4

" $$Zu=$$Z? z!233rG   c                    t         j                  j                  |j                  v rt	        d      t         j                  j
                  |j                  v rt	        d      t         j                  j                  |j                  v rt	        d      |t        j                  j                  k(  rt	        d      || j                  v rt	        d      |j                         dk(  rt	        d      t         j                  j                  j                  |      st	        d      |j                  r\t         j                   j#                  |j                         t         j                   j$                  j&                        rt	        d      y	y	)
za
        Validate that the cache key generated by _cache_key will be
        reasonable.
        zdata dependent outputzdynamic output shapezinplace viewzunsafe viewliftzinductor::resize_storage_bytes_znon-builtinCompositeImplicitAutogradN)ro   Tagdata_dependent_outputtagsr  dynamic_output_shapeinplace_viewaten_unsafe_viewr  lift_fnsrC  _libraryutils
is_builtinis_viewrp   %_dispatch_has_kernel_for_dispatch_keyDispatchKeyro  )rQ   ra   r  r  s       rH   r`  z"FakeTensorMode._validate_cache_keyj  s'    99**dii7&'>??99))TYY6&'=>>99!!TYY.&~664$$,,,&}554== &v..99;;;&'HII~~##..t4&}55
 <<EHHJJIIK--GG
 ''BCC
<rG   c                   t        |t              rC| j                  ||j                         |       | j                  ||j	                         |       y|D ]y  }t        |t
              r| j                  |      st        d      |j                  t        d      t        |      rt        |j                   d      t        |j                         j                         t              rt        d      t        |      }|j                  || |       t        |t               rt        d      t        |t              r|j#                  ||       t        |t$        t&        f      rt        d      t        |t(        t*        t        f      r| j                  |||       O|j-                  t/        |             |j-                  |       | y)a7  
        Translate the provided args into a form suitable for caching at FakeTensor
        dispatch, i.e., convert unhashable types like lists & dicts into tuples and
        convert FakeTensors into metadata. Raises _BypassDispatchCache to signal
        unsupported cases that should bypass caching.
        Nznot our fakeconstant attributez tensorzsymbolic nbytesznon-fake tensorzsymbolic shape)ry   ri   r  keysvaluesr   r7  r  r   r#   r   r  r  r   r  r  r   r  r   r   r  r  rz   r   )rQ   r  r  r  r  r  s         rH   r  z"FakeTensorMode._prep_args_for_hash  sy    dD!$$VTYY[%@$$VT[[]EB 	#C#z*'',.~>><<+./CDD %.#**W/EFF c113::<fE./@AA237&&vtU;C(*+<==C(%%fc2C'8!45*+;<<C$t!45((e<
 d3i(c"7	#rG   c                   |t        ddd      S t        |t              st        d      |j                  t        d      |j
                  rt        d      t        |      rt        d      |j                         D ]$  }t        |      t        |      k(  st        d       t        t        |            D ],  }t        ||         t        |      k(  st        |dd      c S  d}	|j                  rCt        |      D 
cg c]  \  }
}t        |t              s|
 }}
}t        |      dk(  sJ |d	   }	t        |      }t        fd
|j                   D              |_        t        fd|j"                  D              |_        j%                  |j&                        |_        |j(                  dnj%                  |j(                        |_        t        d||	      }| j+                  ||||      }t,        j.                  j1                  |      }t,        j.                  j1                  |      }||k7  rt        d      |S c c}}
w )z
        Make a cache entry object for the given 'output' Tensor. Raises
        _BypassDispatchCache if the output tensor has characteristics that
        prevent caching it.
        N)r  r  r  znon-FakeTensor outputr  zsparse outputzsparse compressed outputzkwarg aliases outputr7   r   c              3  @   K   | ]  }j                  |        y wrn   convert_outputr   r   r  s     rH   r   z3FakeTensorMode._make_cache_entry.<locals>.<genexpr>  s     O1u33A6O   c              3  @   K   | ]  }j                  |        y wrn   r  r  s     rH   r   z3FakeTensorMode._make_cache_entry.<locals>.<genexpr>  s     QA 4 4Q 7Qr  zdispatch_key_set mismatch)r   ry   r   r  r   r   r$   r  r  ranger  r{  	enumerater   r  r  r  r  r  r  r  r^  ro   rp   _dispatch_key_set)rQ   r  r   ra   r  r  rc  kvalidxr  ir   idxsr  rd  synth_outputsynth_key_setkey_sets    `                rH   rb  z FakeTensorMode._make_cache_entry  s-    >&4$QUVV &*-&'>?? ??&&';<< &77'&'ABB MMO 	CD$x2f:%*+ABB	C
 T# 	C$s)}6
** #dT 	 <<"+D/K$!QZ65JAKDKt9>!>AwH*62OOOQQQ"'"6"6x7N7N"O %%- %%h&<&<= 	 $
 44UE3dS 22<@((,,V4G#&'BCCE Ls   I Ic           
         |j                   #||j                      }t        |t              sJ |S |j                  }|yt	        |      rJ 	 	 	 	 	 	 dfdt        fd|j                  D              }t        fd|j                  D              }	 |j                        }
|j                  dn |j                        }t        j                  }| j                  | j                  j                  }t        |       5   |       5  t        j                   ||	|j"                  |j$                  d|j&                        }ddd       ddd       |j(                  r t        j*                  j-                  d       |j.                  r t        j*                  j1                  d       |j2                  rw|t5        t6        |j8                           }t        |t              sJ |j;                         }t        |       5   |       5  j=                  ||
||	       ddd       ddd       t        | |j>                        S # 1 sw Y   xY w# 1 sw Y   	xY w# 1 sw Y   BxY w# 1 sw Y   FxY w)	z?
        Create a new FakeTensor from the cache entry.
        Nc                    t        | t              r*|j                  J | j                  |j                        S t        | t              rJ | S rn   )ry   r:   r   r  r9   )r  r  r   s     rH   check_valuez<FakeTensorMode._output_from_cache_entry.<locals>.check_value0  sJ     %!23222}}S%//::%e_===rG   c              3  0   K   | ]  } |        y wrn   rF   r   r   r  r  s     rH   r   z:FakeTensorMode._output_from_cache_entry.<locals>.<genexpr>:  s     Dk!U+D   c              3  0   K   | ]  } |        y wrn   rF   r  s     rH   r   z:FakeTensorMode._output_from_cache_entry.<locals>.<genexpr>;  s     F{1e,Fr  r"  )r  r   r  r  T)r  r  r  r8   rW   zUnion[IntLikeType]) r  ry   r   r  r#   r  r  r  r  r  
contextlibnullcontextr   suppress_guardsr>  ro   empty_stridedr  r   r  r  rp   	_set_conjr  _set_negr{  r   r  r  r  set_r  )rQ   r  rd  r   ra   r  inplace_argr  r  r  r  r  maybe_suppressr2  view_argstorager  s    ` `            @rH   r^  z'FakeTensorMode._output_from_cache_entry  sB    (u001Kk:666 >> ***	#	,:		 DX^^DDFhooFF$X%<%<eD %%- X33U; 	 ?I>T>T>>%!^^;;N)$/ 	1A 	''nn&44E	 	 HHud+??HHeT*<<Denn56Hh
333..0G-d3 C^5E C

7NE6BC C $x77/	 	 	 	(C C C CsH   
I:II
I4I('I4I	II%(I1	-I44I=c                
   	 | j                  ||||      }	 ||t        t        ||       y|J |J y# t        $ r}t        d| d| d| d|       |d}~ww xY w# t        $ r}t        d| d| d|       |d}~ww xY w)z
        Helper to validate that the output synthesized from the cache matches
        the output created by normal dispatch.
        z*FakeTensor cache crosscheck failure: func=z, args=z	, kwargs=z: Dispatch raised=N)ra  	Exceptionr  r"   r!   )rQ   rc  ra   r  r  r  true_outputre  s           rH   r_  z'FakeTensorMode._crosscheck_cache_output`  s    	--dE4HK
	'f.@"9k6B"***~%~  	<TF CvYvh.@E 	  	<TF CvYvh0 	s0   6 A A 	AAA	B(A==Bc                   |xs i }t               5  t        j                  d|||       d d d        |t        v rt        |   |      S t        j	                         t
        j                  k  r(t        j                  ddt        z  |       t               }|t        v rt        |       5   ||i |cd d d        S | j                  r| j                  ||||      S | j                  ||||      S # 1 sw Y   xY w# 1 sw Y   IxY w)Nz%s %s %sz'%sFakeTensorMode.__torch_dispatch__: %s )r.   r  rM  r  getEffectiveLevelloggingDEBUGrO   rK   _DISPATCH_HANDLE_DIRECTLYr>  r'  rf  ra  )rQ   ra   r  r  r  incrs         rH   rA  zFakeTensorMode.dispatch  s     2] 	6IIj$f5	6 ***4066  "gmm3II93;PRV +,D ,,-d3 -T,V,- - --dE4HH&&tUD&AA/	6 	6"- -s   C.(C:.C7:Dc           
        &'()*+,-./ t         j                  ||f      \  *&t        *      }|r@*D cg c]  }t        |      st	        |       }}t        j                  d|       t        S *D cg c]  } j                  |      s| }	}t        d |	D              xs t        d *D              , j                  ( j                  v }
|
r|	rt              rƉ,s|	st        d |	D              s
J  d       *D cg c]!  } j                  |      r|j                  n|# }}t         j                  |&      \  }} |i |}t	        |      t         u rH j#                  |      r7t%               5  |j'                         }d d d        (j)                   |d      S |
rQt+        |      dk(  rt+        |      d	k(  sJ | d
|        t	        |d         t         u r(j)                   |d         S d} j,                  rrt.        j0                  j2                  j4                  j6                  k(  rd|v r
|d   dk7  rd}t.        j0                  j8                  j:                  j6                  k(  rd} j=                  (*&      \  *}	~~t        d |	D              }t.        j>                  j@                  jB                  vrt.        j>                  jD                  jB                  vr|rt+        |	      dk7  r,s|s*D cg c]!  } j                  |      r|j                  n|# }}t         j                  |&      \  }}t%               5   |i |}d d d        t         jG                        }|D cg c]  }tI        |t               s| }}t         fd|D              }|rt         jK                  t         ( fd|      S |D ]  }(jM                  |        t         j                  *&      \  }} jO                  |	||       d  fd}ddl(m)'m*+ tW               --. jX                  r_t        d |	D              rMt        + /fd*D              s7*D cg c]
  } ||       }}t         j                  |&      \  }} |i |.nG jX                  r;tZ        j                  d|	* j\                  r j\                  j^                  nd        d!'+-. fd},r-ta               jc                        }| | | g|i |      S ddl2m3} |vr ji                        sqddl2m5} |v r6,stm              r)t        d |	D              r 5   |   |i |cd d d        S  5   jn                  |i |}|t        ur|cd d d        S 	 d d d        djp                  jr                  v r;tu        d      r/tw              s$ 5   | jx                  |i |      cd d d        S t.        jz                  j|                  j~                  j                  js                               j                  j                  }|r|t.        jz                  j                  j                         )t.        jz                  j                  j                  )fd      5   5   ||i |} ||      cd d d        cd d d        S t        D ]-  \  } }! |       s |! g|i |}"|"t        us% ||"      c S  	 d"	 	 	 d#&*, fd}#t              s |#       }$ ||$      S 	 t               5   |i |}d d d         | j                  *|jc                  d                  S c c}w c c}w c c}w # 1 sw Y   xY wc c}w # 1 sw Y   7xY wc c}w c c}w # 1 sw Y   ?xY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   nxY wd d d        # 1 sw Y   %xY w# 1 sw Y   xY w# t        $ r}% |#|%      cY d }%~%S d }%~%wt        $ r tZ        j                  d        w xY w)$Nz,FakeTensorMode unrecognized subclass(es): %sc              3  4   K   | ]  }|j                     y wrn   )r  )r   r  s     rH   r   z0FakeTensorMode._dispatch_impl.<locals>.<genexpr>  s      !
./A))!
   c              3  <   K   | ]  }t        |t                y wrn   )ry   r   )r   rL  s     rH   r   z0FakeTensorMode._dispatch_impl.<locals>.<genexpr>  s     :1Av&:s   c              3  8   K   | ]  }|j                   d u  y wrn   r   )r   r   s     rH   r   z0FakeTensorMode._dispatch_impl.<locals>.<genexpr>  s      +,

$&   z. should not have fake inputs without constantsTr   r   r7   r  Fr  r   c              3  8   K   | ]  }|j                   d u  y wrn   r   r   re  s     rH   r   z0FakeTensorMode._dispatch_impl.<locals>.<genexpr>       Qa1::T1Qr  c              3  @   K   | ]  }j                  |        y wrn   )may_turn_const)r   r   rQ   s     rH   r   z0FakeTensorMode._dispatch_impl.<locals>.<genexpr>"  s     P!t2215Pr  c                ,    j                  | d      S )NTr  )r   )r   	converterrQ   s    rH   <lambda>z/FakeTensorMode._dispatch_impl.<locals>.<lambda>'  s    i88qPT8U rG   c                f   t        | t              r| j                  S t        | t              rj                  J | j
                  j                  | j
                  j                  j                  j                  j                        j                  j                  j                              S | S rn   )ry   r   rq  r,   r   nodepytypeexprxreplace
var_to_valunbacked_var_to_val)r   rQ   s    rH   maybe_to_real_tensorz;FakeTensorMode._dispatch_impl.<locals>.maybe_to_real_tensor5  s    !Z(}}$A|,~~111vv}}FFKK(()B)BCLL::  rG   )compute_unbacked_bindingsfree_unbacked_symbolsc              3  8   K   | ]  }|j                   d u  y wrn   )rq  r  s     rH   r   z0FakeTensorMode._dispatch_impl.<locals>.<genexpr>L  s     M!AMM-Mr  c              3     K   | ]F  }t        |t              xr0  |      xxr$ j                  d uxr t        fdD               H y w)Nc              3  N   K   | ]  }|j                   j                  v  y wrn   )r   r  )r   srQ   s     rH   r   z:FakeTensorMode._dispatch_impl.<locals>.<genexpr>.<genexpr>S  s      VAAT^^%G%GGVs   "%)ry   r   r   r   )r   rL  r  rQ   symss     rH   r   z0FakeTensorMode._dispatch_impl.<locals>.<genexpr>N  sd        q&) W!6q!99Wd2W VQUVVWs   AAz,propagate_real_tensors skipped %s(%s, %s) %sc                n    dd l dfd}ur"t        ||         j                  | d       | S )Nr   c                <   t        | t              r|| _        y t        | t              rs |       rjt        | j                  j
                  j                        r?j                  J j                  j                  | j                  j
                  |       y y y y rn   )	ry   r   rq  r,   r  r  Symbolr   set_unbacked_var_to_val)r   real_tr  rQ   sympys     rH   gozOFakeTensorMode._dispatch_impl.<locals>.maybe_propagate_real_tensors.<locals>.gom  sv    a,$*AM<05J15M!!&&++u||<#~~999>>qvv{{FS = 6N0rG   T)peek)r   r  r  r   rW   rX   )r  r3   r   )fake_outr  r  r  r  nilreal_outrQ   s     @rH   maybe_propagate_real_tensorszCFakeTensorMode._dispatch_impl.<locals>.maybe_propagate_real_tensorsj  s9    T s""h1
 *$..(NOrG   )
meta_tabler   c              3  4   K   | ]  }t        |         y wrn   )r#   r  s     rH   r   z0FakeTensorMode._dispatch_impl.<locals>.<genexpr>  s     PQa 00Pr  zprims::prim_meta_implc                      S rn   rF   )ctxs   rH   r  z/FakeTensorMode._dispatch_impl.<locals>.<lambda>  s     rG   c                    t         j                  j                  j                        ry sj	                        st              | t              } t        |       S rn   )ro   rx  ry  can_generate_trivial_fake_implcan_run_unsafe_fallbackrf   run_fallback_kernel)error	args_specr  ra   has_symbolic_sizesrQ   s    rH   maybe_run_unsafe_fallbackz@FakeTensorMode._dispatch_impl.<locals>.maybe_run_unsafe_fallback  s]     ~~##BB4H!)E)Ed)K2488}4T:&tT9iOOrG   z*failed while attempting to run meta for %sr-  )r   rI   rW   zOptional[Union[T, Tensor]])r  rI   rW   rI   rn   )r  zOptional[RuntimeError]rW   r$  )Lr   tree_flatten_check_for_subclass_check_for_subclass_argr   r  rM  r  r7  r   r!  rw  rE  r   r   tree_unflattenr   r  r.   cloner   r  r:  ro   r  ru  _to_copyr  prims
device_put%validate_and_convert_non_fake_tensorsrp  nondeterministic_seededrr  rt  r   ry   tree_map_onlyr   invalidate_written_to_constantsr  r  r  r  r   r  r   r  get_fast_op_implsr   r   r  cpp_meta_supports_symintr   r   	decomposer   rC  hasattrstride_incorrect_opr  rx  simple_registry	singletonfind	fake_implkernelFakeImplCtxset_ctx_getterop_implementations_checkshas_metar>  rv  r  rC  +wrap_meta_outputs_with_default_device_logic)0rQ   ra   r  r  r  has_unrecognized_typesr   r  r   flat_arg_fake_tensorsis_lift_funcrL  const_flat_args
const_argsconst_kwargsr  avoiding_device_initall_constantflat_outflat_out_tensorsr   r  real_flat_args	real_argsreal_kwargsr  	fast_implr  r   rh  maybe_fake_implr  run_impl_checkop_implop_impl_outr  fallbacknot_implemented_errorr  r  r  r  r  r  r  r  r  r  s0   ``                                    @@@@@@@@@@rH   ra  zFakeTensorMode._dispatch_impl  sl     &22D&>B	9 "5Y!?!!*".Ea.HQ" "  %%>@R "!,5 Mq9I9I!9L M M  !
3H!
 
 ;:	:: 	 ..	t}}, !6+D1&) 0E  GEFG  CL=>d..q1

q8O  (.'<'<_i'X$J
3l3CCyF"t':':3'? !] &))+C& 11$41PP v;!#D	QJ4&&8JJ6DG}& 11$Q@@  %!!		//777&8$-'+$uyy11999'+$ .2-W-W)Y	.
*) & Q;PQQII--TYY>		&&dii7)*a/&( CL=>d..q1

q8O  (.'<'<_i'X$J  8J7,78 ))#.H+3Maz!V7LMMP?OPPL++U  ( ;55c:;
 ,,Y	Bf,,T3H$PVW		

 h''M7LMM  #  @II!215INI%+%:%:>9%U"I{Y6+6H(( II>%6:nn22$	 	. )+//5I$3Id4TT4TV4TUU 	;z!$*G*G*M9 **" 05P:OPP  F4.t4dEfEF F  "DNND3F3N*	  + ***./'- 3'D''88   ..88BBGGIIK

)FF 	 ..**66tTBC))88E <t <($9&93F;< < < (A 	E#NGd#%dDB4B6Bn47DD		E -1	P)	P!	P 	P" ~02H/99
	-d3 *$)&)* ,<<46::h+? = 
 	
c	" !N2& &n8 8  Nn J|F F " < < < < <L* *" 	D,-BCC 	MMFM	s   ]5	]58]:]:&]?=^&^	^^#^#*^(#^-:^:_)_*,_<	_*/` :	_7` ^^ -^7:___	_**_47` <` 	`=``=$`=
debugprimsr  ru  xlavision	torchtext
torchaudio	quantizedc                v    | j                   sy|j                  | j                  v xs |j                         dk(  S )NFzfbgemm::gmm)r  	namespace+_can_run_unsafe_fallback_allowed_namespacesrC  rQ   ra   s     rH   r  z&FakeTensorMode.can_run_unsafe_fallback  s:    **
 NNdNNN ,yy{m+	
rG   c                d     g d fd}D cg c]
  } ||       }}|fS c c}w )z
        Checks if the list of tensors are fake tensors.
        If not, try to convert them to fake tensors.
        Returns the original args, kwargs, and a flattened list of (args, kwargs) that are fake tensors.
        c                   t        | t              s| S 	j                  |       st        j                  j
                  j                  v r2t        j                        \  }}t        dt        ||             	j                  s[t        | t              r| j                  	urt        d      t        j                        \  }}t        dt        ||             j                  	|       }n| }j                  |       |S )NzECan't call metadata mutating ops on non-Fake Tensor inputs. Found in zMixing fake modes NYIzuPlease convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in )ry   r   r7  ro   rp  rt  rr  r   r  AssertionErrorr&   r  r   r   r   rz   )
r   r  r  r  r  r  r  r  ra   rQ   s
       rH   validatezFFakeTensorMode.validate_and_convert_non_fake_tensors.<locals>.validate  s   a( ##A&99))TYY6#)#8#8I#NLD&(_`klprvx~`  `A  B  11!!Z0Q[[5L,-DEE#)#8#8I#NLD&(BBMdTXZ`BaAbd 
  00q9!((-JrG   )r   rI   rW   Union[T, FakeTensor]rF   )	rQ   ra   r  r  r  r"  rL  validated_argsr  s	   `````   @rH   r  z4FakeTensorMode.validate_and_convert_non_fake_tensors  s?     35	 	6 099!(1+99444 :s   -c                `      j                   d dd fd}t        ||      S )NFc                R    t         t              s S t        j                        \  j	                         }|r7t        j                   j                  k(   fd       t        t               S +rj                         S j                   xs       S  S )Nc                 (    dj                    d  S )Nz-FakeTensor is wrapped to wrong device, found z, expected r-  )r  re  s   rH   r  zZFakeTensorMode.wrap_meta_outputs_with_default_device_logic.<locals>.wrap.<locals>.<lambda>O  s    KAHH:U`an`op rG   )ry   r   r   r  r7  ro   _checkr  r   rI   r   r#  )	re  r7  r  r  r  r  ra   r  rQ   s	   ` rH   wrapzHFakeTensorMode.wrap_meta_outputs_with_default_device_logic.<locals>.wrap>  s     a($ 224C!* **1-KHH-p Aqz!&) %55dA>>$99a!8=  rG   re  rI   rW   r#  )r!  r2   )	rQ   rh  ra   r  r  r)  r  r  r  s	   ` ``` @@@rH   r   z:FakeTensorMode.wrap_meta_outputs_with_default_device_logic1  s7     ..	 !&	 	B a  rG   r  c                  dd l }|4| j                  }| j                  sJ d       | xj                  dz  c_        |j                  j	                  |d      }|j
                  j                  j                  d      }| j                  J | j                  j                  | j                  j                  ||      ||      }|S )Nr   z1should only called while FakeTensorMode is activer7   intermediate_offsets_or_lengths)valr   )symr   r   )r.  r  r*  rp   _get_nested_intr%  r   EphemeralSourcer   r  create_symbol)rQ   r  ro   r   srcrets         rH   r  z)FakeTensorMode.create_symbolic_nested_inta  s    
 	444L##X%XX#%%*%xx''a8mm""223TU~~)))nn..,, -   / 
 
rG   c                l    t         j                  j                  |j                  v ry|| j                  v S rh   )ro   rp  	view_copyrr  _cpp_meta_supports_symintr  s     rH   r  z'FakeTensorMode.cpp_meta_supports_symint  s,    99$))+t5555rG   c                    |j                         t        k  xr< t        |       xr. | j                  |       xr |j                  j
                  dk(   S rs  )numelCONSTANT_NUMEL_LIMITr#   r7  r  r   r6  s     rH   r  zFakeTensorMode.may_turn_const  sV    GGI-- ,!!$$,$$Q'', HHMMV++		
rG   c                   t        d |D              }t        |      }|r|j                         rt        |||d      \  }}|j	                         D ]u  \  }	}
|	dk7  s|j                  |	      r|	nd}	| j                  |
      s2|j                  |	      sD|
j                  Q| j                  j                  |
j                         w y y y )Nc              3  8   K   | ]  }|j                   d u  y wrn   r   r  s     rH   r   zAFakeTensorMode.invalidate_written_to_constants.<locals>.<genexpr>  r  r  T)r  r  normalize_to_only_use_kwargsinputrQ   )
r   r   
is_mutabler(   rk   has_argumentr7  r   r!  r   )rQ   ra   r  r  r  any_constantschema_infor   
new_kwargskr   s              rH   r  z.FakeTensorMode.invalidate_written_to_constants  s     Q;PQQ%d+K224.4TMAz #((* W1w,+*B*B1*EAF$$Q'#..q1

...JJ1::VW	 5<rG   )r  r   r   r   c                   | j                   }|| j                  }|r|J d       d }| j                  j                  | |||||      S )Nz2cannot set both static_shapes and symbolic_context)r   r   r   r   )r   r  r!  r   )rQ   r   r  r   r   r   r   s          rH   r  zFakeTensorMode.from_tensor  sn     )-	  ..M (DCD(I))::- ; 
 	
rG   )r  r   r  r   r   r&  r  r  r   r   rW   rX   rV   )r   r  rW   zTypeGuard[FakeTensor])rW   r   )rW   r[   r  )rW   r   )rL  zOptional[Type[BaseException]]rM  zOptional[BaseException]rN  zOptional[TracebackType]rW   rX   )rW   r  )
r  r8   ra   r=   r  r  r  r  rW   r  )ra   r=   r  r  r  r  rW   rX   )r  r  r  z?Union[Mapping[str, object], Sequence[object], Iterable[object]]r  r8   rW   rX   )r  r8   r   r  ra   r=   r  r  r  r  rc  r$  rW   r   )r  r8   rd  r   r   r  ra   r=   r  r  rW   r$  )rc  r$  ra   r=   r  r  r  r  r  r  rW   rX   )
ra   r=   r  r  r  r  r  r  rW   r$  ra   r=   rW   r   )
ra   r=   r  r   r  r  r  r4   rW   z%Tuple[List[object], List[FakeTensor]])
rh  r  ra   r=   r  r  r  r)  rW   r1   )r  r  rW   r  r  )
ra   r=   r  zSequence[FakeTensor]r  r  r  r  rW   rX   )r   r   r  r  r   r'  r   r(  r   r   rW   r   )DrC   rD   rE   r  r^   r  r  r   r  r  re  r7  r  r  rR   r4  r7  r*  r:  r?  r5   r'   r  rG  rK  r  rS  rU  rX  rf  r[  r`  r  rb  r^  r_  rA  ra  rl   r  r  r  r   r  ru  r2  r  r  r  as_strided_scatter
as_stridedas_strided_r3  detachview_as_realview_as_complexr  source_Storage_storage_offset(_sparse_coo_tensor_with_dims_and_tensorsr6  r  
lift_freshlift_fresh_copyrw  r  r  r  r  r  s   @rH   r%  r%  -  s   :<E7<JL#%0%5NN5 E3N!&$&!! !##"&((
 (,&+(,(, !]D !%]D  $	]D
 &]D &]D  !]D" 
#]D~DA - -  
 
 "$'5'7  	
 % 
 $(V(V #V #	V
 
V&   	
 	
  )) ) 	)
 %) 
)V 4 4  4 	 4
 % 4 
 4D(D(D (D %	(D
 
(DT,#,# N,# 	,#
 
,#\TT T 	T
 T %T %T 
TlF8F8 #F8 	F8
 F8 F8 
F8P$  	
  % 
F "$'5'7BB B 	B
 %B 
BBH
H
 H
 	H

 %H
 
H
^
 3>	3/	
*5*5 '*5 $	*5
 *5 
/*5X.!.! .! $	.!
 .! 
.!b 04,	2 !,

  ""''  

!!$$		//55==!6
 4??22D4H4H4P4PQH
WW  4W 	W
 %W 
W4 )-#'6:

 &	

 !
 4
 
 

rG   c                    t         j                  j                  |j                  v ri t	               5  d fd}|D cg c]
  } ||       }}t
        j                  ||      \  }} ||i |}	d d d        t               |D ]H  }
t        |
t              st        |
      r j                  |
j                         j                         J d fd}t
        j                  |	      S c c}w # 1 sw Y   xY w)Nc                    j                  |       r\t        j                  | | j                        }| j                  r|j                  | j                                | t        |      <   |S | S )Nr-  )r7  ro   
zeros_likero  r   _coalesced_r  r  )re  r  r   	inp_implss     rH   to_real_tensorz+run_fallback_kernel.<locals>.to_real_tensor  sY    $$Q'&&q?;;OOANN$45%&	"S'"
HrG   c                "   t        |       vr9t        | t              r)t        |       s| j	                         j
                  v rt        | t              r7t        |       v rt        |          S j                  j                  |       S | S rn   )r  ry   r   r#   r   _cdatar!  r   )re  r   rT  orig_not_implemented_exceptionstoragess    rH   map_outz$run_fallback_kernel.<locals>.map_out  s    a5	!q&!!!$  "))X500a !u	! A'' 66GG	STUUHrG   )re  rI   rW   zUnion[T, Tensor]r*  )ro   rp  rt  rr  r.   r   r  setry   r   r#   addr   rW  r2   )r   ra   r  r  rX  rU  rL  r  r  rh  re  rZ  rT  rY  s   `   `       @@rH   r  r    s     yy*,,I 
 "	 1::1^A&:	:,,Y	Bf$!&!"  &)UH 8a  #Q--/6678   ??7A&&C ;" "s   C:C5#C:5C::Dc                  4    e Zd ZddZ	 	 d	 	 	 	 	 	 	 	 	 ddZy)FakeCopyModec                    || _         y rn   )r   )rQ   r   s     rH   rR   zFakeCopyMode.__init__	  s	    "rG   Nc                   |r|ni }|t         j                  j                  j                  k(  r<t	        |d   t
              sJ  || j                  j                  |d   d      fi |S |t
        j                  k(  rt        |      dk(  rt        |      dk(  sJ t        t
        |d         }t        t        t        t        f   |d         }t        |      |v r|t        |         S | j                  j                  |d      }||t        |      <   |S t         j                  j                         5   ||i |cd d d        S # 1 sw Y   y xY w)Nr   T)r     r7   )ro   rp   
TensorBaser  ry   r   r   r  __deepcopy__r  r   r	   r  r   r  DisableTorchFunctionSubclass)rQ   ra   r  r  r  r   memor  s           rH   __torch_function__zFakeCopyMode.__torch_function__	  s9    "r 588&&,,,d1gv...**47$*GKQ  V(((t9>c&kQ&666&$q'*FS*_-tAw7D&zT!BvJ''..,,V4,HC"DFJ668 -T,V,- - -s   4EE)r   r%  rW   rX   )rF   N)
ra   r=   r  r  r  r  r  zOptional[Mapping[str, object]]rW   r   )rC   rD   rE   rR   rf  rF   rG   rH   r^  r^  	  sD    # "$15-- - 	-
 /- 
-rG   r^  c                    t        |       dk(  rt        | d   t              sJ | d   j                  j                  rt        j                  d      S | d   j                  S )Nr7   r   r"  )r  ry   r   r   r7  ro   r  ro  r  s    rH   _device_handlerri  1	  sT     t9>ja*===Aw--||F##Aw"""rG   c                &    t        d | D              S )Nc              3  2   K   | ]  }t        |        y wrn   )r  r   s     rH   r   z&_check_for_subclass.<locals>.<genexpr>J	  s     =a&q)=r   )r   )r  s    rH   r  r  I	  s    =9===rG   c                    t        | t               xrJ t        | t              xr8 t        |       t        uxr% t        |       t        j
                  j                  uS rn   )ry   r   r   r   ro   r  r  )r   s    rH   r  r  M	  sQ    q*%% 	.q&!	.G6!	. G588---	rG   c                d    t        d t        t        | d         j                         D              S )Nc              3  2   K   | ]  }t        |        y wrn   r  r   r  s     rH   r   z<lambda>.<locals>.<genexpr>X	  s      4A4r   r   )r  r   r   r  rh  s    rH   r  r  X	  s,    e 4VT!W-2244 / rG   c                d    t        d t        t        | d         j                         D              S )Nc              3  2   K   | ]  }t        |        y wrn   ro  rp  s     rH   r   z<lambda>.<locals>.<genexpr>[	  s      6A6r   r   )r  r   r   r  rh  s    rH   r  r  [	  s,     6VT!W-4466 1 rG   c                V    t        t        t        | d         j                               S )Nr   )r  r   r   r  rh  s    rH   r  r  ^	  s"    VT!W,,.9 rG   )_device_not_kwarg_ops_is_tensor_constructor_like_tensor_constructorscontains_tensor_typesr  r  r  r  c                    t         j                  d       t         j                  dt        j                         t         j                  dt        j                         t        j
                  } | rjt         j                  d       t        d | D              }t        | j                         d       D ]#  \  }}t         j                  d|d	z   | d
|       % y y )NzFakeTensor cache stats:z  cache_hits: %sz  cache_misses: %sz  cache_bypasses:c              3  2   K   | ]  }t        |        y wrn   )r  )r   rC  s     rH   r   z#dump_cache_stats.<locals>.<genexpr>}	  s     -qCF-r   c                    | d    S rM   rF   )r  s    rH   r  z"dump_cache_stats.<locals>.<lambda>~	  s    AaD5 rG   )r   z    %-*s %sr7   r  )	r  infor%  r  r  r  maxsortedrk   )r  widthrC  r   s       rH   dump_cache_statsr  u	  s    HH&'HH!:!:;HH!>#>#>?,,H$%-H--8>>+A 	;DAqHH]EAI!Aw:	; rG   )rk   rI   rW   zDict[T, Literal[True]])rW   z2Generator[Optional[TorchDispatchMode], None, None])r~   r   rW   zList[Tensor])r   r  rW   zTypeGuard[Tensor])r   r  rW   zOptional[FakeTensorMode])ra   r=   rW   ztorch._C._SchemaInforE  )r   zType[T]r   r1   rW   zList[T])r   r  rW   r   rV   )r   r%  rW   zGenerator[None, None, None])r   r   rW   r  )r   r%  ra   r=   r  r  r  r1   rX  r  rW   r   )r  r  rW   r)  )r  r  rW   r   )r   r  rW   r   )
__future__r   atexitr  r  	functoolsr  r  rK  r+  typingr   collectionsr   r   r   r   r   r	   r
   r   r   r   r   r   r   r   r   r   r   r   r   typing_extensionsr   r   r   ro   r   r   r   r   torch._C._functorchr   r   torch._prims_commonr    torch._subclasses.meta_utilsr!   r"   r#   r$   r%   torch._utilsr&   torch.fx.immutable_collectionsr'   torch.fx.operator_schemasr(    torch.multiprocessing.reductionsr)   torch.overridesr*   torch.typesr+   r,   torch.utils._backport_slotsr-   torch.utils._mode_utilsr.   torch.utils._python_dispatchr/   r0   torch.utils._pytreer1   r2   r3   r4   torch.utils._statsr5   torch.utils._tracebackr6   _fake_tensor_utilsr8   r9   r:   r  r;   torch._guardsr<   
torch._opsr=   r  r>   r?   	getLoggerrC   r  _logginggetArtifactLoggerr  
ValueErrorre  r[   rB   rZ  DimListry  _pytreer   rI   _opsr  ru  r9  rO   rK   r  rZ   r`   rd   rf   rl   contextmanagerrv   r   r   r   	lru_cacher   r   r   r   r   r4  r>  rE  rG  rO  r   r  r  r  r  r   r  r  r  r%  r  _StoragePointerr  r^  ri  r  r  r  r  r  r  r  r  r  r  r  r  r  torch._subclasses.fake_implsrt  ru  rv  rw  r  r  r  r  registerr  rF   rG   rH   <module>r     s
   "       	    # !    & . !  3 3 T 5  % 9 8 ; - 1 7 / F E $ 4 R R #$%Og!..::8EVW	 	 m
			CLzz~~   \   ,   <   <  & - -$(* T. . T	D 	D@
$W Wt T
' 
' @@ @ @,B B7E 7Etw4 w4t	 &79LLM  
(% (%  (%V#L 
$ $  $4 
$    
$9    
$   (X
& X
v$ ='='
='  =' 	='
 %1=' ='D -$  -F#0> 
IINN!!?	IINN " 
IINN!! $ 
IINN!!)) ,  (	IINN''	IINN$$	IINN%% 	 	 	 	; 	;KI  )SV3- 	.s   &R< <SSS