
    9j                       U d 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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mZ ddlmZmZmZmZ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$m%Z%m&Z& dd	l'm(Z(m)Z)m*Z*m+Z+m,Z, dd
l-m.Z. ddl/m0Z0 ddl1m2Z2 ddl3m4Z4m5Z5m6Z6m7Z7 ddl8m9Z9 ddl:m;Z;m<Z<m=Z= ddl>m?Z?m@Z@mAZAmBZB ddlCmDZDmEZE ddlFmGZG ddlHmIZI ddlJmKZK ddlLmMZMmNZN ddlOmPZPmQZQmRZRmSZSmTZTmUZUmVZVmWZWmXZX ddlYmZZZm[Z[ ddl\m]Z] ddl^m_Z_ ddl`maZambZb ddlcmdZd ddlemfZf ddlgmhZhmiZimjZjmkZkmlZl ddlmmnZn dd lompZpmqZqmrZr dd!lsmtZtmuZumvZv dd"lwmxZxmyZymzZzm{Z{m|Z| dd#l}m~Z~ dd$lmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZ dd%lmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZ d&d'lmZmZmZmZmZmZmZ d&d(lmZ d&d)lmZ d&d*lmZmZmZmZmZmZ d&d+lmZmZmZmZ d&d,lmZ d&d-lmZmZmZmZmZmZmZmZmZmZmZ d&d.lmZmZ d&d/lmZmZ d&d0lmZmZ d&d1lmZmZmZmZmZmZmZmZ d&d2lmZmZmZmZmZmZmZ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&d3lmZmZmZ d&d4lmZ d&d5lmZmZ d&d6lmZ d&d7lmZmZmZm Z  d&d8l!m"Z"m#Z#m$Z$ d&d9l%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+ d&d:lm,Z,m-Z-m.Z.m/Z/ d&d;l0m1Z1m2Z2m3Z3 d&d<l4m5Z5m6Z6m7Z7m8Z8m9Z9m:Z:m;Z;m<Z<m=Z=m>Z>m?Z?m@Z@mAZAmBZB 	 ddlCZDerdd>lGmHZH dd?lImJZJmKZK  e
j                  eM      ZNej                  j                  eMd@      ZQddAlmRZR  eRdB      ZS eRdC      ZTeUZVdDeWdEeXfdFZY G dG dH      ZZej                   G dI dJ             Z\ G dK dLe\      Z] e^dM  e_e	      j                         D              Zae^eb   eFdN<    ec       Zdeceb   eFdO<   ej                  j                  j                  Zhej                  j                  j                  Zj G dP dQ      ZkdRedEefdSZldTedUe@dz  dEefdVZm	 	 ddWdXdYedZedz  d[endz  d\edEefd]ZodWdXdYedZedEdfd^Zp	 	 dd_eneS   dWdXdYedZedz  d[endz  d\edEeSfd`Zq	 dd_eneS   dWdXdYej                  j                  daej                  d[endz  d\edEeSfdbZu	 	 dd_eneS   dWdXdYej                  j                  dZedz  d[endz  d\edEeSfdcZvdZedWdXdYej                  j                  d\ewexef   d[endz  d_eneS   dEefddZydTeTdEeneT   dz  fdeZzd_edWdXdZed[endz  dEewexef   f
dfZ{d_eneS   dWdXdYej                  j                  dZed[endz  d\ewexef   dEeSfdgZ|dEeRfdhZ}da~ecex   dz  eFdi<   daebdz  eFdj<   dEecex   fdkZdlexdEeXfdmZdWdXdnexdoej                  dEezfdpZdaecex   dz  eFdq<   daebdz  eFdr<   dEecex   fdsZdlexdEeXfdtZ	 ddoedWdXduedveXdweXdEeVfdxZddydoedWdXduedz  dzeXd{edz  dEefd|ZddydoedWdXduedz  dzeXd{edz  dEefd}Z G d~ d      Zej                         e_         G d d      Zy# eE$ r dZDeeFd=<   Y w xY w)aq  
This module contains classes and utilities for building variable trackers in Dynamo.
Variable trackers are used to convert Python values into symbolic representations
that can be traced and transformed during graph capture.

The key classes are:

- VariableBuilder: Handles source-tracked objects that need guards and proper
  reconstruction in the output graph. Used for inputs, module attributes, etc.

- SourcelessBuilder: Handles ephemeral objects created during tracing that don't
  need source tracking or guards. Used for temporary lists, intermediate values, etc.

Variable trackers enable Dynamo to track the flow of values through the program,
maintain guards for dynamic properties, and reconstruct values in the output graph.
The builders in this module handle converting Python values into appropriate
VariableTracker instances based on their type and usage context.
    N)CallableMutableMapping)
ModuleType)Any
NamedTupleNoReturnoverloadTYPE_CHECKINGUnion)SymInt)enable_python_dispatcher)CURRENT_STREAM_INDEXget_external_object_by_indexregister_user_object)get_metrics_contextis_int_specialization_caseis_torch_symnormalize_count_iterset_feature_use)TracingContext)
flat_apply)call_torchbind)is_opaque_reference_typeis_opaque_typeis_opaque_value_typeshould_hoist)
OpaqueBase)HigherOrderOperator
OpOverloadOpOverloadPacket)
FakeTensorFakeTensorModeis_fakemaybe_get_fake_mode)is_sparse_any	safe_grad)justknobs_check)BackwardState)normalize_source_name)_DynamicScalar
DynamicInt)	_constrain_range_for_size_nested_int_aware_sort
DimDynamicRelaxedUnspecConstraintStatefulSymbolicContextSubclassSymbolicContextSymbolicContextSymIntSymbolicContextTrackedFake)immutable_dictimmutable_list)ExpandedWeight)
OrderedSet)is_traceable_wrapper_subclass"is_traceable_wrapper_subclass_type)ValueRanges)TensorWeakRef   )configgraph_break_hintsmutation_guardreplay_recordtrace_rules) get_registered_device_interfaces)InternalTorchDynamoErrorraise_observed_exceptionunimplemented)GuardBuilderinstall_guardmake_dupe_guard)auto_dynamic
auto_unsetFrameStateSizeEntryInferStrideprocess_automatic_dynamic)SideEffects)AttrProxySource
AttrSourceCallMethodItemSourceChainedSourceConstDictKeySourceConvertIntSourceCurrentStreamSourceDictGetItemSourceDictSubclassGetItemSourceDynamicScalarSourceFloatTensorSourceGetItemSource
GradSourceis_constant_sourceis_from_closure_sourceis_from_global_sourceis_from_nonlocal_sourceis_from_optimizer_source&is_from_unspecialized_nn_module_sourceListGetItemSourceLocalSourceNonSerializableSetGetItemSourceNumpyTensorSourceOptimizerSourceRandomValueSourceSkipGuardSourceSourceSubclassAttrListSourceTupleIteratorGetItemSource"UnspecializedBuiltinNNModuleSourceUnspecializedNNModuleSource) _extract_tensor_dictbuild_checkpoint_variablebuild_invoke_subgraph_variableclone_inputcommon_constant_types	dict_keys"enumerate_items_with_dict_positionget_fake_valueget_locals_to_stealget_static_address_typeis_frozen_dataclassis_functionis_function_or_wrapperis_invoke_subgraphis_lru_cache_wrapped_functionis_namedtupleis_parameter_freezing	is_typingis_utils_checkpointis_wrapper_or_member_descriptoristypenamedtuple_fieldsodict_valuesproxy_args_kwargsrange_iteratorset_example_valuetensor_always_has_static_shapetuple_iteratortuple_iterator_getitemtuple_iterator_len unwrap_with_attr_name_if_wrapperwrap_fake_exception   )AttributeMutationExistingAttributeMutationNewtypestrValueMutationExistingValueMutationNewVariableTrackerVariableTrackerMeta)BuiltinVariableConstantVariable)AutocastModeVariableCudagraphOverrideVariableDynamoConfigPatchVariableErrorOnGraphBreakVariableNullContextVariablePreserveVersionContextVariable)ConstDictVariableDefaultDictVariableMappingProxyVariableSetVariable)WorldMetaClassVariable)BuiltinMethodVariableCollectionsNamedTupleFunction!CollectiveFunctionRewriteVariable'CreateTMADescriptorExperimentalVariable!CreateTMADescriptorStableVariableFunctoolsPartialVariableSysFunctionVariableTritonKernelVariableTritonSetAllocatorVariableUserFunctionVariableWrapperUserFunctionVariable)"LocalMapWrappedHigherOrderVariable TorchHigherOrderOperatorVariable)CountIteratorVariableItertoolsVariable)LazyConstantVariableLazyVariableTracker)BaseListVariableListIteratorVariableListVariableRangeVariableSizeVariableSliceVariableTupleIteratorVariableTupleVariable)AutogradEngineVariableAutogradFunctionContextVariableAutogradFunctionVariableComptimeVariableConstantLikeVariableDebuggingVariableDelayGraphBreakVariableGetAttrVariableGetSetDescriptorVariableIgnoredFunctionVariableLambdaVariableLoggingLoggerVariableMethodWrapperVariableNumpyDTypeVariableNumpyVariableObjectVariablePythonModuleVariableRandomClassVariableRandomVariableSavedTensorBoxStringFormatVariableTorchVersionVariableTypingVariableWeakRefVariable)FSDPManagedNNModuleVariable$UnspecializedBuiltinNNModuleVariableUnspecializedNNModuleVariable)OptimizerVariable)OpaqueObjectClassVariableTorchScriptObjectVariableSDPAParamsVariable)DictKeySetVariableFrozensetVariableOrderedSetClassVariableOrderedSetVariable)EventVariableStreamContextVariableStreamVariable)NumpyNdarrayVariable$supported_const_comparison_op_valuesSymNodeVariableTensorSubclassVariableTensorVariableUnspecializedPythonVariable)DispatchKeySetVariableFuncTorchInterpreterVariableTorchCtxManagerClassVariableTorchInGraphFunctionVariable)TensorWithTFOverrideVariable#torch_function_mode_stack_state_mgrTorchFunctionModeVariable)FrozenDataClassVariableInspectVariableIntWrapperVariableKeyedJaggedTensorVariableMutableMappingVariableSourcelessGraphModuleVariableUserDefinedClassVariableUserDefinedConstantVariableUserDefinedDictVariable!UserDefinedExceptionClassVariableUserDefinedListVariableUserDefinedObjectVariableUserDefinedSetVariableUserDefinedTupleVariablenp)	PyCodegen)InstructionTranslatorInstructionTranslatorBasecudagraph_static_inputs)TypeVarVTTypeAliasTtreturnc                     t         j                  j                  t         j                  j                  j                        5  t        | d      cd d d        S # 1 sw Y   y xY w)Ngrad)torch_logginghide_warnings	_internalsafe_grad_filterhasattr)r  s    _/media/conek/DATA/Code/OCR/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builder.pysafe_has_gradr  W  sC    		%	%enn&>&>&O&O	P "q&!" " "s   AA!c                       e Zd Zy)_missingN__name__
__module____qualname__     r  r  r  \  s    r  r  c                   Z   e Zd ZU edz  ed<   eed<   eed<   ej                  j                  j                  dz  ed<   dZeed<   dZej                  ej                  z  dz  ed<   d	ed
eddfdZedej                  ej                  z  ez  dz  fd       ZddZddZddZdedefdZy)GraphArgNsource_examplepass_arg_as_tensorfake_tensorT	is_tensorexample_strong_refnamevaluer	  c                 2    t         j                  | ||       y N)object__setattr__)selfr$  r%  s      r  r)  zGraphArg.__setattr__  s    
 	4u-r  c                 z    t        | j                  t              r| j                         }|J |S | j                  S r'  )
isinstancer  r<   )r*  rs     r  examplezGraphArg.example  s4    dmm]3A= =H== r  c                     t        | j                  t        j                        r2t	        | j                        | _        t        | j                        sJ y y r'  )r,  r  r  Tensorr<   r#   r!  r*  s    r  __post_init__zGraphArg.__post_init__  s=    dmmU\\2)$--8DM4++,,, 3r  c                 (     || j                          y r'  r  r*  codegens     r  reconstructzGraphArg.reconstruct  s    r  c                      d | _         d | _        y r'  )r  r#  r1  s    r  erasezGraphArg.erase  s    "&r  otherc                     t        |t              sy| j                  |j                  d u S |j                  y| j                  j                  |j                  j                  k(  S )NF)r,  r  r  r$  )r*  r:  s     r  __eq__zGraphArg.__eq__  sT    %*;;<<4''||#;;##u||'8'888r  r	  Nr6  r  r	  N)r  r  r  rj   __annotations__r   boolr  _subclassesr!  r!   r"  r#  r0  r   strr)  propertyr(   r.  r2  r7  r9  r(  r<  r  r  r  r  r  `  s    TM M" ""..99D@@ It >Bu||3d:A. .C .D . !4}DtK ! !-
'9F 9t 9r  r  c                   (     e Zd Zd fdZddZ xZS )BackwardStateGraphArgc                 >    t         |   d t               dd d       y )NF)r  r  r   r!  r"  )super__init__r(   )r*  	__class__s    r  rH  zBackwardStateGraphArg.__init__  s&    "_$ 	 	
r  c                    j                   j                  j                  sJ j                  fd       j	                  dd       j                          j                  j                   j                  j                         y )Nc                  D     j                  t        j                  d      S )Nr(   )load_import_fromr(   r  )r6  s   r  <lambda>z3BackwardStateGraphArg.reconstruct.<locals>.<lambda>  s    G,,]-E-EW r  r   F)txoutputbackward_state_varadd_push_nullcall_functiondup_topstorer5  s    `r  r7  z!BackwardStateGraphArg.reconstruct  se    zz  3333W	
 	a'gjj''::;r  r=  r>  )r  r  r  rH  r7  __classcell__rI  s   @r  rE  rE    s    
<r  rE  c              #      K   | ]8  \  }}|j                  d       s"t        j                  |      rt        |       : yw)_N)
startswithinspectisclassid).0r$  members      r  	<genexpr>r_    s7      /f??3GOOF$; vJ/s   >A ITERTOOLS_TYPE_IDSITERTOOLS_POLYFILLED_TYPE_IDSc            
           e Zd ZdZ	 d5dddededdf fdZd	edefd
Z	d	edefdZ
dedefdZdefdZdedef   deeef   dz  fdZedeeedef   f   fd       Zeej*                  dedeeedef   f   fd              Zd	ej0                  e   defdZd	ej8                  e   defdZd	ej@                  jB                  jD                  de#fdZ$d	ede%fdZ&d	edefdZ'eej*                  dee(ed egef   f   fd              Z)d	edefdZ*d	edefdZ+d	e,e-edf   e.e   e/e0f   defdZ1d	e2defdZ3d	e4defdZ5d	edefd Z6d	e7e8z  de9e:z  fd!Z;d	ejx                  d"eddfd#Z=d	ej|                  j~                  defd$Z@d	edefd%ZAd	e(defd&ZBd	eCdefd'ZDeEd	e(d(ee(gef   defd)       ZFeEd	eCd(eeCgef   defd*       ZFeE	 d6d	e(eCz  ez  ez  d(ddefd+       ZF	 d6d	e(eCz  ez  ez  d(eegef   dz  defd,ZFd	ejx                  ddfd-ZGd	ejx                  defd.ZHd	edefd/ZI	 	 d7d	e(d0eJdz  d1eKdz  defd2ZLd	eCdefd3ZMd	edefd4ZN xZOS )8VariableBuilderz3Wrap a python value in a VariableTracker() instancerN  r  r  allow_lazy_constantr	  Nc                     |J d       t        j                         J d       t        |           || _        || _        |j                  | _        || _        y )NzRConsider SourcelessBuilder for ephemeral objects, usually objects created locally.zExpected active TracingContext)r   try_getrG  rH  rN  r  r$  rd  )r*  rN  r  rd  rI  s       r  rH  zVariableBuilder.__init__  sf     ! 	
`	
! %%'3U5UU3KK	 $7 r  r%  c                 x   t        j                         }	 | j                  |      | j                  j                  j
                  xj                  t        j                         |z
  z  c_        S # | j                  j                  j
                  xj                  t        j                         |z
  z  c_        w xY wr'  )timetime_ns
