g^(h"UddlZddlmZddlmZmZmZddlZddlm Z ddl m Z m Z ddl mZmZddlmZdd lmZdd lmZgZeeed <eeeeee efZeeeeee ee fZej:fd Z dd edeeej@jBfdee efdee efdeeedeee"deeedeeeffdZ#y)N)Iterable)AnyCallableOptional) QConfigAny) MatchAllNodePattern)GraphNode)type_before_parametrizations)_is_observed_standalone_module)QuantizeHandler__all__ct|tr$|^}}|turt|dk(sJdg}n|}g}t|trt |t ry||k(ryt|trt|j|kDryt|trPt |tjjr,|jdk7ryt|j|k(syt|rF|jdk7s|j|ury|jturU|j d|dk7r@yt|t"r|jdk7s|j|k7ry|j|k7ry|syt|t|j k7ryt%fd t'|j |DS) z&Matches a node in fx against a patternz.Expecting getattr pattern to have two elementsTF call_module call_functionr call_methodc3BK|]\}}t||dyw)r )max_usesN) _is_match).0node arg_matchmoduless b/var/www/skyplay_api_hub/venv/lib/python3.12/site-packages/torch/ao/quantization/fx/match_utils.py z_is_match..Js+ D) '4Q77s) isinstancetuplegetattrlentype issubclassrr userstorchnnModuleopr targetcallableargsstrallzip)rrpatternr self_match arg_matchess` rrrs'5!#* [  w<1$ V&V V$K  *d# :|(L w dD !S_x%?*d# :uxx(O 77m #+GDKK,@AZO *  77o %J)F [[G #yy|wqz) J $ 77m #t{{j'@  "  ;3tyy>) "499k: graphrpatternsroot_node_getter_mappingstandalone_module_namesstandalone_module_classescustom_module_classesreturnc |g}ggi}t}fdfdt|jD]} | j|vs| j|vs!|j D]\\} } |j | d} t || | s&| j|vs5g} | | | | || | || }| }||| | | ||J|jD]L} | jdk(st|| j|vs.| | dt| |df|| j<Ndtdtttjjfffd }|jD]]} | jdk(s|| j|st!|| js?| | dt| |d f|| j<_|S) a Matches the nodes in the input graph to quantization patterns, and outputs the information needed to quantize them in future steps. Inputs: - graph: an fx.Graph object - modules: a mapping of fully qualified module name to instance, for example, {'foo': ModuleFoo, ...} - patterns: a mapping from a tuple of nodes in reverse order to uninitialized QuantizeHandler subclass. Outputs a map of node_name -> (node, matched_values, matched_pattern, QuantizeHandler instance, qconfig) For example, { 'relu_1': (relu_1, [relu_1], torch.nn.functional.relu, , QConfig(...)), ... } Nc t|tr||||f||j<yt|tsy|D]}||||||yN)rr namer) last_node match_map node_patternmatched_node_patternr0 match_valuen#_recursive_record_node_in_match_maps rrEz:_find_matches.._recursive_record_node_in_match_map{sa lD )$ ,Il'' ( L(3 ! 3y!-A7K r3ct|tr|^}}t|dk(}g} ||||||dtur*t ||j D]\} }  | | |||t|dkDr8|r|j t|y|jt|y|j |dy|j |y)Nr r) rr r"r!r/r,appendextendlist) r0rr?rBr@sr, is_single_argcurrent_node_pattern subpatternarg record_matchs rrOz#_find_matches..record_matchs gu %HAINM/1 D)-A9 Mqz('*4';YOJ S$8LiXY'(1,!(//6J0KL(//5I0JK$++,@,CD ' ' -r3rT)is_custom_module node_targetrc8|J|vxst||vSr=)r#)rQrr8r7s ris_standalone_modulez+_find_matches..is_standalone_modules7""" 2 2 )GK()() r3)rS)setreversednodesr>itemsgetrr)r#r*rr-dictr&r'r(r)r4rr5r6r7r8r9r@ all_matchedrr0quantize_handler_clsroot_node_getterrBquantize_handlerr?rSrErOs `` @@r _find_matchesr^Ps>$ " ($&!&"$)+IEK&.4% 99I %$));*F191A --#;#?#?#N WdG4)9S79( $6JIV';,g7G($!%I7!!,,( ) 2     GG} $WT[[)*.CCgE $Idii  # S%((//=Q8R   77m # g 6-gdkk.BCgDI $Idii  r3)NNN)$syscollections.abcrtypingrrrr&torch.ao.quantization.qconfigrtorch.ao.quantization.utilsrr torch.fx.graphr r torch.nn.utils.parametrizer graph_modulerr]rrrIr-__annotations__r _MatchResult_MatchResultWithQConfigmaxsizerrYr'r(r#r^r3rrls5 $** 4=&C8-cT4:x'8/IJ $t*hw'*D03{{.l486:15T T #uxx& 'T7O+,T#7H#45 T &d3i0 T (T 3 T$DI.T #| Tr3