Vh <VUddlZddlmZmZddlmZmZmZddlZddl m Z ddl m Z ddl mZmZddlmZddlmZdd lmZmZmZdd lmZmZmZmZdd lmZmZdd l m!Z!dd l"m#Z#gZ$e%e&e'd<dede&dede(dedef dZ)de!defdZ*dejVjXde-e&ejVjXfde#dede-e&e.e&e/ffde-e&eff dZ0dede1e&de&ddfd Z2d!ed"efd#Z3d$e d%e%efd&Z4ded'eee&fdedefd(Z5d)Z6d*Z7d+Z8dede-eee&feffd,Z9ded-efd.Z:y)/N) defaultdict OrderedDict)AnyCallableUnion) _FusedModule)QConfig) BackendConfig DTypeConfig)get_module_to_qat_module)_is_activation_post_process)_add_module_to_qconfig_obs_ctrqconfig_equals QConfigAny)_MODULE_NAME_DICT_KEY_MODULE_NAME_REGEX_DICT_KEY_OBJECT_TYPE_DICT_KEYQConfigMapping) _parent_nameget_qconfig_dtypes) GraphModule)Graph__all__qconfig_mappingcur_module_pathcur_object_typecur_object_type_idxfallback_qconfigreturnc||jjD]\\}}}}||k(s||k(s||k(s|cS|SN)&module_name_object_type_order_qconfigsitems) rrrrr module_name object_typeindexqconfigs ^/home/dcms/DCMS/lib/python3.12/site-packages/torch/ao/quantization/fx/qconfig_mapping_utils.py7_maybe_adjust_qconfig_for_module_name_object_type_orderr)s\"HHNNP     O +/--N  modelc |j}t|dk(r|St|j}|jj D]}|j dk(s|j|vs"|t|j}t|tsKt|jj}|jt|dd}|ddD]Q}t!|jt|d|r)t#ddt|dzdt|z|||t|<y) z Update the QConfigMapping to account for fused modules such as LinearReLU. This assumes the QConfigMapping's attributes have already been converted to OrderedDicts. r call_moduleNz+During fusion, we need to specify the same z!qconfigs for all module types in  zoffending type: )object_type_qconfigslendict named_modulesgraphnodesoptargetstr isinstancerlist_modulesvaluesgettyper LookupError) r+robject_type_dictmodulesnodemaybe_fused_moduleops fused_qconfigr6s r(_update_qconfig_for_fusionrF4sL ';; !5&&()G !!K 77m # w(>!(T[[)9!: 0,?)2299;V=WWXYZ,T"XJ78  (=J &8!9:7Kr*rootrA input_graphnode_name_to_scopecL|j}i}td}|jD]y}d} |jdk(rZt |j \} } t |t|| | |} t| |j|j d} n|jdk(rt||j |} ||j\}}t |||| } |||j }|||j xxdz cc<t|||j || } t| |j|j d} nI|jdk(r_||j\}}t ||j ||} t |||| } t| |j|j d} n|jdk(rt||j rt |t||j |j |} ||j\}}t |\}} |||}|||xxdz cc<t||||| } t| |j|j d} | ||j _nd} | ||j<||S)Nc ttSr!)rintr*r(z0_generate_node_name_to_qconfig..ns C r*get_attr call_functionr. call_methodr-)global_qconfigrr5r6rr7-_maybe_adjust_qconfig_for_module_type_or_namer>rr=_get_object_type_qconfignamer)r r')rGrArHrrIrRnode_name_to_qconfig#submodule_to_object_type_to_cur_idxrBr'r$_qconfig_with_device_checkfunction_qconfig module_path module_typer parent_names r(_generate_node_name_to_qconfigr^]s%33NKV K'!!SD 77j )$++6NKCgk&:!;[.G)GT[[$7) %WW ' 8n  (:$))'D $KCk;KG#Fk"R #  0 G Dk7G )GT[[$7) %WW %*74;;+?@Cgdkk&:!;T[[.G(:$))'D $K*+6NK"Ek"R#  0 <[ IQ N IMk;NPWG)GT[[$7) %,EGDKK ((, %*CTYY'gSDh r* config_dict allowed_keys dict_namec|jD],}||vstd|zdzt|zdz|zdzy)zChecks if the given config_dict has the correct keys Args: `config_dict`: dictionary whose keys we want to check z Expected z to have the following keys: z . But found 'z ' instead.N)keys ValueErrorr8)r_r`raks r(_check_is_valid_config_dictrfsq     L 12l#$" "      r*prepare_qconfig_mappingconvert_qconfig_mappingc t|j|jsJd|j|j|jg}|j|j|jg}t t tg}tt|D]u}||jD]]}|||vsJd||d|d|||$t||||||r=Jd||d|d|||d|||wy) aCompare the qconfig_mapping passed in convert to the one from prepare and check the values Args: `prepare_qconfig_mapping`: configuration for prepare quantization step `convert_qconfig_mapping`: configuration for convert quantization step zWExpected global qconfigs to be the same in the prepare and convert quantization configsz Missing key r/zI in convert QConfigMapping when it was present in prepareNzLExpected convert QConfigMapping to have the same qconfig as prepare for key z; prepare: z ; convert: ) rrRr0module_name_qconfigsmodule_name_regex_qconfigsrrrranger1rc)rgrh prepare_dicts convert_dicts dict_namesirUs r()_compare_prepare_convert_qconfig_mappingsrqs ..0G0V0V a`a  4444::(M 4444::(M #J 3}% & V!!$))+ VD a(( 0jm_AdV4/0 0(!#D)1^a & a(8(>6 V]^hij^k]llmnrmst'*401]1=Md=S? 7: ;@r*);r collectionsrrtypingrrrrvtorch.ao.nn.intrinsicrtorch.ao.quantizationr $torch.ao.quantization.backend_configr r *torch.ao.quantization.backend_config.utilsr torch.ao.quantization.observerr torch.ao.quantization.qconfigrrr%torch.ao.quantization.qconfig_mappingrrrrtorch.ao.quantization.utilsrrtorch.fxrtorch.fx.graphrrr:r8__annotations__rLr)rFnnModuler2tupler>r^setrfrqrrTrrrSrrrMr*r(rs 0'' .)KOF  I c#  !   *&Kk&KN&KRg ((//g #uxx& 'g g $ g S%T "223 g  #z/ g T$'H9< ,$V+$VFT$VN* *%)+%6*ZS#Sx}%S!S SS$!#! %#  *+!H @# @5B @r*