_call_implrN  rO  bytecode_tracing_timingsvariable_builder_call_ns)r*  r%  _t0s      r  __call__zVariableBuilder.__call__  sy    lln	??5)GGNN33LL$LDGGNN33LL$Ls   A/ /A
B9c                    | j                   j                  j                  j                  j	                  | j
                         || j                   j                  j                  v r| j                   j                  j                  |   }t        | j
                  |j
                        }|r| j                  |       t        |t        j                  j                        r+t        |t              r|j                  | j
                         |S | j                   j                  j                  j!                  | j
                        }|r| j"                  st        |t$              s|S | j'                  |      }|j
                  | j
                  |_        dt(        dt*        dt,        fd}| j/                  |      s	 |||      r]|| j                   j                  j                  vr;t1        |      s0| j                   j                  j                  j3                  ||      }d| j
                  j4                  vr-|| j                   j                  j                  | j
                  <   |S )Nr%  vtr	  c                 `    t        |       xs t        | t              xr t        |t              S r'  )r   r,  r*   r   )r%  rp  s     r  _is_deduplicable_sym_variablezAVariableBuilder._call_impl.<locals>._is_deduplicable_sym_variable  s,     U#Hz%'H2R12r  JVP_NESTING)rN  rO  current_tracertraced_sourcesaddr  side_effectsrI   install_guardsr,  r  nnModuler   set_nn_module_stack_sourcevariable_tracker_cachegetrd  r   _wrapr   r   r@  _can_lift_attrs_to_inputsr   track_object_existingr$  )r*  r%  side_effect_result	dup_guard	cached_vtrp  rr  s          r  rj  zVariableBuilder._call_impl  s   %%4488EDGGNN///!%!<!<U!C'5G5N5NOI##I.%1j"$A7 #==dkkJ%%GGNN99==dkkJ	 ''z.0 ! ZZ99BI	2 	2/ 	2d 	2 ..r20;TWW^^8883E:,,BB5"MB  0 00ACDGGNN11$++>	r  rp  c                 N    t        |      t        t        t        t        t
        hv S r'  )typer   r   r   r   r   )r*  rp  s     r  r  z)VariableBuilder._can_lift_attrs_to_inputs/  s'    Bx(% %
 
 	
r  c                     | j                   S r'  r4  r1  s    r  
get_sourcezVariableBuilder.get_source8  s    {{r  guards.c                     | j                         }	 |D cg c]  }|j                  |       }}t        |ddi i S c c}w # t        $ r Y y w xY w)Nskipr   )r  
make_guardNotImplementedErrorrH   )r*  r  r  guardtmps        r  rx  zVariableBuilder.install_guards;  s]    "	9?@6$$U+@C@ 	s##		 A" 		s   A =A A 	AAc                 @    | j                  t        j                        S r'  )_type_dispatch_implr>   trace_numpy)clss    r  _type_dispatchzVariableBuilder._type_dispatchD  s    &&v'9'9::r  r  c                 N   t         j                  t         j                  j                  t         j                  j
                  t         j                  j                  j                  f| j                  ft        t        t        t        j                  t         j                  f| j                  ft         j"                  | j$                  ft&        | j(                  ft*        | j,                  ft.        t0        f| j2                  ft        t4              | j6                  ft8        j:                  | j<                  ft>        j@                  | jB                  ft         jD                  jF                  jH                  | jJ                  ft         jL                  jN                  | jP                  ftR        jT                  | jV                  fg}|r1tX        r+|j[                  tX        j\                  | j^                  f       i }|D ]*  \  }}ta        |t              r|n|fD ]  }||vsJ |||<    , |S r'  )1r  r0  ry  	ParameterrA  r!   functional_tensorFunctionalTensorwrap_tensortuplelistr   collectionsdequeSizewrap_listlike	itertoolscountwrap_itertools_countr   wrap_tuple_iteratorr   wrap_range_iteratorslicerangewrap_slice_rangers   wrap_literalrePatternwrap_regex_patternweakrefReferenceTypewrap_weakrefutilshooksRemovableHandlewrap_removable_handlejitScriptFunctionwrap_jit_functiontypesMappingProxyTypewrap_mapping_proxyr   appendndarraywrap_numpy_ndarrayr,  )r  r  entriesresulttsfnr  s          r  r  z#VariableBuilder._type_dispatch_implH  s    LLHH&&%%00%%77HH	  lK,=,=uzzJ!! __c667S445S445U^S112()3+;+;<ZZ//0""C$4$45[[..0I0IJYY%%s'<'<=##S%;%;</
4 2NNBJJ(>(>?@  	FB%b%0Rre &q		
 r  c                 V    | j                  t        j                         t        |      S r'  )rx  rG   ID_MATCHr   r*  r%  s     r  r  z"VariableBuilder.wrap_regex_patterns  s!    L112#E**r  c                     | j                  t        j                         t        j                  | j
                  || j                        S Nr4  )rx  rG   
TYPE_MATCHr   buildrN  r  r  s     r  r  zVariableBuilder.wrap_weakrefx  s3    L334$$TWWeDKKHHr  c                 "    t        dddg        y )Nz3Attempted to represent unregistered RemovableHandle zDynamo attempted to build a representation of a torch.utils.hooks.RemovableHandle, which is not supported. This happens because the RemovableHandle was created in another frame.gb_typecontextexplanationhints)rF   r  s     r  r  z%VariableBuilder.wrap_removable_handle|  s     	Im	
r  c                    t        |d      s5t        ddt        |      j                   dg t        j
                         | j                  t        j                         t        |d| j                  t                     S )N_torchdynamo_inlinez.wrap_jit_function: missing _torchdynamo_inlineztype: zfDynamo expected a JIT function with a _torchdynamo_inline attribute, but the object does not have one.r  r  mutation_type)r  rF   r  r  r?   SUPPORTABLErx  rG   r  r   r  r   r  s     r  r  z!VariableBuilder.wrap_jit_function  sw    u34H e!5!5 6746)556 	L334*!;;35	
 	
r  c           
      t     j                  t        j                          j                  t        j                         t	        d |D              }|sEt        dd|j                         D cg c]  }t        j                  |      r| c} ddg       dt        dt        d	t        t        t        f   f fd
t        fd|j                         D              }t        |d       }t        | j                         } j"                  j$                  j&                  j)                  ||      S c c}w )Nc              3   F   K   | ]  }t        j                  |        y wr'  r   
is_literalr]  ks     r  r_  z5VariableBuilder.wrap_mapping_proxy.<locals>.<genexpr>  s     F1(33A6F   !znon-const keys in mappingproxynon-const keys: z1Dynamo expects mappingproxy keys to be constants.zFEnsure your mappingproxy keys are constants (e.g. int, float, strings)r  r  vr	  c                     t        j                  |       }| }t        j                         |      }t	        j                  ||      }||fS r'  )r   creater[   r  r   )r  r  key
source_keysource_value	res_valuer*  s         r  build_key_valuez;VariableBuilder.wrap_mapping_proxy.<locals>.build_key_value  sG    "))!,CJ():JGL+221lCI	>!r  c              3   6   K   | ]  \  }} ||        y wr'  r  )r]  r  r  r  s      r  r_  z5VariableBuilder.wrap_mapping_proxy.<locals>.<genexpr>  s     Etq!_Q*Es   r4  )rx  rG   r  MAPPING_KEYS_CHECKallrF   keysr   r  r   r  r   dictitemsr   r   r  rN  rO  rw  track_mutable)r*  r%  	all_constr  r  dict_vtr  r  s   `      @r  r  z"VariableBuilder.wrap_mapping_proxy  s   L334 	L;;<FFF	8*uzz|+j!K[KfKfghKiA+j*klO\		"s 	"s 	"u_o5U/V 	" Eu{{}EE $E$7%gdkkBww~~**88GG+ ,ks   .D5	D5c                     ddl m } |d ft        j                  d ft        j                  d fg}i }|D ]9  \  }}t        |t        t        f      r|n|fD ]  }||vsJ ||t        |      <    ; |S )Nr=   )comptimec                     t               S r'  )r   r  s     r  rM  z.VariableBuilder._id_dispatch.<locals>.<lambda>  s
    +;+= r  c                 t    t        t        fd| j                  i| j                  t        j
                        S )Nr  )r   _dataclasses_fields_lambdar  rx  rG   CLOSURE_MATCHr  s     r  rM  z.VariableBuilder._id_dispatch.<locals>.<lambda>  s5    N.%;;% )),*D*DE% r  c                     t               S r'  )r   r  s     r  rM  z.VariableBuilder._id_dispatch.<locals>.<lambda>  s
    4H4J r  )	r  dataclassesfieldsr  __version__r,  r  r  r\  )r  r  r  r  r  r  r  s          r  _id_dispatchzVariableBuilder._id_dispatch  s    
 	( =>""  JK
  	#FB%b5$-8Rre #& "r!u#	#
 r  c                 Hm   BC ddl m}m}m} ddlm}m}m}  |       rddlm	} ddl
m}	 n G d d      }	 G d	 d
      }dld}
dld} G d d      }dld} |       rddlm}
m}  |       rddlm}  |       rdd l}t%        |d      r|j&                  } j)                         j+                  t-                    }|	 |       S  j/                         j+                  t1                    }|	 |       S t3        t4        j6                        rt-              t4        j8                  j:                  j<                  t4        j8                  j:                  j>                  t@        fvrft-              tB        jD                  vrKt-              jF                  t4        j6                  jF                  u stI              r jK                        S tM              r jO                  tP        jR                         tU        t-                    D cg c]7  }tW        jX                  t[        |      t]         j^                  |            9 }}ta        | j^                  tc                     } te        jf                  t-                     j^                  |      } jh                  jj                  jl                  jo                  |      S tq        tr        tt        jv                  tt        jx                  f      r} jO                  tP        jz                         t}        d D              BBs9 jh                  jj                  j~                  j                   j^                         dt        dt        dt        dt        t        t        f   fB fdCts        Cfdt              D              }tq        tt        jv                        r]t]         j^                  d      }t        |t-               t         jh                  |      j                         j^                        }n!t        |t-               j^                        } jh                  jj                  jl                  j                  |      S t3        t4        j8                  j                        r j                        S t        j                        r j                        S t3        t4        j                  j                        rIt         j^                        } jh                  jj                  jl                  jo                  |       |S tq        t        t        f      rt        d  D              r t        d!d"d#d$gt        j                  %        jO                  tP        jz                          jO                  tP        jR                         t        }tq        t              rst]         j^                  d&      }t        |j                  tP        j                                jh                  jj                  j~                  j                  |       t        }t              }t        |      D cg c]0  \  }}tW        jX                  |t         j^                  |            2 }}} || j^                        } jh                  jj                  jl                  jo                  |      S tq        t              rut}        d' D              rcD cg c]"  }t        jY                   jh                  |      $ }} jO                  tP        j                         t        | j^                        S t3        t        j                  t4        j                  t4        j                  j                  j                  f      r6 jO                  tP        j                         t         j^                        S t        j                        r6 jO                  tP        j                         t         j^                        S t              rlt        fd(t4        j                  jB                  j                  D              r6 jO                  tP        j                         t         j^                        S t3        t        j                        r6 jO                  tP        jz                         t         j^                        S t              rt         j^                        S t              rt         j^                        S t        j                        r t        j                   j^                        S t3        t        j                        rt]         j                         d)      } t         jh                  |      j                         }g }t]         j                         d*      } t        j                        D ]<  \  }}!|j                   t         jh                  t        | |            |!             > i }"t]         j                         d+      }#j                  j                         D ]g  \  }$}t        j                  |$      s#t        d,d-|$ d.g t        j                  %        t         jh                  t        |#|$            |      |"|$<   i t         j                         j                  tP        jz                        |#j                  tP        j                        | j                  tP        jR                               t[        d/d       }%t        |||"|%0      S t              r7 jO                  tP        j                         t         j^                        S t        Ct3        t        j                        r' j                  t        j                              S t        j                         rWt        sJ tq        t"        j$                        r`t'        j(                  j*                        rt        t]         j^                  d1      j                  tP        j,                               nt'        j.                        r! jO                  tP        j0                         nt'        j(                        r! jO                  tP        j,                         nJt              r  jO                  tP        j                         n jO                  tP        jz                         t3         j^                        S t        j4                        r7 jO                  tP        j                         t7         j^                        S t        j8                        rt3        t        j:                        r^ jO                  tP        jz                         t]         j^                  d2      }&t        |&j                  tP        j                               n jO                  tP        j                         t=         j^                        S t?        j@                        rM jO                  tP        j,                         t?        jX                   jh                   j^                        S tq        t4        jB                  jD                  jF                        r8 jO                  tP        j0                         tI         j^                        S t3        t4        jB                  jD                  jJ                        r4d }'	 jL                  }'g }( j^                  j                  tP        jz                        g})t3        |'t              rt]         j^                  d3      }*|)j                  |*j                  tP        jR                               t        |'      D ]<  \  }}|(j                   t         jh                  t        |*|            |             > t        |)   jh                  jj                  jl                  jo                  tQ         j^                  tS        |(      4            S t3        t"        j$                        rtq        t[        d5d       t4        jB                  jD                  jF                        rt[        d6d7      d8k(  rt[        jT                  d8d       k(  rt        t]         j                         d1      j                  tP        j,                               tW        tI        jT                  t]         j^                  d59            d8t-              :      S t3        t4        j                  jX                        r7 jO                  tP        j                         t[         j^                        S t4        j                  j\                  j^                  j`                  u r/ jO                  tP        j,                         tc         fd;      S t3        |      rte        jf                        S t3        |      rti        jj                        S t3        |      r#tm        jn                  jp                        S t              r|t        jr                        et        jt                        rd< jh                  jj                  _        t        jr                        jy                   j^                        S t        r.t3        t        jz                        r j}                        S t3        t~              rot4        j                  j                  u rt        d=d7d>g %        jO                  tP        jz                         t        j                   j^                        S t3        t4        j                  j                        r} jO                  tP        j                         t]         j^                  d?      }+ t         jh                  |+      j                        },t        jX                   jh                  |,      S t3        t4        j                        r jO                  tP        jz                         t3         j^                  t              rt        }-nt         j^                        }- jh                  jj                  j                  d@t        |-fi       }.t        |.j                         t        |. j^                  |-A      } jh                  jj                  jl                  jo                  |      S t3        t4        j                  j                        rK jO                  tP        jz                         t        jX                   jh                   j^                        S t3        t4        j                  j                  j                        r+ jO                  tP        j                         t              S t3        t4        j                        r jO                  tP        jz                         t         j^                        }- jh                  jj                  j                  d@t        |-fi       }/t        |/j                         t        |/|- j^                        S tq        t        j                        rOt'        j                  dBd       6 jO                  tP        jz                         t         j^                        S t        j                        rg jO                  tP        jz                         t         j^                        } jh                  jj                  jl                  jo                  |      S t3        t4        j                  j                        r\ jO                  tP        j                         t         j^                         _/        t         j^                  t                     S t3        t4        j                        r, jO                  tP        j                         t              S t        j                        rt         j^                        S t        u r* jO                  tP        j                         t               S t1              t        v rJt1              t        vr8 jO                  tP        j0                         t         j^                        S t3        t              rwt3        t              }0t         j^                  |0      }1t1               jh                  jj                  j                  j                  v r9 jh                  jj                  j                  j                  t1                 }2n jh                  jj                  j                  j                  j                  |1t        j                  C      }3 jh                  jj                  j                  j                  |3j                  |1D      }2t3        |2t              st        dEt-        |2              jh                  jj                  j                  j                  t        j                  dFdG j                        t-        |2      |2|1      }4t        |1|2dHd dH|2I      |4j                  j                  dJ<   |2j                  j                  }5t3        |5t         j                        s
J |5 dK        jh                  jj                  j                  j                  t        |2|1d              t	        jX                   jh                  |4|2      S t              rt3        t4        j                        r j^                  nt         j^                        }1d }6j                  j                         re jh                  jj                  j                  j                  t        j                  j                        |1t        j                  L      }6nkt3        t4        j                        r1 jh                  jj                  j                  j                         }6nt        dMd7dNg t        j                  %       |6J t3        |6t              st        dEt-        |6              jh                  jj                  j                  j                  t        j                  dFdG j                        t-        |6      |6|1      }4t        |1|6dHd dH|6I      |4j                  j                  dJ<   |6j                  j                  }5t3        |5t         j                        s
J |5 dK        jh                  jj                  j                  j                  t        |6|1d              t3        t4        j                        r|6n|6dOk(  }7t	        |4|7      S t3        |	|f      r9 jO                  tP        j                         t        d d  j^                        S |
u rt        dOP      S |u rt        dP      S |j                  u rt!               S |u rt#              S t3        t4        j$                  j&                  j(                        rt3        t4        j$                  j&                  j*                        r j-                        S  jO                  tP        j                         t/        j0                  j2                  j4                  j6                  g j^                  Q      S t9        j:                        rt'        j.                        r! jO                  tP        j0                         n7t'        j(                        r  jO                  tP        j,                         t9         j^                        S t'        j                  dRdH      rB jO                  tP        jz                         t=        dS j^                  t                     S t?              rB jO                  tP        jz                         t=        dT j^                  t                     S t@        jB                  u s)t@        jD                  dUk\  r,t@        jF                  u rtI         j^                        S tK              r[t'        j                  dVdH      rB jO                  tP        jz                         t=        dV j^                  t                     S tt        jL                  u r7 jO                  tP        j                         tO         j^                        S t3        t"        jP                        rNtS        jT                        r7 jO                  tP        j                         tS         j^                        S tW              rtX        jZ                  tX        j\                  fv r] jO                  tP        j                         tW        t_        tX         j^                        j`                  t-              :      S tK              r_tc              \  }8|8t]         j^                  |8       _/        t        jd                        jy                   j^                        S tf        jh                  u r6 jO                  tP        j                         tk         j^                        S tq        tf        jh                        rtm        jn                        ri jO                  tP        jz                         tm         j^                        } jh                  jj                  jl                  j                  |       |S t3        t"        jp                  tr        jt                  f      rj jO                  tP        jv                         ty         j^                        } jh                  jj                  jl                  jo                  |       |S t3        t"        jz                        rt}              S t3        t"        j~                        rt              S t        t-              t,              r=t        t              r+ jO                  tP        j                         t              S t        t-              t,              rt4        j                  j                  j                  t4        j8                  j                  t4        j8                  j                  fv r3t        jd                        jy                   j^                        S t4        jB                  j                  u rA jO                  tP        j0                         t        j                   jh                        S t        t4        j6                        r\t4        j6                  urJjF                  t4        j6                  jF                  u r$t              st         j^                        S t         j^                        s  jO                  tP        j0                         t3        t,              r3t        t        j                        rt         j^                        S t              rt         j^                        S t         j^                        S t        j:                  t-                    rZddWl/mS}9mT}: t3        t4        j                  j                  j                        rj                  t4        j                  j                  j                        r jh                  jj                  j                  j                  t        j                  dFdG j                        t-               j^                        };t         j^                  dHd dH      |;j                  j                  dJ<   t        jX                  |; j^                        S t        t-                    r" jO                  tP        j                         n t        t-                    r@ jO                  tP        jz                          jO                  tP        j                         nt%        dX      s j-                        S tW        j                   t         jh                   |: j^                              j                         j                                      tW        j                   t         jh                   |9 j^                              j                                      t4        j                  j                  j                   jh                  jj                  j                        }<t        t-                    rt        t-                    s}<};n/tB        j                  rt         j^                        s,t         j^                        st         j^                        r> jh                  jj                  j                   j                   j^                        S  jh                  jj                  j                  j                  t        j                  dFdG j                        t-              |< j^                        };t         j^                  dHd dH|<      |;j                  j                  dJ<   t        jX                  |;|< j^                        S t3        tr        tt        jx                  f      rt-              j                  tr        j                  u r_ jO                  tP        jz                          jO                  tP        jR                          jh                  jj                  j~                  j                   j^                         dt        dt        dt        dt        t        t        f   f fdYCts        CfdZt              D              }t        |t3        tt        jx                        rtt        jx                  ntr        tc                j^                  [      }=d<|=_k        t        |= j^                  \      } jh                  jj                  jl                  jo                  |      S t3        t              r jO                  tP        jz                          jO                  tP        jR                         t        t        j                              D cg c]G  }tW        jX                  t        j                  |      t         j                         |            I }}ta        | j^                  tc                     }te        | j^                  ]      } jh                  jj                  jl                  jo                  |      S t3        t              r jO                  tP        jz                          jO                  tP        jR                         t        t        j                              D cg c]G  }tW        jX                  t        j                  |      t         j                         |            I }}t        | j^                  tc                     }>t        |> j^                  ^      } jh                  jj                  jl                  jo                  |      S t3        t        t        f      ra jO                  tP        jz                          jO                  tP        jR                         t        tr        j                              }t        t        j                  |            D cg c]F  }tW        jX                  t        j                  ||      t         j                         |            H }}t3        t              rt        }?nt3        t              sJ t        }? |?| j^                  tc                     }@t        |@ j^                  _      } jh                  jj                  jl                  jo                  |      S t        t-              t              rg jO                  tP        jz                         t         j^                        } jh                  jj                  jl                  jo                  |      S t              rg jO                  tP        jz                         t         j^                        } jh                  jj                  jl                  jo                  |      S t3        t              rt}        d` D              rD cg c]"  }t        jY                   jh                  |      $ }}t         j                         j                  tP        jR                         j                         j                  tP        j                               t        | j^                        S t        dadbD $cg c]  }$t        j                  |$      r|$ c}$ dcddg%       y t        j                        rddel|m}}A j                  %j                  j,                  Aj                  k(  r j                  j                        S j                  j,                  Aj                  k(  rt        j                  df j^                  j                  t        j                          j                  j                  t        j                  t	        t        dHg      h      i      S j                  j,                  Aj                  k(  rmt        j                  df j^                  j                  t        j                          j                  j                  t        j                  j      S tO        dkj                         tq        t              r7 jO                  tP        jz                         t         j^                        S  j-                        S c c}w c c}}w c c}w # tN        $ r Y &w xY wc c}w c c}w c c}w c c}w c c}$w )mNr   )
has_triton has_triton_experimental_host_tma%has_triton_tensor_descriptor_host_tmar=   )CudagraphOverrideContextManagerDynamoConfigPatchProxy(ErrorOnGraphBreakDecoratorContextManager)	Autotuner)JITFunctionc                       e Zd Zy)*VariableBuilder._wrap.<locals>.JITFunctionNr  r  r  r  r  r         r  r  c                       e Zd Zy)(VariableBuilder._wrap.<locals>.AutotunerNr  r  r  r  r  r    r  r  r  r	  c                       y r'  r  r  r  r  create_1d_tma_descriptorz7VariableBuilder._wrap.<locals>.create_1d_tma_descriptor      r  c                       y r'  r  r  r  r  create_2d_tma_descriptorz7VariableBuilder._wrap.<locals>.create_2d_tma_descriptor  r  r  c                       e Zd Zedd       Zy)/VariableBuilder._wrap.<locals>.TensorDescriptorNc                       y r'  r  r  r  r  from_tensorz;VariableBuilder._wrap.<locals>.TensorDescriptor.from_tensor  s    r  r=  )r  r  r  staticmethodr  r  r  r  TensorDescriptorr
    s     r  r  c                       y r'  r  r  r  r  set_allocatorz,VariableBuilder._wrap.<locals>.set_allocator  r  r  )r  r  )r  r  r4  r  )r  tuple_vtc              3   F   K   | ]  }t        j                  |        y wr'  r  r  s     r  r_  z(VariableBuilder._wrap.<locals>.<genexpr>E  s     Jq,77:Jr  ir  r  c                     	j                         }rt        j                  |      }|}n"t        ||       }t	        j                  ||      }t        ||      }t	        j                  ||      }||fS r'  )r  r   r  rT   r   rW   )
