8h<&SSKrSSKJr SSKJr SSKrSSKJ r SSK J r SSK J r SSKJr SSKJrJrJr SS KJr \ "S S 9S \S \4Sj5r\ "S S 9S\S \4Sj5r\ "S S 9S S S.S\ S \S\S\\\S4S\S \\ \\S4\\S444 Sjjj5r\ "S S 9S\ S\ S\\S4S\\S4S \ 4 Sj5r\ "S S 9S\ S \S S4Sj5r\ "S S 9S!S\ S\\\S4S\S\S \ 4 Sjj5r g)"N) SimpleQueue)Optional) compatibility)Graph) GraphModule)Node)legalize_graphNodeListNodeSet)lift_subgraph_as_moduleF)is_backward_compatiblenodesreturnc8[RUS5n[5nUHDnURHnXA;dM X==S- ss'M XS:XdM3UR U5 MF /nUR 5(dyUR 5nURU5 URH1nXA;dM X==S-ss'XS:XdM UR U5 M3 UR 5(dMy[U5[U5:XdS5eU$)Nrz@topological sorted nodes doesn't have same length as input nodes) dictfromkeysrall_input_nodesputemptygetappenduserslen)r indegree_map candidatesnoden sorted_nodess [/var/www/fran/franai/venv/lib/python3.13/site-packages/torch/fx/passes/utils/fuser_utils.py topo_sortr!s==*L$/MJ%%A "a'"&   " NN4  L  ~~D!A 1$?a'NN1%     u:\* *J *  partitionc^[U5m/nTH/nURHnUT;dM URU5 M M1 S[S[4U4SjjnU"U5(agg)N root_nodesrc>[5nUnU(a\UR5nURU5 UT;agURHnXA;aM UR U5 M U(aM\g)NTF)setpopaddrr)r%visitedqueuecurrent user_node partition_sets r bfs_find_cycle*validate_partition..bfs_find_cycle>si5 %iikG KK -'$]] ' Y'+er"FT)r'rrr bool)r#outputsrr-r/r.s @r validate_partitionr3-se  NMGI -y)$84g r"always_return_tuplegm module_namepartition_lookup_tabler5.c^^^^UHnURRULdUSUR535eUR(a US35eXPRR;aMlUSUR535e [ U5(dS5eTc[ RU5m[5m0m0mS[S[4UUUU4SjjnUHnTRXV5nUTU'M 0nUH%nURHn U T;dM TUX'M M' [UR55n U(aTRU 5 O%TR[U 5S:XaU S OU 5 TR!5 [#UTS US 9up[TR%55n [UR%55nXU4$) a Fuse nodes in graph_module into a GraphModule. Args: gm (GraphModule): target graph_module nodes (List[Node]): list of nodes in `gm` to fuse, where the node must be topologically sorted module_name: class name for the fused GraphModule partition_lookup_table (Optional[Dict[Node, None]]): optional dict of nodes to speed up lookup always_return_tuple (bool): whether to always return a tuple, even if there is only one output Returns: fused_gm (GraphModule): fused graph module, where its node is a copy of `nodes` in `gm` original_inputs (Tuple[Node, ...]): input nodes to `nodes` in original `gm` original_outputs (Tuple[Node, ...]): consumer nodes of `nodes` in original `gm` z* doesn't belong to passed in graph module z# has been removed from owning graphz is not found in graph module z*Invalid partition, found dependency cyclesxrc>URS:XaUT;aTU$UT;aNTRURURS9n[R"UR 5UlUTU'TU$)Nget_attr) type_expr)op placeholdernametypecopymeta)r:placeholder_nodenode_mapnode_to_placeholderr8subgraphs r remap_inputs)fuse_as_graphmodule..remap_inputss| 44:   & &A;  ' ''33AFFaff3M $(IIaff$5  !%5  ""1%%r"rr) comp_name class_name)graph owning_module _get_name_erased_find_nodes_lookup_tabler3rrrr node_copyrtuplevaluesoutputrlintr keys)r6rr7r8r5rrHnew_nodeoutput_mappingr-outsfused_gm_original_inputsoriginal_outputsrErFrGs ` @@@r fuse_as_graphmoduler_^sDzz''2- f>r||~>N O -<<MD6)L!MMxx888 f22<<>2B C 8  e $ $R&RR $%!%u!5wH "$H&&&&*%%d9! (*NI 66'/~$$ &&( )D 3t9>Qt< MMO) H{KH )..A.F.F.H(IO*/~/B/B/D)E &6 66r"sub_gm orig_inputs orig_outputsc&URRnURXA5 URR XBSS9nURR 5n[ U5S:Xa7[URS[5(dUSRUSS9 U$[U5HAupx[RRU5URn URU SS9 MC [SU55UR S'U$)N)argskwargsrrT)propagate_metac3X# UH oRRSS5v M" g7f)valN)rCr).0 orig_outputs r insert_subgm..s'( AM+    - -s(*rh) __class____name__ add_submodulerM call_module output_noder isinstancerdrSreplace_all_uses_with enumeratetorchfxProxyrrC) r6r`rarbsubmodule_name module_noderqirj proxy_outs r insert_subgmr|s%%..N^,((&&~PT&UK,,**,K <Aj1A1A!1De&L&LQ--k$-O I( 5NA{3A6;;I  - -i - M6 #(( AM( #  Ir"c^[U5HnURRU5 M gN)reversedrM erase_node)r6rrs r erase_nodesrs$ D! r" partitionsprefixc [U5HNupE[[U55nU[U5-n[ UUUUUS9upn [ XX5 [ X5 MP [U5 U$)Nr4)rtr!liststrr_r|rr ) r6rrr5 partition_idr#rrxr`rarbs r fuse_by_partitionsrst$-Z#8  i1 #l"33,?     3 - )\ R;B%$9"2 Ir"r~)fused_F)!rBr+rtypingr _Optionaltorch.fxrutorch.fx._compatibilityrtorch.fx.graphrtorch.fx.graph_moduler torch.fx.nodertorch.fx.passes.tools_commonr r r torch.fx.passes.utilsr r!r1r3rrrSr_r|rrrr"r rs (1 -JJ9e,X(-<e,-(-t---`e, ;? r7 !& r7r7 r7r7&d4:&67 r7  r7 ;dCi(%c *: :;r7-r7je, tSy!c "   -:e,"K""T"-" e, % T$*%&    -r"