r  r  r  baser  r  r  r  r  r*  s
           r  r  z.VariableBuilder._wrap.<locals>.build_key_value`  so     (*11!4C!"J!3D!!<J-44Q
CC0zB/66q,G	I~%r  c              3   :   K   | ]  \  }}} |||        y wr'  r  r]  r  r  r  r  s       r  r_  z(VariableBuilder._wrap.<locals>.<genexpr>o  (      Aq!  1a(   default_factory)r  r  )user_clsr  c              3   P   K   | ]  }t        |t        j                           y wr'  )r,  r  r0  r]  xs     r  r_  z(VariableBuilder._wrap.<locals>.<genexpr>  s     >1:a.>s   $&z$Attempted to wrap a set with tensorsz+Python set containing torch.Tensor elementszDynamo cannot trace sets of tensors. To get a stable ordering, Dynamo needs to convert the set into a list and the order might not be stable if the set contains tensors.z,Use a dictionary where the keys are tensors.r  _dictc              3      K   | ]p  }t        |      t        j                  u xr |j                  d k(  xs> |t        j
                  j                  j                  v xs t        |t        t        f       r yw)r  N)r  r  BuiltinMethodTyper  r  r  _pytreeBUILTIN_TYPESr,  r   r    r  s     r  r_  z(VariableBuilder._wrap.<locals>.<genexpr>  sr      .
  aE333O8O > U[[((666> 1z+;<=>.
s   A6A8c              3   &   K   | ]  }|u  
 y wr'  r  )r]  r  r%  s     r  r_  z(VariableBuilder._wrap.<locals>.<genexpr>  s      %
ERK%
s   funcargskeywordsz,functools.partial() with non-literal keywordznon-literal keyword: z3functools.partial() expects literal/string keywords
cache_hash)original_cache_hash__func__dtypesaved_tensors)r  r,  __self__r  r  apply)r^  py_typec                     t        t        j                  j                  j                  j
                        j                   j                   j                  j                  j                  j                         fi       S r'  )r   r  _dynamoexternal_utilsFakeCompiledAutogradEngineexec_final_callbacksrR  rN  rO  rw  get_ca_final_callbacks_varr1  s   r  rM  z'VariableBuilder._wrap.<locals>.<lambda>}  sZ    ,MM00KK``-GGWW^^00KKMO r  Tz9Attempted to wrap torch._higher_order_ops.invoke_subgraphzJDirectly using invoke_subgraph is not supported. Use nested_compile_regionstreamrR  )r  user_object_indexenter_result)r  dynamic_dim)hintr  Expected SymInt, got [^a-zA-Z0-9]+rX  F)r"  r#  grapharg is not a basic Symbol.)r:  z!Attempted to wrap unbacked SymIntz+Unbacked SymInt input is not supported yet.r   )rank)target_valuesr  __script_if_tracing_wrapper__original_fn__wrapped__)      r  )FlattenScriptObjectSourceScriptObjectQualifiedNameSource__obj_flatten__c                     j                         }t        ||       }t        j                  ||      }t	        ||      }t        j                  ||      }||fS r'  )r  rT   r   r  rX   )	r  r  r  r  r  r  r  r  r*  s	           r  r  z.VariableBuilder._wrap.<locals>.build_key_value  sW     (/a8
)00J?8zJ/66q,G	I~%r  c              3   :   K   | ]  \  }}} |||        y wr'  r  r  s       r  r_  z(VariableBuilder._wrap.<locals>.<genexpr>  r  r  )r  r  r  )r  r  )r  r  )list_vtr  )set_vtr  c              3   F   K   | ]  }t        j                  |        y wr'  r  r  s     r  r_  z(VariableBuilder._wrap.<locals>.<genexpr>  s     Aa#..q1Ar  znon-const keys in dict_keysr  z.Dynamo expects dict_keys keys to be constants.zCEnsure your dict_keys keys are constants (e.g. int, float, strings))_DimHintTypez%s marked %s via IntWrapper	warn_only)
constraint)dynamismr  rS  zUndefined dynamism r=  (  torch.utils._tritonr  r  r  
decoratorsr  r  r  triton.runtime.autotunerr  triton.runtime.jitr  $triton.tools.experimental_descriptorr  r  triton.tools.tensor_descriptorr  tritonr  r  r  r}  r  r  r\  r,  r  r0  ry  	parameterUninitializedBufferUninitializedParameterr7   r>   nontraceable_tensor_subclasses__torch_dispatch__r9   r  r~   rx  rG   SEQUENCE_LENGTHr   r   r  getattrrQ   r  r   r   r   
get_vt_clsrN  rO  rw  r  r   r  r  defaultdictOrderedDictr  r  guard_on_key_orderrv  r   r  r   ru   r   rc  r  r   r  rz  wrap_moduler   r  r  	overridesTorchFunctionModer   setr8   anyrF   r?   r  r   rH   r  DICT_KEYS_MATCHr   r  	enumeratere   	frozensetSourcelessBuilderEQUALS_MATCHr   enumEnumDispatchKey_C
_functorchTransformTyper  r   r   is_reorderable_logging_functionBUILTIN_MATCHcallabler2  ignore_logging_functionsr   loggingLoggerr   r   rp   r|   rq   r   should_wrap_in_hopr  	functoolspartialr  r%  r&  r  r[   r'  r  
USER_ERRORrW   r   r   r   r   genericr  asarrayrB   is_numpyr  
MethodTyperZ  
isfunctionr*  r  r[  CLASS_MATCHr   is_numpy_dtyper   is_numpy_type_infoiinfor   r   can_rewriteautogradfunctionFunctionMetar   FunctionCtxr,  RuntimeErrorr   r   r-  r   _ImperativeEnginer   r3  r4  _exec_final_callbacks_stubr   r   changesr   error_on_graph_breakr   fwdbwdlookup_callableis_callable_allowed!has_user_defined_allowed_in_graphcreate_with_sourcenumberwrap_unspecialized_primitiver   _higher_order_opsinvoke_subgraphr   makecudaStreamContextr7  r   StreamrV   r   r   create_proxyr   r   noder   _SDPAParamsr   pyfunctorchFuncTorchInterpreterr   Eventr   
contextlibnullcontextgetattr_staticr   r   is_matching_objectoptim	Optimizerrg   r   r   DispatchKeySetDISPATCH_KEY_SET_MATCHr   r   is_group_member_typer   r`  ra  r   r*   r+   rY   root_tracerdynamic_scalar_nodes	shape_envcreate_unspecified_symbolrealr.   DYNAMICcreate_symintnoder   AssertionErrorcreate_graph_inputr  subr$  r  metaexprsympySymboltracked_fakesr4   r   r   rU   has_hint$create_unspecified_symint_and_symbolintr;  SymBoolcreate_unbacked_symintr   r   r  r   r   ampautocast_modeautocast_UnmanagedAutocastwrap_user_definedr   device
fast_dtype_enabled_cache_enabledr   is_matching_clsr   r}   sysexc_infoversion_info	exceptionr   r{   
namedtupler   r!  r   is_supported_builtin_methodrz   floatfromhexhexr   r  r   lookuprandomRandomr   r   is_supported_random_objr   rA   DummyModuleMODULE_MATCHr   GetSetDescriptorTyper   MethodWrapperTyper   
issubclassBaseExceptionr   r  r  BackwardHookr  Buffer _unsafe_preserve_version_counterr   constructorr:   r   r^   r   r   r   r   rG  rH  _libraryfake_class_registryFakeScriptObjectreal_objtracing_with_realr   CONSTANT_MATCHr   OPAQUE_OBJ_GUARD_FN_MATCHrealize_all_typequalified_namerI  maybe_to_fake_obj	fake_moder   install_free_tensorsr_   r`   rb   register_attr_or_module__new__should_reconstruct_allr   r  __len____getitem__rc   r   r   fromkeysr   r   r   ry   r   rt   r   r   torch.export.dynamic_shapesrO  rS  STATICwrap_symintvallogdebugr3   r/   AUTOr(  r   )Dr*  r%  r  r  r  r  r  r  r  r  r  r  r  r  
triton_modtype_dispatchid_dispatchr$  rO  r  r  factory_sourcevarset_var_clsinternal_dict_sourceLr  r  r  func_srcfunc_objr&  args_sourceargr'  keywords_sourcer  r)  	dt_sourceactual_saved_tensorsr,  r  saved_tensors_sourcestream_source
stream_varindexstream_proxyevent_proxyis_intr  r  symsym_node_proxysym_expr
new_syminttracing_symint	attr_namerG  rH  proxyfake_script_objr  rL  
set_vt_clsrM  rO  r  r  sD   ``                                                                @@r  r~  zVariableBuilder._wrap  s>(   	
 	
	
 	
 <:6  			 	
	 ,- 12G<'z?3 * 8 8 ++-11$u+>$ u-- '')--bi8"tU++ uell+U ""66""99	 U6#H#HH U..%,,2Q2QQ07''.. < <= .d5k:-
 	 $**E4(%dkk48-F - %{{35H
 F-88eE{{!F
 77>>..DDUFSSED+"9"9;;R;RST 7 78JEJJI " 1155dkkB&&&#&&78&  A%H F
 e[445!+DKK9J!K,K$LODGG^$L--%  ;; +!%[;; 77>>..<<UFKKuxx/##E**((/$$U++u@@A+E$++FCGGNN''==eSIJEC,->>>BI>
 G*66  7 78 < <=%KeZ('1$++w'G$(33L4P4PQ 11556JK0 UA
 &aL	 Aq $**=dkk1ME  !t{{;F77>>..DDUFSSE9%# .
 .
 +
" DIIa&--dggq9IEI 9 9:$U4;;??YY))588+>+>+L+LM
  5 56,U4;;GG>>uE  : :;$U4;;??e_ %
"'--"6"6"O"O%
 "
  5 56*5EEw~~. 7 78(t{{CC ',DKK@@&1EE/BB5I5;;4;;OOy001!$//"3V<H9tww9%**EHD$T__%6?K#EJJ/ 3KODGG];-JKCP
 H():JGO,,. 
1'2215! N"7s ;$Y= 1 < <=	oGG.B
 !,,\-D-DE**<+G+GH&&|'C'CD #*%t"D+$>Q  u 5 56!{{  ^
5"** =**2::e+<==!!%(I2eU--.%%enn5!"4;;
;FF(66
 '##L$<$<=##E*##L$>$>?%##L$9$9:##L$;$;< t{{;;''. 5 56%eDKK@@++E2%*##L$;$;<&t{{G<	i22<3H3HIJ##L$9$9:'dkkBB.::5A : :;4;;{{ 
 E5>>22??@ 8 89+{{  u~~66BBC#' ',':':$ Mkk,,\-D-DEFF.6'1$++'O$(33L4P4PQ &&:; DAq!(( GG]3G%K 6"77>>..DD/;;"0"?  ue../z40%..2I2I2V2V z2.'9$?? 4??,j9DD ..
 #(NN%dkk*E U  uxx99: 5 56)%DD}}++FFaab  : :;!  56,U]];;GH,U-G-GHH>?,UYY		BBe_!<!<U!C!O..u5CG@..u5HHdkk I   Jubii044U;;23//???W l	  7 78388t{{SSuzz778 5 56&t{{H=M@-@NJ(//DDu||, 7 78$++':;,,UDKK@77>>66!=xL l//7 eDKK5C 77>>..DDUCPP 4 46 7 78%,,TWWeT[[IIu//;;PPQ 5 56/66u{{+ 7 78(<E''..55,	K k..6 {{	  5*001&&undCK 7 78&dkk::&99%@ 7 78.uT[[IF77>>..DDUFSSu{{445 5 56)$++6DK${{79 
 u334 C CD)%00#88?)%DDj  5 56*,,uI++5	!>> 8 89$U4;;??~.z2F(f=F%yDGGNN66KKK ww~~11FFr%yQggnn..HHJJ! * 2 2 I 
 ww~~//AA! B 
 "$/(+@d)MNN "WW^^77JJ'dii8T
	 K N 4<#'4N$$Z0 yy~~Hh5 *345 GGNN((//D&$0OP"))$''>4HH%  eU\\2 %dkk2 
 Jzz""$GGNN,,QQEJJOO,$.$6$6 R   eU]]3!%!9!9!P!P!RJ " C "$Q> 1 = =>	 )))j&1$'<T*=M<N%OPP!WW^^77JJ'dii8Z 	 K N 4<#-4N$$Z0 "++Hh5 *345 GGNN((//JPT0UV )=
:QR?  #>>BBY78 5 56'{{	  ..:BB..:BB&222466m#-e44uyy66??@%!8!8!K!KL--e44##L$9$9:+((,,	#  ;;  *99%@u%##L$<$<=##E*##L$>$>?/dkkJJ##E+H%P 7 78.{{79	  +51 7 78.{{79	  cll"'ES]],B&uT[[AA#E*w/E/E(%0
  7 78.%{{79	  k,,, 5 560t{{KK5**
#??F 5 56(t{{CCEemmUYY-G$G 5 56"dkk:U 
 $E*?FE9 $(i@%%e,??dkk @   fmm# 5 56&dkk::E6==)n.T.T/
  7 78#E$++>FGGNN''55eVDM  0 0-2K2KLM 9 9:){{F
 GGNN''==eVLMu99: ,E22u667
 )//UT*z%/O 5 564U;;UT*!!..""  #))%0CC$++ D   GGG##L$<$<=5AA$''JJ 5%,,/- ,,0O0OO
 ;5A-eDKKHH)$++6 ##L$<$<=%&:eTYY+G/;; 
 e$0;; 
 ,{{  '66tE{C  %!C!C!T!TU ~~11CCEJ22EEFF+S$))<K;;	 F  /7KK/


+ 177;;  $DK0##L$?$?@)$u+6##L$;$;<##L$J$JKU$56 --e44 $//O!@!M 446 $//TODGG-Ft{{-ST--/ $nn@@RR((%O $DK0d5k9R"',,%dkk2*4;;79$++Fww~~==499T[[ >  
 22EEFF+S$))<K#;;	 F  /7KK#/


+ -33{{  ut[%<%<=>U##t||3  7 78 < <= GGNN--11$++>
&
&
&#&
&78
&  A%H F
 ( "%)@)@A  ++35{{	G .2G*,UGDKKXF77>>..DDUFSSu% 7 78 < <= u}}U34
 	 $**%%eQ/():A>F  %{{35H
 .F 77>>..DDUFSSt$ 7 78 < <= t||E23
 	 $**$$UA.,T__->BF  #{{35G
 -UGDKKXF77>>..DDUFSSY/0 7 78 < <=T]]5)*A t||A/
 	 $**$$Q*:4??;LaPF  %%(
!%333.
t{{:O:QF ,E&UF77>>..DDUFSSU^4 7 78+E$++FF77>>..DDUFSS ' 7 78,U4;;GF77>>..DDUFSSy)A5AA HMM!*11$''1=MMOO%001M1MNOO%001J1JK *%DD9.5/gaHXHcHcdeHf/g.hi P]	  2259@~~%)<)<@S@S)S''		22$$(<(<<		1KK$$&&
 ''II'//1#:U#K (   $$(9(99		1KK$$&&
 ''		J<N<N'OO"%88H#IJJE6" 7 78!%<<))%00y-@2 Jv   ^.&B N 0hs_   <CY*15CY/'CY5CY: {>ACZ@%ACZE3ACZM1'CZPCZP:CZY:
CZZCZc                 .   ddl m} | j                  t        j                         t        j                  |      r!t        |t        | j                              }nkt        ||      rHt        |      t        vr7| j                  t        j                         t        || j                        }nt        || j                        }t        j                   t        |            s|S | j"                  j$                  j&                  j)                  ||      S )Nr   )_CONSTANT_BASE_TYPESr4  )user_definedr   rx  rG   r  r   r  ri   r  r,  r  rs   CONSTANT_SUBCLASS_MATCHr   r   rO   "cls_supports_mutation_side_effectsrN  rO  rw  r  )r*  r%  r   r  s       r  r  z!VariableBuilder.wrap_user_defined.  s    6L334--e4
 %U?4;;3OPFu23U#88 D DE0t{{KF.uT[[IF==d5kJMww~~**@@OOr  c                 .   t         j                  rPt        |      t        j                  u r5| j                  t        j                         t        j                  |      S | j                  t        j                         t        |t              rt        d |D              ro| j                  j                  j!                         rK| j                  t        j                         t#        |D cg c]  }t        j                  |       c}      S t%        |      D cg c]4  \  }}t'        j                  |t)        | j+                         |            6 }}}| j,                  j.                  j0                  j3                  d      }t5        | j                  t6              r| j                  j8                  t;        |      v r| j                  }t5        |t<              sJ | j,                  j.                  j>                  jA                  tC        jD                  dd| jF                        t        |      ||      }d|jH                  jJ                  d<   tM        tN        | j,                  ||d |	      }g }	t%        |jP                        D ]  \  }}
t)        ||d
      }|
| j,                  j.                  jR                  |<   tU        ||         |
jV                  jH                  jJ                  d<   tY        jZ                  t        j\                  t_        ||               }|	ja                  |jc                  |              te        |	ddi tg        ||d
d d
      }||jH                  jJ                  d<   |D ]  }|ji                            tk        jl                  |      || j                        }t        |t<        tn        jp                  f      r0| j,                  j.                  jr                  ju                  ||      S |S c c}w c c}}w )Nr%  c              3   F   K   | ]  }t        j                  |        y wr'  r  )r]  items     r  r_  z0VariableBuilder.wrap_listlike.<locals>.<genexpr>U  s     H$$//5Hr  r4  r*  r=  rX  T	steal_arg)
target_clsrN  r  example_valuesubclass_typer  Fr  r  index_is_slicetensor_dictr  r   )r   r!  r"  r>  );r>   specialize_intr  r  r  rx  rG   r  r   r  ra  r   r  r  r  guard_sourceis_unspecialized_nn_moduler   rm  r   r[   r  rN  rO  local_scoper}  r,  rd   
local_namerw   r  r  r  r  r  r$  r  r  wrap_fx_proxy_clsr   r  input_source_to_varro   r  r~  r  TENSOR_MATCHr<   r  r  rH   r  realizer   cls_for_instancer  r  rw  r  )r*  r%  r'  r  rO  maybe_gmr  tensor_list_proxylist_variabler  tensor_variablesource_ir  r>  rp  r  s                   r  r  zVariableBuilder.wrap_listlikeE  s`      T%[EJJ%> ; ;<#**77 	L889 5% H%HH((CCE ; ;< E!RD"2"9"9$"?!RSS %U+

 4	  &&$T__%6:
 
 77>>--11&9KK
kk$$(;H(EE [[FeT*** $ : : M M'dii8U	 !N ! 8<""''4-)77'#"M F&/0C0C&D 
:"?(fAeT?N228<AU!HB%%**//> ")) --]585L h11%89
: 6**#( H 7?""''
3(  

 :!2259&U%$ 1 12377>>..<<UFKKy "S
s   .P9Pc                    | j                  t        j                         t        t	        |            D cg c]@  } t        | j                  t        | j                         |            t        ||            B }}t        || j                        }| j                  j                  j                  j                  ||      S c c}w r  )rx  rG   TUPLE_ITERATOR_LENr  r   rc  rN  rl   r  r   r   r  rO  rw  r  )r*  r%  r  rO  r  s        r  r  z#VariableBuilder.wrap_tuple_iterator  s    L;;<
 -e45	
  WODGG%?@QST%UV&ua0
 
 'vdkkBww~~**88GG
s   ACc                 B   | j                  t        j                         t        j                  |      D cg c]  }t        j                  |       }}t        || j                        }| j                  j                  j                  j                  ||      S c c}w r  )rx  rG   RANGE_ITERATOR_MATCHcopydeepcopyr   r  r   r  rN  rO  rw  r  )r*  r%  r  r  r  s        r  r  z#VariableBuilder.wrap_range_iterator  s    L==> 15e0D(
+,##A&(
 (
 &eDKK@ww~~**88GG	(
s   Bc                    t        |      \  }}t        j                  |      rt        j                  |      s| j                  |      S | j	                  t
        j                         t        t        j                  |      t        j                  |      | j                        }| j                  j                  j                  j                  ||      S r  )r   r   r  r  rx  rG   COUNT_ITERATOR_MATCHr   r  r  rN  rO  rw  r  )r*  r%  current_itemstepr  s        r  r  z$VariableBuilder.wrap_itertools_count  s    1%8d''5 ++D1))%00L==>&##L1##D);;

 ww~~**88GGr  c                 r   dD cg c]@  } t        | j                  t        | j                         |            t	        ||            B }}| j                  t        j                         t        |t              r"t        || j                  | j                        S t        || j                        S c c}w )NstartstoprG  r4  )rc  rN  rQ   r  rb  rx  rG   r  r,  r  r   r  r   )r*  r%  r  r  s       r  r  z VariableBuilder.wrap_slice_range  s    
 /	
  GODGGZ0A1%EFq!
 
 	L334eU# DD t{{;;
s   AB4r  c                 x   ddl m} t        j                  d| j                  j
                  t        |              |||       || j                  j                  j                  v rT| j                  j                  j                  |   }|j                  |j                  j                  j                  d   d<   y y )Nr=   )mark_static_addressz Marking static input %s, id: %s)r  r.  _dynamo_static_input_type)rV  rM  static_inputs_logr  r  r$  r\  rN  rO  rw  rO  r  r  r  )r*  r%  r  rM  r  s        r  mark_static_inputz!VariableBuilder.mark_static_input  s    4.0@0@"U)	
 	E/ DGGNN///''..--e4C // IINN./JK 0r  c                    ddl m} t        |j                        dk(  r7t	        dt        |      dt        |       dg t        j                  d       t        ||      rt        j                  |j                  d	d
      r<t        j                  |j                  dd       }t        | j                  d| d      S | j                  t        j                          t#        | j                  d      | _        | j%                  |j&                        S t)        |      t*        j,                  j.                  j0                  u r(t	        dt3        |      dg t        j4                         t7        |t*        j8                  j:                  t*        j8                  j<                  t*        j8                  j>                  f      r9t@        jB                  s)t	        dt3        |      ddgt        jD                         tG        |dd
      rtG        |dd
      st	        dddg        | j                  t        j                          tI        || jK                               }tM        jN                  t)        |            s|S | jP                  jR                  jT                  jW                  ||      S tY        jZ                  || jP                  j\                        r1t7        |t*        j^                  j`                  jb                  jd                        r*|jg                         }ti        | j                        | _        tk               }tm        |jn                        rI|jn                  jp                  tr        u r-	 |jo                         D ]  \  }}| ju                  ||        	 tm        |j~                        rI|j~                  jp                  t        u r-	 |j                         D ]  \  }}	| ju                  |	|        	 |r-|| jP                  jR                  j                  | j                  <   |j                  j                  d      r|j                  j                  d      rtG        |j                  dd
      r| j                  }
| jP                  jR                  j\                  rt@        j                  rt        | j                        }
t        ||
      }t        |
j                  t        j                                n| j                  }
| jP                  jR                  j\                  rt@        j                  rt        | j                        }
t        ||
      }t        |
j                  t        j                                | jP                  jR                  j                  || j                         tM        jN                  t)        |            s|S | jP                  jR                  jT                  jW                  ||      S t        |j                  t*        j8                  j                  j                  j                        r:| j                  t        j                          t        || jK                               S | jP                  jR                  j                  || j                  | jK                               S # tv        $ r?}ty        t)        |      | jP                  t{        |j|                               Y d }~fd }~ww xY w# tv        $ r?}ty        t)        |      | jP                  t{        |j|                               Y d }~Sd }~ww xY w) Nr=   )OptimizedModuler   zUninitialized nn.Modulez6Attempted to trace an uninitialized nn.Module of type .z?Ensure your nn.Module instance has called `super().__init__()`.r  _torchdynamo_disableF_torchdynamo_disable_msgzHOptimized `nn.Module` is wrapped with `torch.compiler.disable` (reason: ))r  msg	_orig_modz+torch.jit.script/freeze modules unsupportedzDynamo does not support tracing into torch.jit.script or torch.jit.freeze modules because they execute in the TorchScript runtime, not Python. Replace the ScriptModule submodule with the original eager nn.Module.z#Attempted to wrap RNN, GRU, or LSTMz*Dynamo does not support RNN, GRU, or LSTM.zgSet torch._dynamo.config.allow_rnn=True to enable experimental support for RNN, GRU, and LSTM in Dynamo_is_fsdp_managed_module_fsdp_use_orig_paramszFSDP with use_orig_params=Falser  z3Dynamo only supports FSDP with use_orig_params=Truer4  rN  )r&  )ztorch.nn.modulesz	torch.ao.ztorch.nn.modules.container_dynamo_marked_static)S
eval_framerS  len__dict__rF   r   r?   r  r   rZ  r  forwardr   r  rx  rG   r  rQ   rg  rY  r  r  r  _scriptRecursiveScriptModulerB  FUNDAMENTALr,  ry  RNNGRULSTMr>   	allow_rnnr  rb  r   r  rO   r#  rN  rO  rw  r  r@   is_dynamic_nn_moduleexportfxexperimentalproxy_tensor
_AttrProxyget_baserP   r   ry  named_parametersr*  !og_module_named_parameters_fn_ptrrQ  	TypeErrorrE   r  r&  named_buffersog_module_named_buffers_fn_ptr
nn_modulesr$  r  rY  rI  r  rm   r   rH   r  rn   r    add_fqn_info_for_inlined_modulesr  paralleldistributedDistributedDataParallelr  )r*  r%  rS  rX  r  freezingrX  peb
new_sources              r  rg  zVariableBuilder.wrap_module  s   0u~~!#1TU\]bUcTddef&11U	 %)%%emm5KUS
 ,,MM#=t /;;bcfbgghi 
  7 78$T[[+>DK##EOO44;%))++AAAEE
,&22
 uuxx||UXX\\588==IJ$$=E
H}&22	 53U; 5"95A= U	"  7 780t?PQFAA$u+N77>>..DDUFSS00G
 %!6!6!C!C!N!NO(-dkk:,.H //0**337XXR % 6 6 8 B1..q.AB ,,-''004RRR % 3 3 5 B1..q.AB
  8=))$))4   ++,MN((334PQ*A5I![[
ww~~,,0K0K!CDKK!PJ=eJWj33L4K4KLM![[
ww~~,,0K0K!<T[[!IJ6uZPj33L4K4KLMGGNN;;E4;;OAA$u+N77>>..DDUFSSOOUXX..::RR
  7 780t?PQQ77>>99		( :  c ! R,T!WdggDLQQR ! R,T!WdggDLQQRs0   7+[ +\ 	\4\		\	]4]]c                    t        |      t        u r>t        |t              sJ t        | j                  j
                        rKt        j                  d| j                  j
                         | j                  |t        j                        S t        | j                  j
                        rKt        j                  d| j                  j
                         | j                  |t        j                        S t        j                  s6t        || j                        rd }| j                  j                   j#                         s$| j                  j                   j%                         rd}t'        | j(                  | j                  j
                  t+        j,                  |      | j                  j                   j%                                | j/                  t1        j2                  t4        j6                  |             t9        j:                  || j                        S | j=                  || j>                        S | j=                  |      S t        |      t@        u rOt        |t@              sJ t        jB                  s| j=                  || jD                        S | j=                  |      S t        |      tF        tH        fv r)t        |tF        tH        f      sJ | j=                  |      S | j/                  t4        jJ                         t9        j:                  || j                        }t        |tL        tN        f      r0| j(                  jP                  jR                  jU                  ||      S |S )N&%s marked dynamic via source whitelistrT  '%s marked unbacked via source whitelista  torch.compile considers integer attributes of the nn.Module to be static. If you are observing recompilation, you might want to make this integer dynamic using torch._dynamo.config.allow_unspec_int_on_nn_module = True, or convert this integer into a tensor.r1  )recompile_hintr%  r  )+r  r  r,  is_dynamic_sourcer  r$  r  r  r  r.   r  is_unbacked_sourceUNBACKEDr>   r/  r   r0  "is_unspecialized_builtin_nn_moduler1  rN   rN  rL   make_scalarrx  r~  r  rG   rp  r   r  _wrap_lazy_constant_wrap_symint_for_lazyr  specialize_float_wrap_symfloat_for_lazyr@  rB  r  r  rj  rO  rw  r  )r*  r%  r  r  s       r  r  zVariableBuilder.wrap_literal  s   ;#eS))) !1!12		BDKKDTDTU''
8J8J'KK!$++"2"23		CT[[EUEUV''
8K8K'LL(( .eT[[A%)N00SSU;;33NNP5 ' .((+77>37;;3K3K3f3f3h	 ''!))(55n
 ,22t{{SS//t7Q7QRR++E22%[E!eU+++**//t7S7STT++E22%[T3K'edC[111++E22 ; ;<%,,5MF%$-ww~~22@@OOMr  c                 $    | j                  |      S r'  )r  r  s     r  r  z%VariableBuilder._wrap_symint_for_lazy  s    &&r  c                 $    | j                  |      S r'  )wrap_symfloatr  s     r  r  z'VariableBuilder._wrap_symfloat_for_lazy  s    !!%((r  wrap_fnc                      y r'  r  r*  r%  r  s      r  r  z#VariableBuilder._wrap_lazy_constant      
 r  c                      y r'  r  r  s      r  r  z#VariableBuilder._wrap_lazy_constant  r  r  c                      y r'  r  r  s      r  r  z#VariableBuilder._wrap_lazy_constant  r  r  c                     | j                   sJ| ||      S | j                  t        j                         t	        j
                  || j                        S t        j
                  || j                        S )zCWrap a primitive constant, deferring guard installation if allowed.r  r4  )rd  rx  rG   r  r   r  r  r   r  s      r  r  z#VariableBuilder._wrap_lazy_constant	  s_     ''"u~% ; ;<#**t{{KK#**5EEr  c                 v    t        |      r.t        |      | j                  j                  u rt	        dd      y y )Nz*Cannot wrap a Tensor that has already beenz"wrapped by this instance of Dynamo)r#   r$   rN  r  rD   r  s     r   assert_not_wrapped_by_this_graphz0VariableBuilder.assert_not_wrapped_by_this_graph	  s;    5>1%8DGG<M<MM*<4  N>r  c           	         | j                         }|| j                  j                  j                  vsJ t	        |      d u}|s]t        |t        j                  j                        s|r7|j                  j                         r| j                  |t                      d}t        j                  xr% t        |      xs t!        |      xs t#        |      }|xr0 t               xs$ t        j$                  j                  j&                  }|s6|j                  j)                         s|r]|j                  j+                         sC| j-                  |       | j                  j                  j/                  || j0                  |      S t	        |      dk(  rb| j3                  t4        j6                         | j-                  |       | j                  j                  j/                  || j0                  |      S t9        |      rX| j-                  |       | j                  j                  j/                  |t;        j<                  dd| j0                        |      S || j                  j                  j>                  v }|r#| j                  j                  j>                  |   S i }tA        |      }|| j3                  t4        jB                         t	        |      dk(  r| j3                  t4        j6                         | j-                  |       t        |t        jD                        rS|jF                  rGt        |t        jH                  jJ                  jL                  jN                        stQ        ddd	g 
       t        |t        jD                        rPtS        |      rE| j                  jT                  rt        jV                  stQ        dddg tX        jZ                  
       t]        |      rt_        |      u|j`                  t_        |      j`                  k7  rSt_        |      }	|	tc        |	j`                        nd}
tQ        dd|j`                   d|
 dg tX        jd                  
       tg        || j                  d|      }| j                  j                  jh                  jk                  t;        j<                  dd| j0                        tm        |      ||      }to        | j                  ||       tq        d)| j                  ||||d|}| j                  j                  j                  js                  ||       |ju                         r-tg        |jv                  | j                  ty        |d      d       t4        jz                  }t        |t|              rt        |      rt4        j                  }t        j                  j                         xr. t        |t        j                  j                  j                        }|sA| j3                  t        j                  |t        |t              r|n
t        |                   t        |      r|rO| j3                  t4        jB                         |j                         d   }|ddgk7  rt        d      |j                         d   }t        |      dk(  rf|d   |j                  j                  k(  rJ|d   |j                  j                  k(  r.|d   |j                  j                  k(  r|d   |j                  k(  st        d      t        ty        | j                  d       j                  t4        j                               t        ty        | j                  d!      j                  t4        j                               no| j3                  t4        j                         | j3                  t4        jB                         t        t        |      j                  t4        j                               |j                         \  }}|D ]  }t        ||      }tm        |      }t        |t        j                  j                  j                        rtm        |j                        }t        |t        jD                        sBt        |      s7t        tm        |      j                  d"tm        |      j                   d#      ty        | j                  |      }t        j                   t        | j                  |      |              || j                  j                  j>                  |<   d$|j                  j                  vsJ t        |      |j                  j                  d$<   |j                  j                  j                  d%   }t        |      | j                  j                  urt        d&      t        ||d'|      }||j                  j                  d(<   |S )*NrN  Tr4  guardedr=  rX  z&Attempted to wrap strided NestedTensorr  3torch.compile does not support strided NestedTensorr  zAttempted to wrap sparse Tensorz-torch.compile does not support sparse TensorsNonez.dtype mismatch between tensor and its gradientztensor dtype: z; grad dtype: zmInconsistent dtype between tensor and its gradient. This can happen in FSDP and crashes meta tensor creation.rN  r"  r  )rN  r  r*  r+  r  _base)rN  r  r"  r%  r   _local_tensordevice_meshzDExpecting DTensor inner attrs to be ['_local_tensor', 'device_mesh']r      r=   rE  z\Expecting DTensor flattening ctx to be (placements, tensor_meta, shard_order, requires_grad)_specrequires_gradz found in tensor attrs of z[.__tensor_flatten__(). Only tensors and reference-type opaques are allowed in tensor attrs.r.  r*  z(Wrapped Tensor must be this graph's fakeFr>  r  )kr  rN  rO  rw  rx   r,  r  ry  r  r0  r1  rQ  r   r>   r  r_   r`   rb   r2  prepare_freezingis_specialized_nn_moduleis_fsdp_moduler  r  r$  rx  rG   r  r]   r  r  r5  infer_subclass_typer  r0  	is_nestednestedr  nested_tensorNestedTensorrF   r%   ri  capture_sparse_computer?   SPARSE_TENSORr  r&   r+  rB  r  wrap_to_fake_tensor_and_recordr  r  r  cache_real_value_when_exportwrap_fx_proxyr  _is_viewr  rQ   r6  r\   ra   NOT_NONE_MATCHrw  is_availabletensorDTensorr~  r  rf   r<   r9   __tensor_flatten__r  r^  r  
placementstensor_metashard_orderr  rH   r  r  DTENSOR_SPEC_MATCHrp  TENSOR_SUBCLASS_METADATA_MATCHrk   rb  r  r  r  r  r   r  r   r  rc  r  r  ro   r  r$   r  rD   r  )r*  r%  r  is_static_inputshould_install_free_tensormake_graph_attributeis_duplicate_tensoroptionsr+  safe_grad_valgrad_strr*  tensor_proxyr<  
guard_type
is_dtensorinner_attrsflattening_ctxattrsrX  attrinner_value
inner_typeinner_sourcefake_tensor_valuer>  s                             r  r  zVariableBuilder.wrap_tensor	  sg   " DGGNN777771%8Duehh0016..IIK""50E0G"H"O &,%@%@ &
!&) >&v.>5f= 	#  /  
!#Lu}}';';'L'L 	 &  99;?S''66811%877>>99tyy :   #5)Y6  5 5611%877>>99tyy :   f%11%877>>99'dii8 :   %(J(JJ77>>55f==+E2$ 7 78"5)Y6 5 56 	--e4 uell+uell&<&<&J&J&W&WX@Q	 uell+e$WW^^6+H+H
 9K8)778	 % % ,y/555%e,M3@3Ls=../RXHH(^H:NL6)556" 7dggf
 ww~~11DDFF#S$))4K	 E 
 	%TWWlEB' 
ww''
 
 	##99%Q>> +77!&'2	 "..
fj).Fv.N%44J&&335 
*5$$++33;

  !! &f.?@ *51	 )/ ##L$;$;<#668;?M"BB&^  "'!9!9!;A!>'1,&q)U[[-C-CC&q)U[[-D-DD&q)U[[-D-DD&q)U-@-@@&v  t{{G4??$77 t{{O<GG$11 ##L$O$OP##L$;$;<*62==l>W>WX //1HE1 %eT2 "+.
NN66GG "&k&:&:!;J!2:>&,5588R;//0 1++   *$++t<#//:ODGG\:;G)0 6E**62L$5$5$:$::::0DU0K}- ,1166;;OL019J9JJ*+UVVFE52CD-5z*r  c           	         t         J t        |t         j                        sJ t        | j	                               }ddlm} |j                  j                   }|r	 d|j                  _        d }t        j                         5  	 |j                  |      }|rddlm}  ||      }d d d        |J t)        j*                   t-        | j.                  |      |             t1        || j.                  d|	      }| j.                  j2                  j4                  j7                  t9        j:                  d
d| j<                        t?        |      ||      }	tA        | j.                  |	|       d|i}
tC        dtD        | j.                  |	|d d|
}|| j.                  j2                  jF                  |<   |jH                  jJ                  jL                  d   }tO        ||d|d|      }||	jJ                  jL                  d<   | jP                  |_(        |S # t        $ r* t        |j                  t         j                        sJ Y w xY w# t"        $ r$}t%        dt'        |      dg |       Y d }~d }~ww xY w# 1 sw Y   xY w)Nr   )_utilT)clone_preserve_stridesz)failed to convert numpy.ndarray to TensorzHException encountered when attempting to convert numpy.ndarray to Tensor)r  r  r  r  from_excFr  r=  rX  r4  r  )r)  rN  r  r*  r+  r*  r   r!  r"  r#  r>  r  ))r   r,  r  rf   r  torch._numpyr  flags	writeable
ValueErrorr  nditerr   temp_restore_stack_try_convert_to_tensortorch._prims_commonr  r  rF   rB  r   r  rc  rN  r  rO  r  r  r  r  r$  r  r  r4  r   r5  r  r  r  r  r  )r*  r%  r  r  readonlytensor_valuer  r{  r*  r  r  numpy_ndarray_variabler>  s                r  r  z"VariableBuilder.wrap_numpy_ndarray.
  s^   ~~%,,,"4??#45&{{,,,9(,%
 0CCE 	$;;EBJ#9,#GL	  ''' 	''(H(H(VW6ww	
 **==FF#S$))4	 > 
 	%TWWe\BV$!2 "
+ww'"
 "
 6L**62.4499>>O #%+
 '/


# )-%%%O  9 "%**bii8889 ' GJ j 	 	sB   H I:!I
/II
	I7I2,I:2I77I::JrS  r  c                 &   t        |      t        u sJ | j                  | j                  j                  j
                  v r-| j                  j                  j
                  | j                     S | j                  j                  j                  }d }t        j                         j                  ra|j                         }t        |       | j                  j                  j                  j                  t        || j                  d              n-t!        | j#                               s|dt$        j&                  j(                  j*                  r@| j-                  t.        j0                         t3        j4                  || j                        S | j                  j                  }t7        | j                  |t9        j:                  |      | j                  j<                  j?                               }tA        | j                  j                        }| j                  }	tC        |	tD              r|	jG                         }	||}
nt(        jH                  r)|jJ                  tL        u rtO        dd       tQ               }
ntC        |	tR              r6|	jT                  *tW        |	jT                        j                  |ddi      d   st(        jX                  stZ        j\                  }
nS|jJ                  tL        u rtO        dd       | j-                  t.        j0                         t3        j4                  |      S t(        j^                  r||j`                  nd }|jc                  || j                  |
|      }tC        |td              stg        d	t        |             | j                  j                  j                  j                  t        || j                  |             n[t!        | j#                               sJ | j-                  t.        j0                         t3        j4                  || j                        S tC        | j#                         th              rJ tk        | j#                         jm                  t.        jn                               d
| j#                         i}| j                  j                  jp                  js                  tu        jv                  dd| j                        t        |      || j#                               }|jx                  jz                  }tC        |t|        j~                        s
J | d       || j                  j                  jp                  j                  |<   t        j4                  | j                  ||fi |}|| j                  j                  j
                  | j                  <   t!        | j#                               s6t        | j#                         |dd d|      |jx                  j                  d<   |S )Nr  r  dynamo.automatic_dynamic_shapesTr   Fr%  )r  r:  excluded_valuer<  r  r=  rX  r4  r?  r  r>  )Dr  r  r$  rN  rO  unspec_variable_mapr  r   r}  #force_unspec_int_unbacked_size_liker  r,   r  r  r4   r  r]   r  r  r2  r>   r/  rx  rG   r  r   r  rN   rL   r  r0  r1  r)   r,  rS   rn  automatic_dynamic_shapesscalarrJ   r   $get_automatic_dynamic_shapes_mark_asrd   rS  r  assume_static_by_defaultr.   r  !automatic_dynamic_exclusion_guardexcluded_scalarr  r   r  rh   rH   r  r  r  r  r  r  r  r  r  r  bound_symbolsr   r  r  )r*  r%  rS  r  r  frame_state_entrywrapped_valuer$  normalized_source_namebase_sourcer:  r  r  r  r  
unspec_vars                   r  r  zVariableBuilder.wrap_symint
  s    E{c!!!99:::77>>55dii@@GGNN,,	8<CC%<<>M%m4GGNN((//M4;;= $DOO$56EMM$8$8$G$G
 ##L$?$?@'..U4;;OO;;##D 9#//6+/;;+C+C+^+^+`	! &;4;;;K;K%L"++K+}5)224#&//%,,< A4HBD;4((4--.223IAu:V 44(00 %++|;#$EuM##L$?$?@'..U;; ;;%1 "11 	  &JJ{{'.	 K M mV4$'<T-=P<Q%RSSGGNN((//M4;;@ &doo&7888 ; ;<#**t{{KKdoo/1BCCCdoo'22<3J3JKLT__./**==FF#S$))4??$	 > 
 !%%**(ELL1WhZ?V3WW1=B""00:$++DGGUMUWU
8B**4995!$//"34*2!#( #0+EJJOOJ' r  c           
      
   | j                   | j                  j                  j                  v r-| j                  j                  j                  | j                      S t	        | j                  | j
                  j                   t        j                  |      | j
                  j                  j                               }t        j                  j                  j                  st        | j                               st!        j"                  |      sqt!        j$                  |      s\t        j&                  j                  j(                  j*                  s.t-        dd      s"t        j.                  rR|j0                  t2        ur@| j5                  t6        j8                         t;        j<                  || j
                        S t        j>                  |t        j@                        }t        jB                  jD                  jG                  |      r@| j5                  t6        j8                         t;        j<                  || j
                        S tI        | j                         tJ              rJ tM        | j                         jO                  t6        jP                               tS        | j                               }||d}tU        || j                  d|      }| j                  j                  jV                  jY                  t[        j\                  dd	| j                         t_        |      ||
      }ta        | j                  ||       tc        td        f| j                  ||d d|}tI        |td              sJ || j                  j                  j                  | j                   <   | j                  jf                  r=tI        | j                         th              stk        d| d| j                                d }	|jl                  jn                  jp                  d   }ts        |      sJ |}	|	jt                  | j                  jt                  u sJ d|	jt                   d       tw        | j                         |d|	d|      |jn                  jp                  d<   ty        | j                   | j                  j                  jz                  ddgt}        |gi              }
| j                  j                  j~                  j                  t        |
j                  | j
                  d              t               j                  ddd       |
S )Nr  z.pytorch/compiler:unspecialize_float_killswitchFr  )r+  )r  	raw_valuer  r=  rX  r4  rN  r  r*  r+  =Dynamo attempts to add additional input during export: value=	, source=r*  fake mode (^) from fake tensor metadata doesn't match mode({self.tx.fake_mode}) from InstructionTranslatorTr  r>  call_methodr'  tensorify_float_attempt)	overwrite)Er$  rN  rO  r  rN   r  rL   r  r0  r1  r  r2  r>   r  r]   r  mathisnanisinf	_inductorr[  
cudagraphsr'   r  r  rJ   rx  rG   r  r   r  r  float64rt  ru  is_gradtrackingtensorr,  rh   rH   r  r  rZ   r  r  r  r  r  r  r  r4  r   ri  rd   r  r  r  r  r#   r  r  r  r  r   r  r  r4   sym_numr   rj  )r*  r%  r  r  r  r  r*  r  r  r  r-  s              r  r  zVariableBuilder.wrap_symfloat  s$    99:::77>>55dii@@5GGKK++E2'+{{'?'?'Z'Z'\
 MM  11!$//"34zz% zz%  %%,,77OQVW//%,,L@  ; ;<#**t{{KK U%--@
 8844]C ; ;<#**t{{KK doo/1BCCCdoo'22<3J3JKL #4??#45#%8 7dggv
 **==FF#S$))4	 > 
 	%TWWe]C&'
ww'
 

 *&ABBB8B**499577>>*T__->"L OP]^ghlhwhwhygz{  !"((--22?C}%%%) **dgg.?.?? 	
+556 7? ?	
?  '/OO#),'



# GG'DGGNN'' #J<4
 	$$++K		4;;PT,UV!!";TT!Rr  c                    | j                   | j                  j                  j                  v r-| j                  j                  j                  | j                      S t	        j
                  |      }t        | j                         t              s6t        | j                         j                  t        j                               d| j                         i}|j                  d|i       t        || j                  d| j                               }| j                  j                  j                  j!                  t#        j$                  dd| j                         t'        |      || j                               }t)        | j                  ||       t+        t,        f| j                  ||d d|}|| j                  j                  j                  | j                   <   t/        | j                               s| j                  j0                  r=t        | j                         t2              st5        d	| d
| j                                d }|j7                         r|j9                         }n#|j:                  j<                  j>                  d   }tA        |      sJ |}|jB                  | j                  jB                  u sJ d|jB                   d       tE        | j                         |d|d|      |j<                  j>                  d<   |S )Nr  r  Fr  r=  rX  r4  r  r  r  r*  r  r  Tr  r>  )#r$  rN  rO  r  r  r  r,  r  rh   rH   r  rG   r  updater  r  r  r  r  r  r  r4  r   r]   ri  rd   r  is_python_constantas_python_constantr  r  r  r#   r  r  )r*  r%  r  r  r*  r  r  r  s           r  r  z,VariableBuilder.wrap_unspecialized_primitive  s   99:::77>>55dii@@U+$//+->?$//+66|7N7NOPT__./U+,6dggt?P
 **==FF#S$))4??$	 > 
 	%TWWe]C&'
ww'
 

 9C**4995!$//"34ww~~j1BK&P$STaSbbklpl{l{l}k~  !%,,. * = = ? !+ 0 0 5 5 : :? K=))) -$..$''2C2CC /99: ;C CC
 +3!#'-#0+EJJOOJ' r  )Tr'  NN)Pr  r  r  __doc__rj   r@  rH  r(  r   rn  rj  r  r  r   r   r  rB  rx  classmethodr  r~  cacher  r  r  r   r  r  r  r   r  r  r  r  r  r   r  r   r  r  r  r  r~  r  r   r  r  r   r   r  r   r  r   r  r  r  r  r   r   r  r0  rQ  ry  rz  rg  r  r  r  r  r	   r  r  r  r  r.   r3   r  r  r  rU  rV  s   @r  rc  rc    s   = %)	7#7 7 "	7
 
7&f  > >? >@
O 
 
F hsCx&8 T#s(^d=R  ;tFHS#X,>$>? ; ; __'d 'tFHSRUXDV<V7W '  'R+

3 +<P +
I'"7"7"< I I
[[&&66
	

s 
/J 
"H H HB __	c8.4oEFF	G  6R13 R1? R1h"Ps P P.p5c?DI|ZOPp	pd	H 	HO 	HH HO HH# H/ H <eem <8U <u|| D T $h h_ hT?& ?_ ?B'3 '? ')U ) )  3%01 
	   5'?23 
	   U{T!C'  
	  <@FU{T!C'F 3%01D8F 
	Fell t W W/ WrS& S& S&p '+04	AA t#A '-	A
 
AFO5 O_ Ob># >/ >r  rc  objc           	         d }t        | t              r| j                  }n"t        dd|  dt	        |        d|  dg        |J g }t        j                  |      D ]X  }d }| j                  r,t        | j                  d      }t        ||j                        }|j                  t        ||             Z t        |      S )	Nzdataclass fields failurezobj: z; variable type: z$Dataclass fields handling fails for z*. Expected it to be a user-defined object.r  __dataclass_fields__r4  )r,  r   r%  rF   r  r  r  r  rQ   rW   r$  r  r   )r	  r%  r  fieldr  base_srcs         r  r  r    s    E#01		.C5 1$s)=>seCmn		
 E##E* F::!#**.DEH&x<F.uVDEF r  r%  r  c                     t        | t        j                        rIt        | t              s9t        j                  |       rt        |       |u s| j                  st        |       } | S r'  )r,  r  r0  r!   _is_functional_tensorr$   r  rr   )r%  r  s     r  _clone_inputr    sQ    %& uj) ++E2'.);  &ELr  rN  r  r  r*  r+  r  c                     | |||d|}|t        ddt        i|S t        ddt        i|}|j                  |        |S )Nr  r)  r  )r4  r   r   install_global)rN  r  r*  r+  r  kwargsr  s          r  r  r    sd     &&	
 F  ENEfEE #U.JUfUb!r  c                     | j                   r_t        j                  j                         5  t	        || j
                        |j                  j                  |j                  <   d d d        y y # 1 sw Y   y xY wr'  )	ri  r  rt  DisableTorchFunctionSubclassr  r  tracerreal_value_cacher  )rN  r  r*  s      r  r  r    sb     
yy XX224 	8Dr||9ELL))%**5	 		 	 	s   8A--A6r)  c                    |t        | ||||fi |}n;t        |t        j                        rt	        | ||||fi |}nt        ||||||       }t        |t        j                  j                  j                  t        j                  j                  j                  f      r>|j                  j                  dk7  r%|j                  j                  j                  |       |S )Nplaceholder)_wrap_fx_proxyr,  r  r0  _wrap_fx_preexisting_tensorhandle_traced_outputr2  	variablesr   r   r  oprO  rt  record_proxyable_vt)r)  rN  r  r*  r+  r  outs          r  r4  r4  T  s     )E=-
CJ
 
M5<<	0)E=-
CJ
 #2ug}j

 	''66''77	
 JJMM]*
		  44S9Jr  r  c                 z   ddl m} t        |t        j                        sJ dt        |              t        ||      sJ d|v r-|d   (|j                  j                  j                  |d          |j                  j                  dk(  r#d|j                  j                  v sBJ d| d       d|j                  j                  vsJ |j                  j                  d           t        j                  j                  j                         5  t        |      |j                   u rnt#        |||       |j$                  r^t        j&                  j)                         5  t+        ||j                         |j,                  j.                  |j                  <   d d d        d	| t0        t2        fv i}d
|v r|d
   J |d
   |d
<   t5        |fd|i|}|j6                  j
                  dk7  r%t        |      |j                   urt9        d|       d d d        t;        | |||||      S # 1 sw Y   xY w# 1 sw Y   %xY w)Nr=   r  z1_wrap_fx_preexisting_tensor expected tensor, got r  r  r*  zplaceholder z* doesn't have 'example_value' in node.metar"  r  rN  r  zN`tensor` needs to be a `FakeTensor`wrapped by this instance of Dynamo. Found: )symbolic_convertr  r,  r  r0  r  rO  r  r  r  r  r  r2  r  +_disable_saved_tensors_hooks_during_tracingr$   r  r  ri  rt  r  r  r  r  r   r   r  r  rD   construct_tensor_variable)r)  rN  r  r  r+  r  r  r  s           r  r  r  |  s9    =fell+ 
;DL>J+ b34447wx0<
		 12 zz}}%%**//1 	
5'!KL	
1 ejjoo5 	
zz/0	
5
 
			H	H	J "v&",,6(UF;yy XX::< @LAELL11%**= Z"$@ABF w&78+<+HHH&x0F83FLrLVLF==''r||;*>>DXG ?"H %Bv}g 5 " "s&   AH18H%A7H1%H.	*H11H:c                    ddl m} t        ||      sJ d|v r-|d   (|j                  j                  j                  |d          d|j                  j                  vsJ |j                  j                  d           t        j                  j                  j                         5  t        |j                  |d      }d d d        t        ||||||       S # 1 sw Y   xY w)Nr=   r"  r  r*  T)allow_non_graph_fake)r#  r  r,  rO  r  r  r  r  r  r2  r  r$  rv   r  )r)  rN  r  r*  r+  r  r  s          r  r  r    s     =b34447wx0<
		 12%**//1Xejjooo6V5WX1 
			H	H	J R 'uzz2DQ	R  
 R Rs   %CCc                     dd l }dd l}dd l}t        | |j                        rt        |       rD|j                  rt        j                  s(t        dt        |       dg t        j                         t        |||| ||      }t        |t              sJ |j                  j                   j#                  ||t$               |S t'        |j(                  j*                  d      ra|j(                  j*                  j,                  dk(  r>t        |j(                  j*                  j.                  |j0                  j2                        s,|j(                  j*                  |j4                  j6                  u rFt9        |j(                  j*                        t        usJ t;        |j(                  j*                        S |j(                  j*                  |j0                  j<                  u s,|j(                  j*                  |j>                  j@                  u rtC        |       S tE        | |jF                        r@tI        d | D              r.| D cg c]  }tK        jL                  |       }	}tO        |	fi |S t        | tP        tR        f      rtU        |j(                  |        g }
tW        |       D ]  \  }}|%|
jY                  tK        jL                  d#i |       -|jZ                  j]                  d	t^        j`                  ||fi 
      }d|v r9t        | tR              sJ |d   }|jc                         }te        ||d      |d<   n|}|
jY                  tg        d$||||d|        t        | |jF                        rtO        |
|fi |S tE        | tP              rti        |
fi |S tE        | tR        tj        f      rtm        |
fi |S to        |       sJ dt9        |               ti        |
|jq                  dts                           } tu        jv                  t9        |             | fd|i|S | "|j(                  j*                  |jx                  u rtK        jL                  d#i |S t        | |jz                  |j|                  |j~                  f      rT|j                  j                  j                  | |       tU        |j(                  |        t        jL                  ||| fi |S t        | |j                        r|j(                  j*                  t        u s=|j(                  j*                  t               D cg c]  \  }}|j                   c}}v r[tU        |j(                  |        d }|j(                  j*                  t        u r|j(                  j                  d   }t        || |fi |S t        | |j                        r|j(                  j*                  t        u s=|j(                  j*                  t               D cg c]  \  }}|j                   c}}v r[d }|j(                  j*                  t        u r|j(                  j                  d   }tU        |j(                  |        t        || |fi |S t        j                  |j(                  j*                        r*t        |j(                  j*                  |j                        s=|j(                  j*                  t               D cg c]  \  }}|j                   c}}v r$tU        |j(                  |        t        || d fi |S |j(                  j*                  dk(  r;|j(                  j                  dk(  r"tU        |j(                  |        tK        | fi |S | lt        | |j                        rV|j(                  j*                  dk(  r=|j(                  j                  dk(  r$tU        |j(                  |        t        || d fi |S t        | t              rv|j(                  j*                  |j                  |j                  |j                  t        t^        j`                  |j                  j                  |j                  t^        j                  |j                  j                  j                  |j                  j                  j                  |j                  j                  j                  t        |j                  dt              t        |j                  dt              |j                  fv s2|j(                  j                  dk(  rE|j(                  j*                  dk(  r,tU        |j(                  |        tK        jL                  | fi |S t        | |j                  j>                  j                        r%ddlcmd} tU        |j(                  |         ||fi |S t        | t              rr|j(                  j*                  |j0                  j                  |j0                  j                  j                  |j                  j>                  j                  |j                  j>                  j                  |j                  j>                  j                  |j0                  j                  |j0                  j                  |j0                  j                  |j0                  j                  |j0                  j                  |j0                  j                  dgtS        t        j                               z   v r,tU        |j(                  |        tK        jL                  | fi |S t        | t        t        t        f      r|j(                  j*                  t        u sN|j(                  j*                  t        u s2|j(                  j                  dk(  rE|j(                  j*                  dk(  r,tU        |j(                  |        tK        jL                  | fi |S t        | t              s|j(                  j*                  dv r,tU        |j(                  |        tK        jL                  | fi |S t        | |j                  j                  j                        rt        jL                  ||       S t        t9        |             rzt        t9        |             rt        jL                  | |       S |j                  j                  j                  |j                  j                  |       }t        jL                  ||      S t        ddt        |        d |j(                  j                   d!|j(                  j*                   d"g        y c c}w c c}}w c c}}w c c}}w )%Nr   z4Attempted to wrap sparse Tensor with VariableTrackerzBtorch.compile does not support sparse Tensors with VariableTrackerr  )mutation_type_clsr  	set_statec              3   <   K   | ]  }t        |t                y wr'  )r,  r  r  s     r  r_  z'handle_traced_output.<locals>.<genexpr>!  s      3 
1c3s   rR  )kindtargetr&  r  r  Fr,  )r)  rN  r  r*  z)expected namedtuple or structseq but got r  r  r  queryr  record_eventget_rankget_world_size
bit_lengthr   r   
is_integerr'  )r  	__round__ztorch.* op returned non-Tensorzexample_value type: z; op: z
; target: zUtorch.* ops that return a non-Tensor cannot be traced into the Dynamo FX graph outputr'  r  )torch._functorch.vmaptorch._subclasses.fake_tensortorch._utilsr,  r0  r%   ri  r>   r  rF   rB  r?   r  r%  r   rO  rw  
_track_objr   r  r  r-  r  r-  rt  	Generatorr  set_rng_stater  r   _DisableFuncTorchr  _is_in_bad_forkr   r   r  r  r   r  r   r  r  r   rm  r  r  r  operatorgetitemrB  r[   r4  r   r6   r   r~   r}  r   r   rc  manual_seedr   SymFloatr  rt  track_produced_symintsr   r  r   rC   current_streamr&  r   r  r   rZ  r[  r  r  r  sym_intsym_maxsym_minrb  _utils_element_sizeseedmodru  vmap_validate_and_get_batch_sizepredispatch_vmap_increment_nesting_vmap_decrement_nestingrw  r  _constrain_as_sizebackends
SDPAParamssdpar   r@   _are_functorch_transforms_activeis_batchedtensoris_flash_attention_availablecan_use_flash_attentioncan_use_efficient_attention_get_cudnn_sdp_enabled_get_flash_sdp_enabled_get_mem_efficient_sdp_enabled_get_math_sdp_enabled_get_overrideable_sdp_enabled_is_autocast_availabler   r  r  r   r   r  r  r  r   r   r   r  r  r   )r*  rN  r  r  r+  r)  r  r  r  sizesunpackedr  r  proxy_ir  	options_ir  rX  device_interfacer  r   r  s                         r  r  r    s
    !(-. '		!>!>NM*`8)778	 (E=-
 #///
		))3*> 	* 	
 


!!:.JJ&&+5uzz((115883E3EF:: : ::EJJ%%&c111+EJJ,=,=>>

UXX777::

 : ::(77	uzz	*s 3$13 0 6CC!((+CCE-W--	ME4=	1%**m4. (	FAs{$++<G<  ,,33(#++	 4  w& &mT:::$X.F 'I*7#',+Ih' !(I % #-%&) $
=(	R mUZZ0  %;7;;M5) 5W55MD.#9:4G44 / ;D<O;PQ/ %%kk/;K;MNH L+66tM7JK!  
 
	%**"3"3u7H7H"H&&7w77	MELL%..%--#P	Q
		  77uM%**m4%%b%J'JJ=%,,/JJ!==			#C#EA 	'' 
 	%**m4:: <<JJOOA&Ee]EEWEE=%++.JJ!==			#C#EA 	'' 
 :: <<JJOOA&E%**m4UM5DGDD

))*uzz((%++6			#C#EA 	 
 	%**m4UM4C7CC			g	%%**--=*H%**m4999!}ekk2JJ/JJMM]*%**m4UM4C7CC	M3	'

MMMMMMLL&&JJLL!!>>((@@((@@E%%z8<E%%'7B $$#
	
, JJMM]*uzz/@/@L/P 	%**m4&&}@@@	M5>>#6#6#A#A	B,%**m4!%3733	M4	(

HH55HH00NN<<NN77NN;;HH++HH++HH33HH**HH22HH++
 388:
;<	<" 	%**m4&&}@@@	MC#5	6

^+::
*JJMM]*uzz/@/@F/J%**m4&&}@@@	M5	)UZZ->->BV-V%**m4&&}@@@	M5>>#E#E#V#V	W )//
 	
 
]+	,] 34,33   ..<<NNII
 )//
 	

 	4*7=+A*B&Wabgblblbsbsatuo		
[ DXs   u)=u.u42u:c                 &   t        |       t        j                  t        j                  j                  t        j
                  j                  j                  t        j
                  j                  j                  fv st        |       ry t        |       S r'  )r  r  r0  ry  r  rA  r!  r!   r  r  r9   r%  s    r  r  r    sj    E{%%00++<<	 
 
'u	-$ E{r  c                    | j                  |      }t        |t        j                  j                  j
                        r|j                  |j                  u r|r|}nt        |t        j                  j                        rt        j                  j                  }nOt        |t        j                  j                        rt        j                  j                  }nt        j                  }||d<   |S )N
class_type)
specializer,  r  rA  r!  r!   r  ry  r  r  r0  )r)  rN  r*  r+  specialized_propstensor_types         r  get_specialized_propsrj    s     #--m< 	=%"3"3"?"?"J"JK##r||3'Kuxx'9'9:((,,Kuxx7((//K,,K*5,'r  c                 $   t        ||j                        }t        |j                  |       |j                  j                  dk7  r&|j
                  j                  j                  ||       |j                  t        | |||              | |fi |S )z
    Actually construct a tensor variable after all the pre-processing from
    wrapping a pre-existing or newly created tensor value.
    r  )
r  r  r   r  r  rO  rt  rB  r  rj  )r)  rN  r  r*  r+  r  s         r  r%  r%  /  sv     !=Mejj-0
 zz}}%
		  77uMNN(RVWe'w''r  c                      t         j                  dk(  rt        j                  S t         j                  dk(  rt        j                  S t        dt         j                         )Ndynamicunbackedz+invalid automatic_dynamic_shapes_mark_as = )r>    automatic_dynamic_shapes_mark_asr.   r  r  r  r  r  r  r  r  K  sT    ..);!!!		0	0J	>"""9&:a:a9bc
 	
r  _DYNAMIC_SOURCES_DYNAMIC_SOURCES_CONFIG_HASHc                  F   t        t        j                  j                  j                        } t
        t        | k(  rt
        S t        j                  j                  j                  j                  dd      j                  d      D ch c]  }|r| c}a| at
        S c c}w N r  ,)	hashr  compilerr>   dynamic_sourcesrp  rq  replacesplitcurrent_hashss     r  get_dynamic_sourcesr~  Z  s     --==>L #(D(T
 &&66>>sBGMMcR 	

 $0    Bsource_namec                     t               }|D ]6  }|| k(  st        j                  ||       st        j	                  d| |        y y)NzA%s was marked dynamic due to dynamic source allowlist pattern: %sTF)r~  r  matchr  r  )r  rx  patterns      r  r  r  n  sM    )+O" k!RXXg{%CIIS
  r  r$  r{  c           
         |j                         }t        |      s|j                         }|j                         }d g|z  }t	        |      D cg c]
  }||   | f }}|j                  t               i }	|D ]>  \  }
}| }|	j                  |
|
      ||<   |	j                  |
||   z  t        |             @ ng }t        | |t        j                  t        |      t        |                  S c c}w )N)r  )sizer%   stridedimr  sortr-   r}  
setdefaultrM   rN   rL   make_tensorr  )rN  r$  r{  ex_size	ex_strider  r  r  pending
candidatesi_strideneg_is               r  record_automatic_dynamicr  {  s     ffhGHHJ	eeg#/4Sz:!IaL1"%::/0
& 	IOHeA"x:F1I!!(WQZ"7QH	I $
''gfF  ;s   C1_UNBACKED_SOURCES_UNBACKED_SOURCES_CONFIG_HASHc                  F   t        t        j                  j                  j                        } t
        t        | k(  rt
        S t        j                  j                  j                  j                  dd      j                  d      D ch c]  }|r| c}a| at
        S c c}w rs  )	rv  r  rw  r>   unbacked_sourcesr  r  ry  rz  r{  s     r  get_unbacked_sourcesr    s     -->>?L $)F,)V  
 &&77??RHNNsS 	

 %1!r  c                     t               }|D ]6  }|| k(  st        j                  ||       st        j	                  d| |        y y)NzC%s was marked unbacked due to unbacked source allowlist pattern: %sTF)r  r  r  r  r  )r  r  r  s      r  r  r    sN    +-# k!RXXg{%CIIU
  r  r  static_shapes
outer_onlyc                 F  12 | j                   rGt        | t        j                  j                  j
                  j                        st        dddg        |j                  }|j                  j                  j                  j                  | d       }|r|j                  ni }d }| j                         r$t        |d      }	t!        | j"                  ||	|      }t%        |       r|st!        | |||d      }
t        |
t&              sJ i }t)        |       j+                  |       \  }}|D ]l  }t-        | |      xxt        j.                  d x"\   } t        ||      }t!        ||||      ||<   E  xt0        d x\    U }t3        d	t)        |              t5        |
j6                  |
j8                  |
j:                  |
j<                  ||
j>                  |
j                  |
      S |rtA        |      svt'        tB        jD                  g| jG                         z  tB        jH                  g| jG                         z  d g| jG                         z  d g| jG                         z  |||      S ddl%m&2 tO        2fd| jQ                         D              rt'        | jQ                         D cg c]2  }t        |tR              rtB        jT                  ntB        jD                  4 c}tB        jH                  g| jG                         z  d g| jG                         z  d g| jG                         z  |||      S tW        |||       }tY        |       }i 1dtZ        dddt\        dd f1fd}ddl/m0} |j                  jb                  `|j                  jb                  D ]G  }t        ||      r|jd                  |k(  s  ||jF                  |jf                  |j                         I g }g }g }g }g }ti        | jG                               D ]]  }|t-        | dtk                     v }|t-        | dtk                     v }|t-        | dtk                     v } |t-        | dtk                     v }!|t-        | dtk                     v }"|jm                  t-        | di       j                  |g              to        |j                        }#|}	t        |	tp              r|	js                         }	| sFt        |	tt              r|	jv                  zty        |	jv                        j                  |#|di      |   rPtz        j}                  d|       t~        g| jG                         z  }$t        |$|<   |t        j                  |$      z  }t        j                  xr |j                  |      }%t        j                  xr |j                  |      }&tA        |      rtz        j}                  d|       d}%t        |      rtz        j}                  d |       d}%|%xs |&}'1j                  |      }|d }(d })| rt        j                  sd })t        | d!      r| j                  D *cg c]  }*|*jF                  |k(  s|* c}*j                         }+|+j                  |+j                  t        d"      }(ndd#l%mO},  |,t        |+j                  |+j                  $      d%      }(nt        d"      }(n|rt        d"      }(n|"s+|'r)t        d&d       |%rt        d"      }(|&rbt        d"      })nU|"st        j                  st        d&d       d }(d })n2|\  }(}-d })| d'| d(}.|-|j                  j                  j                  |.<   |jm                  |(       |jm                  |)       |st        |      rtB        j                  }/nr|(| s|!s 2| jQ                         |         r|'rt               }/nFtB        jT                  }/n5|st        j                  s|"rtB        jD                  }/ntB        j                  }/|)tB        jT                  }0ntB        jH                  }0|jm                  |/       |jm                  |0       ` t'        ||||||||t-        | d)d       t-        | d*d       |j                  +      S c c}w c c}*w ),NzGEncountered strided NestedTensor in automatic dynamic dim determinationr  r  r  r  T)r  r  z#expected Tensor or OpaqueBase, got )dynamic_sizesdynamic_stridesconstraint_sizesconstraint_stridesview_base_contexttensor_source#shape_env_to_source_to_symbol_cacheinner_contexts)r  r  r  r  r  r  r  r   )is_nested_intc              3   T   K   | ]  }t        |t              xr	  |        ! y wr'  )r,  r   )r]  r}  r  s     r  r_  z%_automatic_dynamic.<locals>.<genexpr>  s(     
La:a 9q)9%99
Ls   %(r  constraint_rangeStrictMinMaxConstraintr$  r	  c                     | v r=ddl m} |    \  }} ||j                  |j                  z  d      }|xs |}||f| <   y ||f| <   y )Nr   r  FvrrQ  )%torch.fx.experimental.symbolic_shapesr  r  )	r  r  r$  r  old_constraint_rangeold_namenew_constraint_rangenew_namedim2constraints	           r  update_dim2constraintz1_automatic_dynamic.<locals>.update_dim2constraint)  sm     . T-;C-@* (#9#&&)=)@)@@$   '4H"6"@N3"2D"8N3r  )_RelaxedConstraint_dynamo_strict_unbacked_indices_dynamo_unbacked_indices_dynamo_dynamic_indices_dynamo_weak_dynamic_indices_dynamo_static_indices_specialize_onFz#automatic dynamic %s marked dynamic)r  r  r  _dynamo_dynamic_rangerP  r  )lowerupperr  r  z.size()[]_dynamo_shape_ids_dynamo_unbacked_bounds)r  r  r  r  specialize_onr  r  r  	shape_idsunbacked_boundsexcluded_sizes)Yr  r,  r  r  r  r  r  rF   r$  rO  tracing_contexttensor_to_contextr}  r  r  rQ   _automatic_dynamicr  r9   r0   r  r  rb  r0  r   r  r1   r  r  r  r  r  r  r.   r  r  INFER_STRIDEr  r  rk  r  r   r  r  r\  r  rB  r  r  export_constraintst_idr  r  rj  r  r)   rS   rn  rd   rS  r  r  r  rK   rJ   rL   	make_sizer>   r  is_size_dynamicis_stride_dynamicr  allow_ignore_mark_dynamicr  r  popminmaxr/   r  r;   r   r  source_name_to_debug_namer  r  r  DUCKr  )3r{  rN  r  r  r  r$  prior_policyr  r  r  outer_contextr  r  rX  r  r  r  
unexpectedr}  r  r  r  r  rR  r  r  r  r  r  r  marked_strict_unbackedmarked_unbackedmarked_dynamicmarked_weak_dynamicmarked_staticr  	mark_sizeautomatic_dynamic_sizeautomatic_dynamic_strideautomatic_dynamicconstraint_sizeconstraint_stridedr	dim_ranger  name_dim_namedynamic_sizedynamic_strider  r  s3                                                    @@r  r  r    s    	{{:	5<<!!//<< 	]M		
 ;;D99,,>>BB1dKL<H88b (
 15zz| 1.qwwKW$Q'
*r6=T
 -)@AAA 7--a0q 	D!T"2U\\^2#-fd#;L+=#R},N4( 3
 "Z\ "(=d:>N=OP 	 ''55)99*;;,??/'550=0a0a)	
 		
 .t4&%,,-7'445?"Vaeeg- $v// 0S
 	
 D

L1668
LL&  '1F&;
""ARARR (445?"Vaeeg- $v// 0S
 	
 1T1= a5DN99$<9DG9	9$ ?	yy##/))66 	J*&89$&%NNJ$?$?		 MOM1557^ E/!"g0#%'
 "
 wq*DceLLga)BCEJJ71.Lce#TTWQ(@#%HHWQ(8"=AA!RHI
 "7v{{!Ck=1%..0K{K0$$0[))*../E5zRSTU II;TB#quuw.I'IaL!4!>!>I!NN ++T0A0Q0QRS0T 	 ++V0A0S0STU0V 	! T"II>E%)"d#II?F%)"2N6N
 $''*
"O $f&F&F$(!156%&%<%<!!!!ce  !}},1F*AE*R +A*immT&++
 '>&NO'"9E"J"'8 A4H)&=&MO+(?$(O%$V-L-L#$EuM"&$(!%/"OU $xs!,HFKBII99(C0!!"3406%..L'"QVVXa[)
 !CE)11f==%,,L &??L('//N'44N\*~.KE/N ##')-#+,O!0$7#<dC(77 z!s   *7bb!b)parent_contextr"  r  c                N   t        j                         }	 t        | ||||      |j                  j                  xj
                  t        j                         |z
  z  c_        S # |j                  j                  xj
                  t        j                         |z
  z  c_        w xY w)Nr  r"  r  )rh  ri  $_wrap_to_fake_tensor_and_record_implrO  rk  !wrap_to_fake_tensor_and_record_ns)r{  rN  r  r"  r  rm  s         r  r  r    s|     ,,.C
3r&In
 			**LLLLNS 	
L		**LLLLNS 	
Ls   A$ $A B$c          
          t               t        j                  t        j                  j                  t
        fv s&t         t        j                        st               rJ t         |      \  }}|st         |      n-t        t              sJ j                  }|j                  |   t        j                  dj                  t!         j"                        t                      t%               5  j&                  J t)         fd      }d d d        vt        t
              rf|j*                  x}	Xt        t,              rj.                  }
nt1              }
j2                  j4                  j7                  t9        |	|
             t              rf|j;                         \  }}|D ]N  }t=        ||      }t=         |      }t        |      }t?        ||t        |t        j                               P j2                  j@                  jB                   <   tE        |      r|jF                  r|jI                         n|jK                         }|jM                         d|jN                  z  |jM                         |jQ                         dj2                  jR                  <   n8|jM                         |jQ                         dj2                  jR                  <   |r~|rjU                         sltW              saj2                  j4                  j7                  t9        |             j2                  jX                  t[                  j7                         |S  S # 1 sw Y   _xY w)N)r  zwrap_to_fake %s %s %s %sc                  @    j                   j                         S )N)r  symbolic_context)r  r  )r{  r  r  rN  s   r  rM  z6_wrap_to_fake_tensor_and_record_impl.<locals>.<lambda>)  s$    00!%5 1  r  r  )r   )r  r  values_sizevalues_stride)r  r  ).r  r  r0  ry  r  r!   r,  r9   r   r  rQ   r^  r  r  r  r$  r  shaper   r  r   	item_memorZ   r  rR   rO  r  r  r4   r  rb  r  r  r  r%   	is_sparse_valuesvaluesr  ndimr  input_source_to_sizes_stridesr  r]   tracked_fakes_id_to_sourcer\  )r{  rN  r  r"  r  r  _reasoninner_context_namefake_esym_valitem_sourcer  rX  r  
fake_innerinnerr  r  r  s   ```               @r  r  r    s    	QELL%(("4"4jAAa&(+!!!!? "
w 1!RO fj111!'-<<=OP		&KK!''NG	
 &' 	<<+++(F	 6:.",,,9 &"34$kk26:II##**G[2BC )0002HE1 
$VT2
4()&$7.'(U\\B#3	
 :J		!!33A6  *0)9)9V^^%v}}F ,%{{}!'?BII33F;  --/?BII33F; "v'F'F'H&v.II##**FF,<= II00A7>>vFO	 	s    M>>Nc            
          e Zd ZdZddZeedddeee	      ee
e	e	f      z  eee	df      z  eee	      z  defd	              Zeedddee	   defd
              Zeedddee	df   defd              Zeedddeez  ez  ez  defd              Zeeddde	defd              Zeddde	defd       Zededefd       Zede
eede	gef   f   fd       Zy)ro  a  
    Like builder, but stateless and does not require a source. Useful for simple type->VT objects, or objects
    that are being created/evaporated during inlining (ex: consider a locally made list of tensors we then iterate over
    .), such a list should not show up as an artifact from inputs, nor in reconstruction, nor in the graph. However,
    there may be reasons to represent it as a ListVariable internally.

    NOTE - Objects produced here are born UNGUARDED due to the nature of sources!

    NOTE - This class is very new! It will have some rough edges, but it was created to stem the bleeding of giant
    if/else type->VariableTracker trees that were cropping up all over dynamo.
    r	  Nc                     t        d      )NzUse SourcelessBuilder.create()r  r1  s    r  rH  zSourcelessBuilder.__init__}  s    =>>r  rN  r  r%  .c                      y r'  r  rN  r%  s     r  r  zSourcelessBuilder.create  s     r  c                      y r'  r  r  s     r  r  zSourcelessBuilder.create  s    SVr  c                      y r'  r  r  s     r  r  zSourcelessBuilder.create  s     r  c                      y r'  r  r  s     r  r  zSourcelessBuilder.create  s     r  c                      y r'  r  r  s     r  r  zSourcelessBuilder.create  s    PSr  c           
      >   t        |      }t        j                  j                  |      }|r	 || |      S t	        |t
              r|S t        t        |            r0t	        |t        j                        st        j                  ||      S t        t        |            rTt        j                  j                  j                  | j                   j"                  |      }t        j                  ||      S t	        |t$        j&                        rt)        |      S t+        j,                  |      rt+        j                  |      S t/        |      rVt1        j2                  |      At1        j4                  |      rd| j                   _         t1        j2                  |      |      S t/        |      r[t9        j:                  |      rF t1        j2                  |j<                        |j<                        }t?        |dt        |            S tA        |      r t1        jB                  |      |      S t	        |t        j                  t        jD                  t        jF                  jH                  jJ                  f      rt)        |      S t	        |t         tL        jN                  f      r?tQ        t        |      t               rtQ        |tR              rtU        |      S t9        |      S t	        |tV        jX                        rt[        |      S t	        |tV        j\                        rt	        |j<                  t         tL        jN                  f      rvt_        |j<                  |j`                  jb                        |k(  sJ t        j                  | |j<                        }	 |je                  | |j`                  jb                        S | j                   jh                  jj                  j                  tm        |j<                              }|)t        jn                  jp                  js                  |j`                  |      S t	        |t        jt                  jv                  jx                        rt{        |      S t	        |t        j|                  j~                  j                        rt)        |      S t	        |t        j                        rt        |      S t	        |t        jn                  jp                  j                  j                        r	 t+        j                  t        |            S t	        |t        t        j                  j                              r3t        jn                  jp                  j                  j                  |      S t	        |tV        j                  tV        j                  f      rt        |      S t        |      rzt        t        |            D cg c]"  }t        j                  | t_        ||            $ }	}t        |	t                     }
 t        j                  t        |            ||
      S t	        |t        j                        ro|j                  j                  | j                   j                  v rC| j                   j                  |j                  j                     }t        j                  | |      S t	        |t              r9dD cg c]"  }t        j                  | t_        ||            $ }}t        ||       S t        |t              rt        |      S t        d|j                   d|j                   d	| g t        j                  
       y # tf        $ r Y Gw xY w# t        jn                  j                  j                  t        jt                  j                  j                  j                  f$ r= t        j                  |j                  j                         |j                  gi       cY S w xY wc c}w c c}w )NTr  r/  r.  )r  rI  z%Unexpected type in sourceless builderrT  z3SourcelessBuilder.create does not know how to wrap r  )lr  ro  _type_handlersr}  r,  r   r   rq  rr  r   r  r   r  r  r  r  rO  r  r  _HAS_DEFAULT_FACTORY_CLASSr   r   r  ry  rB   r  r  r  r   is_supported_new_methodr-  r   r{   r  rs  rt  ru  rv  abcABCMetar  r  r   r  r  r   r  rb  r*  r  var_getattrr  rw  id_to_variabler\  r2  r  UserMethodVariablerj  graph_moduleGraphModuler   r  r"  TreeSpecr  r  r   lazyLazySymNodeFormatStringrB  exc	UserErrorrk  symbolic_shapesGuardOnDataDependentSymNoder   fmt_varr  sym_node_varr  flex_attention_backwardhigher_order_ops&FlexAttentionBackwardHighOrderVariableGenericAlias	UnionTyper   r~   r   r   r   r   rc  r   r  r  r  r   r  r   r   r(  r   rF   r  r  r?   
DYNAMO_BUG)rN  r%  
value_typefast_handlerr  r	  
cls_obj_vtobj_vtr$  rO  r  r  r  r  s                 r  r  zSourcelessBuilder.create  s)   %[
(77;;JGE**e_-L!$u+.z%7S,33E5AA%d5k2#nn@@RR		##UO -33 
 {EEF,U33((/#**511e_!<!<U!C!O..u5>B		;5;..u5e<<e_!9!Q!Q"

 >+--enn=ennMC"3	4;GG#E*,;%%e,U33YY))588+>+>+L+LM
 -U33ckk23$u+t,E=1Q8??+E22u667(//u//0%..4*=>u~~u~~/F/FG5PPP.55b%..I
%11"enn6M6MNN
 //>>BB2ennCUV% ==22EE  uxx44@@A077u{{22;;<,U33rzz*'..u}}66;;SST'..s5z:: tE$;$;$S$STU==**;;bb   2 2EOODE!%((5! .d5k: "((WUD-ABF  %V;K;MNHC+66tE{C  uell+

299#:#::II++EJJOO<E"))"e44u% 3 "((WUA->?E  !++E6"!%((;!,,-Qz/F/F.GHMj\Z1%001		
u + * !!++%%55QQ  ,22MM446''( 	 s1   !%]1 ^  $'`%'`1	]=<]= B``c                 \    t        j                  |       sJ t        j                  |       S )Nr%  )r   r  r  r%  s    r  wrap_constant_literalz'SourcelessBuilder.wrap_constant_literal  s'    **5111&&U33r  r  c                     t         j                  i } t        D ]  }d | |<   
 fd| t        <   fd| t        <   fd| t
        <   fd| t        <   fd| t        <   fd| t        j                  <   | t
           | t        j                  <   | t
           | t        <   | t           | t        <   fd| t        j                  <   d	 | t        j                   <   d
 | t"        j$                  <   d | t&        j(                  <   d | t        j*                  <   d | t        j,                  <   d | t        j.                  j0                  j2                  <   d | t        j4                  j6                  j8                  <   d | t        j4                  j6                  j:                  <   d | t        j4                  j6                  j<                  <   dddt>        dt>        fd}t@        jB                  D ]  }|| |<   	 | S )Nc                     t        |      S r'  r   r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>*  s    ,<U,C r  c           	      b    t        |D cg c]  } | |       c}t                     S c c}w Nr.  )r   r   rN  r%  r  r  s      r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>+  s)    +$)*qVB]*:J:L+
 *   ,c           	      b    t        |D cg c]  } | |       c}t                     S c c}w r+  )r   r   r,  s      r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>.  s*    1C$)*qVB]*:J:L2
 *r-  c                     t        |j                         D ci c]  \  }} | |       | |       c}}t        |      t                     S c c}}w r+  )r   r  r  r   rN  r%  r  r  r  s       r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>1  sF    +<6;kkmDdaVB]F2qM)DK*,,
 Ds   A
c           	      b    t        |D cg c]  } | |       c}t                     S c c}w r+  )r   r   r,  s      r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>6  s)    <$)*qVB]*:J:L,
 *r-  c           	      N    t        |D cg c]  } | |       c}      S c c}w r'  )r   r,  s      r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>9  s#    M$)*qVB]*-
 *   "c           	      N    t        |D cg c]  } | |       c}      S c c}w r'  )r   r,  s      r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda><  s#    $)*qVB]*2
 *r3  c                     t        t        |j                         D ci c]  \  }} | |       | |       c}}t        t	                           S c c}}w r+  )r   r   r  r  r   r0  s       r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>D  sJ    =Q:?++-H$!QAr1-H.0>
 Hs   Ac                     t        |      S r'  )r   r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>L  s    6u= r  c                 ,    t        |t                     S r+  r   r   r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>N  s    8Q!1!39
 r  c                     t               S r'  )r   r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>Q  s
    4G4I r  c                     t        |      S r'  )r   r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>R  s    7KE7R r  c                 ,    t        |t                     S r+  )r   r   r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>T  s    ;Q!1!3<
 r  c                 ,    t        |t                     S r+  )r   r   r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>X  s    :%5%7 r  c                 ,    t        |t                     S r+  r8  r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>^      7%5%7 r  c                 ,    t        |t                     S r+  r8  r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>c  r>  r  c                 ,    t        |t                     S r+  r8  r  s     r  rM  z6SourcelessBuilder.make_type_handlers.<locals>.<lambda>h  r>  r  rN  r  r%  r	  c                     |S r'  r  r  s     r  passthroughz9SourcelessBuilder.make_type_handlers.<locals>.passthroughm  s    Lr  )"ro  r  rs   rj  r8   r  r  r  r  r  r  re  r5   r6   r  r  r  rZ  r  r  r  r   r  ru  r  r  distributionsconstraints_Real	_Interval
Constraintr  r   all_subclasses)handlersr  rB  r  r  s       @r  make_type_handlersz$SourcelessBuilder.make_type_handlers!  s    #))  	 ' 	DACHQK	D
 




 
 -5TN(()#+D> #+D> ,
''( > 	++,'
""# #J%R!!"*
%%& 	!!--BBC 	$$00667 	$$00::; 	$$00;;<	3 	A 	! 	 '55 	(C'HSM	(r  r=  )r  r  r  r  rH  r	   r  r  rj  r   r  r  r  r   r  r   r   r@  r  r  rB  r   r   r(  r'  r   rJ  r  r  r  ro  ro  p  s   
? 'CH~
tCH~

uS#X
  tCy/ 
   V.VtCyV\V  V'05c3h	   '04s
U0BS0H	   S.SsSS  S{
. {
s {
 {
 {
z 4V 4 4 4 Ph/5FGG! P Pr  ro  c                   6    e Zd ZdZd	dZedddedefd       Zy)
"SourcelessUserDefinedObjectBuilderz
    SourceLessBuilder does not return a UserDefinedObjectVariable, but in some
    cases it might be ok to return UserDefinedObjects. In such case, use this
    builder.
    r	  Nc                     t        d      )Nz/Use SourcelessUserDefinedObjectBuilder.create()r  r1  s    r  rH  z+SourcelessUserDefinedObjectBuilder.__init__  s    NOOr  rN  r  r%  c                     t        |      }t        |t              rt        |t	                     S t        |t        j                  j                        rt        |t	                     S t        |t	                     S r+  )r  r  r   r   r   r,  r  ry  rz  r   r   )rN  r%  r"  s      r  r  z)SourcelessUserDefinedObjectBuilder.create  sa    %[
j.1)%?O?QRRuxx/0%5%7  -UBRBTUUr  r=  )	r  r  r  r  rH  r  r   r   r  r  r  r  rL  rL  x  s<    P 	V* 	V3 	V? 	V 	Vr  rL  r  r'  )F(  r  r  r  r  rB  r  rq  r~  rZ  r  r{  r  r>  r  r  r  rh  r  r  collections.abcr   r   r   typingr   r   r   r	   r
   r   r  r  r   torch._dispatch.pythonr   #torch._dynamo.graph_bytecode_inputsr   r   r   torch._dynamo.utilsr   r   r   r   r   torch._guardsr   "torch._higher_order_ops.flat_applyr   !torch._higher_order_ops.torchbindr   torch._library.opaque_objectr   r   r   r   torch._opaque_baser   
torch._opsr   r   r    r7  r!   r"   r#   r$   torch._subclasses.meta_utilsr%   r&   torch._utils_internalr'   %torch.fx.experimental._backward_stater(   torch.fx.experimental._dynamismr)   torch.fx.experimental.sym_noder*   r+   r  r,   r-   r.   r/   r0   r1   r2   r3   r4   torch.fx.immutable_collectionsr5   r6    torch.nn.utils._expanded_weightsr7   torch.utils._ordered_setr8   torch.utils._python_dispatchr9   r:   torch.utils._sympy.value_rangesr;   torch.utils.weakr<   r  r>   r?   r@   rA   rB   rc  rC   r  rD   rE   rF   r  rG   rH   rI   pgorJ   rK   rL   rM   rN   rw  rO   r  rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r]   r^   r_   r`   ra   rb   rc   rd   re   rf   rg   rh   ri   rj   rk   rl   rm   rn   r  ro   rp   rq   rr   rs   rt   ru   rv   rw   rx   ry   rz   r{   r|   r}   r~   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r  r   r   r   r   r   r   r   builtinr   constantr   ctx_managerr   r   r   r   r   r   dictsr   r   r   r   rw  r   	functionsr   r   r   r   r   r   r   r   r   r   r   r  r   r   iterr   r   r  r   r   listsr   r   r   r   r   r   r   r   miscr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   	nn_moduler   r   r   	optimizerr   script_objectr   r   rS  r   setsr   r   r   r   streamsr   r   r   r  r   r   r   r   r   r   r   r   r   r   torch_functionr   r   r   r!  r   r   r   r   r   r   r   r   r   r   r   r   r   r   numpyr   ModuleNotFoundErrorr?  torch._dynamo.codegenr  torch._dynamo.symbolic_convertr  r  	getLoggerr  r  r  getArtifactLoggerrP  r  r  r  r  DimListr(  r@  r  r  	dataclassr  rE  rn  varsr  r`  r  rj  ra  ry  rz  rr  rs  ro  rp  rc  r  r  r  r  r  r4  rj  Proxyr0  r  r  r  rB  r  r  rj  r%  r  rp  rq  r~  r  r  r  r  r  r  r  r  r  ro  rJ  r	  rL  r  r  r  <module>r~     s  &              	 
    4  L L    ; 
  ) 9 <  * H H  B 1 ? A E
 
 
 J ; / 8 * T T ? S S A A  '                 B! ! ! ! ! ! ! ! !D   % &   0    ; ;	 	 	          4  
 ) ( O O $ $   J I I       
       $
 //  g!NN44'  
 m$CL
"V " "
	 	 G9 G9 G9T<H <, &/ /Y--// & IcN  +.% s3x /
 "'!>!> $)HHOO$D$D !E( E(PP O    0 (= # * !%!%	# : $;	
  2
#
,/
@C
	
z !%!%$[!$#$ $ :	$
 $;$ $ $Z "&D[!D#D 88>>D LL	D
 $;D D DX !%!% [! #  88>>  :	 
 $;     HN
N
#N
 88>>N
 #s(^	N

 $;N
 [!N
 N
bq T!Wt^ :#  $;	
 
#s(^2([!(#( 88>>( 	(
 $;( #s(^( (8
j 
 %) #c(T/ (+/ cDj /SX (
3 
4 
#+.38<<@ &* 3s8d? ),0 sTz 0c#h (
C 
D 
( ^
^#^ ^ 	^
 ^ ^P	 "&


#
 TM	

 
 $J
 	
2 "&v
v#v TM	v
 v $Jv 	vrB BJ $5#G#G#I   V VyA  B
s   $]- -^ ?^ 