+ &j`aR=t90t^RIt^RIHtHt^RIHtHt^RIH t H t H t ^RI t .R>Ot]P!]4t] !R4t] !R4t]] ] P(]3,]3,t]]]3,t]],t]!RR44t]]]]],.]],3,t]!R R 44t]!4t]^k/s]^k]] ]]3,]],3,t!/s"]^k]]]] ]]3,,3,t#/s$]^k/s%]^k/s&]^kR R lt'R R/RRllt(RRlt)RRlt*RRlt+RRRRRR/RRllt,RRRRRR/RR llt-R!R"lt.R#R$lt/R%RR&R/R'R(llt0R)R/R*R+llt1R,R/R-R.llt2R/RR0RR,R/R1R2llt3R,R/R3R4llt4R5R6lt5R0R/R7R8llt6R0R/R9R:llt7R0R/R;R<llt8R#)?N)CallableIterable) dataclassfield) Concatenate ParamSpecTypeVarPRc8a]tRt^toRtRtRtV3RltRtVt R#) _OverrideNodezTrack function override data.FTcb<V^8dQh/S[;R&S[;R&S[;R&S[;R&S[;R&S[;R&#)dsl_name op_symbol dispatch_key override_fnunconditional_overrideactive)str_OpFnbool)format __classdict__s"n/Users/jameslopez/projects/CWCArchive/cwc-podcast/.venv/lib/python3.14/site-packages/torch/_native/registry.py __annotate___OverrideNode.__annotate__sQM  N   !( N) __name__ __module__ __qualname____firstlineno____doc__rr__annotate_func____static_attributes____classdictcell__rs@rr r s' $)Frr ca]tRt^,toRt]!]R7t]!]R7t]!]R7t V3RlRlt RV3RlRllt V3RlRlt V3R lR lt V3R ltR tVtR #) _FilterStatez+Manages filtering state for override nodes.)default_factoryc&<V^8dQhRS[RS[/#)rnodereturn)r r)rrs"rr_FilterState.__annotate__4s-Drc VPVP9dR#VPVP9dR#VPVP 9dR#R#)z Check if a node is enabled based on current filter state. Args: node: The override node to check Returns: bool: True if the node should be enabled, False if filtered out FT)r _dsl_namesr _op_symbolsr_dispatch_keys)selfr-s&&r check_enabled_FilterState.check_enabled4sF ==DOO + >>T-- -    3 3 3rc <V^8dQhRS[S[S[,,R,RS[S[S[,,R,RS[S[S[,,R,RS[RR/#)r dsl_namesN op_symbols dispatch_keys remove_keysr.)rrr)rrs"rrr/IsnII#&-I(3-'$.IXc]*T1 I  I  Irc V'dV;P\\V44,unV;P\\V44,unV;P\\V44,unR#V;P\\V44,unV;P\\V44,unV;P\\V44,unR#)a Update filter sets as (current | new) or (current ~ new). Args: dsl_names: DSL names to add/remove from filter op_symbols: Operation symbols to add/remove from filter dispatch_keys: Dispatch keys to add/remove from filter remove_keys: If True, remove keys from filter; if False, add them Note: Uses set.discard as it doesn't raise an exception if the element wasn't in the set to begin with. N)r1set_resolve_iterabler2r3)r4r8r9r:r;s&&&&&rupdate_FilterState.updateIs(  OOs#4Y#?@ @O   $5j$A B B    3'8'G#H H  OOs#4Y#?@ @O   $5j$A B B    3'8'G#H H rcF<V^8dQhRS[S[S[S[3,,/#rr.)r=tupler)rrs"rrr/fs!   s5c?';  rc X\VPVPVP4#)z Build a set of dictionary keys based on the current filter state. Returns: set[tuple[str, str]]: Set of (op_symbol, dispatch_key) tuples )_build_key_setr1r2r3)r4s&rbuild_disable_key_set"_FilterState.build_disable_key_setfs* OO        rc <V^8dQhRS[/#rB)r)rrs"rrr/ssrc fRpVR, pVR, p\VP4Fwr#VRV RV R2, pK VR, p\VP4Fwr$VRV RV R2, pK VR, p\VP4Fwr%VRV RV R2, pK V#) z-Return string representation of filter state.zFilter State: z === DSL: ===  z:  z === OP SYMBOL: === z === DISPATCH KEYS: === ) enumerater1r2r3)r4sidslopkeys& r__str___FilterState.__str__ss   0FA 4s"SE$ $A1 %%t//0EA 4s"RD# #A1 )) 3 34FA 4s"SE$ $A5rcn<V^8dQh/S[S[,;R&S[S[,;R&S[S[,;R&#)rr1r2r3)r=r)rrs"rrr/,s8C5  S6  H9 rrN)F)r r!r"r#r$rr=r1r2r3r5r?rFrSr%r&r'r(s@rr*r*,sX5 5J!#6K$S9N*II:   Orr*c .V^8dQhR\\\,,R,R\\\,,R,R\\\,,R,R\\\\3,,/#)rr8Nr9r:r.)rrr=rC)rs"rrrsg!!Xc]"T)!hsm#d*!#&-! sCx !rc|a\4oRV3RllpV!V\4V!V\4V!V\4S#)a+ Build a set of dictionary keys based on filter criteria. Args: dsl_names: DSL names to include in key set op_symbols: Operation symbols to include in key set dispatch_keys: Dispatch keys to include in key set Returns: set[tuple[str, str]]: Set of (op_symbol, dispatch_key) tuples cnV^8dQhR\\\,,R,R\RR/#)rentriesNgraph_lib_dictr.)rr _MappingType)rs"rr$_build_key_set..__annotate__s2 % %x}$t+ %=I %  %rc<\V4pVF+pW19gK W,FpSPV4K K- R#)z;Helper to add matching keys from graph_lib_dict to key_set.N)r>add)rYrZresolved_entriesentryrRkey_sets&& r_append_to_set&_build_key_set.._append_to_sets<-W5%E&)00CKK$1&r)r=_dsl_name_to_lib_graph_op_symbol_to_lib_graph_dispatch_key_to_lib_graph)r8r9r:rbras&&& @rrErEs> %(EG % %945:67="<= Nrprint_inactiveFc(V^8dQhR\RR/#)rrgr.N)r)rs"rrrsdtrcb\P4Fwwrp\RV: RV: 24\V4FlwrEVP'g V'gK!RV RVP : RVP : 2pV'dVRVP: 2, p\V4Kn K R#)zn Print all override graphs for debugging purposes. Args: print_inactive: Whether to print inactive nodes zop=z, key=rKz: node.dsl_name=z, node.unconditional_override=z node.active=N)_graphsitemsprintrMrrr)rgrQrR node_listrOr-rNs$ r_print_override_graphsrns!(  9 Wv +GA{{{nns"3T]]$44St7R7R6TU!>T[[N++Aa ,!0rcdV^8dQhR\R\R\PP/#)rrrr.)rtorchlibraryLibrary)rs"rrrs+cAVAVrcW3pV\9d)\PPRRV4\V&\V,#)z Get or create a torch.library.Library instance for the given key. Args: op_symbol: The operation symbol dispatch_key: The dispatch key Returns: torch.library.Library: The library instance atenIMPL)_libsrprqrr)rrrRs&& r_get_or_create_libraryrws:  #C %]]**66<Hc :rchV^8dQhR\PPR\R\RR/#)rlibr-rr.N)rprqrrr r)rs"rrrs6   &3CF rcxVPVPVPVVP'*RR7R#)z Register a single node implementation with the library. Args: lib: The torch.library.Library instance node: The override node to register dispatch_key: The dispatch key for registration T) with_keysetallow_overrideN)implrrr)ryr-rs&&&r_register_node_implr~s7HH  333  rcV^8dQhR\\\,,R,R\\,/#)riterableNr.)rr)rs"rrrs,hsm 3d :x}rcpVf.#\V\4'd\V\4'dV3#V#)z Resolve various input types to a consistent iterable of strings. Args: iterable: String, iterable of strings, or None Returns: Iterable[str]: Consistent iterable output ) isinstancerr)rs&rr>r>s5 h ) )Z#-F-F{ Orenable_dsl_namesenable_op_symbolsenable_dispatch_keyscV^8dQhR\\\,,R,R\\\,,R,R\\\,,R,RR/#)rrNrrr.rlist)rs"rrrsY22DIo,2T#Y-2S /D0 2  2rc\PRVVV4\PVVVRR7\ VVV4pVF1pVwrVV\ 9gK\ WV\ V,\R7K3 R#)aO Re-enable overrides by removing them from filter state and reregistering. Args: enable_dsl_names: DSL names to re-enable enable_op_symbols: Operation symbols to re-enable enable_dispatch_keys: Dispatch keys to re-enable Note: This function uses reverse filter state management (removing from filters to enable). z>Re-registering ops by dsl: %s, op_symbol: %s, dispatch_key: %sT)r; filter_stateN)loginfo _filter_stater?rErj_register_overrides_from_graph)rrrrarRrrs$$$ rreenable_op_overridesrs$HHH  %3%G"% '> +M  rdisable_dsl_namesdisable_op_symbolsdisable_dispatch_keyscV^8dQhR\\\,,R,R\\\,,R,R\\\,,R,RR/#)rrNrrr.r)rs"rrr=sY))T#Y-)d3i$.)c?T1 )  )rc\PRVVV4\PWV4\P 4pVF2pVwrVV\ 9gK\ VV\ V,\R7K4 R#)aC De-register overrides by updating filter state and reregistering graphs. Args: disable_dsl_names: DSL names to disable disable_op_symbols: Operation symbols to disable disable_dispatch_keys: Dispatch keys to disable Note: This function uses filter state management to selectively disable operations. z>De-registering ops by dsl: %s, op_symbol: %s, dispatch_key: %srN)rrrr?rFrj_cleanup_and_reregister_graph)rrrrarRrrs$$$ rderegister_op_overridesr=ss$HHH *@UV%2$G$G$IG"% '> ) *   rcFV^8dQhR\R\\,/#)rrr.r)rs"rrriscrc\4p\P4F6wwr#pVF(pVPV8XgKVP V4K4 K8 \ V4#)zGet list of operations registered by a specific DSL. Args: dsl_name: Name of the DSL to query. Returns: Sorted list of operation names registered by the DSL. )r=rjrkrr^sorted)r operationsr_nodesr-s& rget_dsl_operationsrisTJ!(D}}(y)"1 * rc nV^8dQhR\R\R\R\\\3,RR/#)rrrrrRr.N)rrC)rs"rrr{sI#L#L#L#L#L sCx #L  #Lrc`RRlpV!\W4V!\W4V!\W#4R#)z Update the registration mapping dictionaries. Args: dsl_name: The DSL name op_symbol: The operation symbol dispatch_key: The dispatch key key: The dictionary key tuple cV^8dQhR\\\\\\3,,3,R\R\\\3,RR/#)r registrationsymbolrRr.N)dictrrrC)rs"rr/_update_registration_maps..__annotate__sO # #3U38_ 556 # #38_ #  #rcdVPV4pVf V.pW0V&R#VPV4R#)z;Helper to add key to registration list or create new entry.N)getappend)rrrR entry_lists&&& r_get_new_entry_or_append;_update_registration_maps.._get_new_entry_or_appends7 "%%f-  J#-   c "rN)rdrerf)rrrrRrs&&&& r_update_registration_mapsr{s*& #3XC4iE7Krallow_multiple_overridercV^8dQhR\R\R\R\R\\,R\R\RR /#) rbackend lib_symbolrrr}rrr.N)r _OpOverrideFn _OpReplaceFnr)rs"rrrsg/I/I /I/I/I /I , & /I"/I!/I /Irc VR8wd\RV R24hW#3p\PV.4pVP\ VVVVVR74V\V&\ WW7R7R#)a Register a passed override function to the dispatcher. Actually a graph-building operation; real registration happens later. Args: backend: The backend name (DSL name) lib_symbol: Library you're overriding symbols in (must be "aten") op_symbol: Name of the operation you're overriding dispatch_key: Dispatch key to override impl: Implementation function for the override allow_multiple_override: Allow overriding an existing override unconditional_override: Implementation doesn't have a fallback and doesn't require torch.DispatchKeySet as the first argument Raises: ValueError: If lib_symbol is not "aten" rtz.Unsupported lib_symbol (must be "aten", got: "")rrrrr)rRN) ValueErrorrjrrr r) rrrrr}rrrRop_graphs &&&&&$$ rregister_op_overriderst8VI*UVWXX  #C{{3#H OO%#9  GCLg,Hrforce_reregisterctV^8dQhR\\,R\\,R\R\/#)roriginal_graph new_graphrr.)rr r)rs"rrrs:''''M"' '  'rcV'dR#W8g#)a' Determine if a graph needs reregistration based on changes. Args: original_graph: The original graph before modification new_graph: The graph after modification force_reregister: If True, always reregister regardless of changes Returns: bool: True if reregistration is needed Tr)rrrs&&$r_should_reregister_graphrs"  &&rrc pV^8dQhR\R\R\\,R\R,RR/#rrrgraphrNr.rrr r*)rs"rrrsD           %    rcbW3pV\9d\VV'd\VVVVR7R#R#)a= Clean up existing library and reregister a graph. This is the common pattern used across reorder, deregister, and reenable operations. Args: op_symbol: The operation symbol dispatch_key: The dispatch key graph: The graph to register filter_state: Optional filter state for conditional registration rN)rvr)rrrrrRs&&&$ rrrs:$  #C e| #J &   %   rkeys_to_processreregister_overridesc V^8dQhR\R\\\\3,,R,R\R\ R,RR/#)rtransformation_fnrNrrr.)UserOrderingFnr=rCrrr*)rs"rrrsZNN%NsCx)D0N N % N  Nrc VeTM\\P44p\V4FwrVWV3\9dK\\WV3,4pV!WVV4p\T\4'g.\PR\T4PTT4KwT\YV3&T'gK\YxRR7'gK\!TTTTR 7K R# \\ \ \3d\PRTTRR7K\d\PRTT4EKi;i) ad Apply a transformation function to graphs and optionally reregister. This is the core pattern used by reorder_graphs_from_user_function and can be reused for other graph transformation operations. Args: transformation_fn: Function to transform each graph keys_to_process: Keys to process, or None for all graphs reregister_overrides: Whether to reregister changed graphs filter_state: Optional filter state for conditional registration Note: If transformation_fn raises an exception for a specific graph, that graph will be skipped and processing will continue with remaining graphs. NAGraph transformation failed for %s/%s. Preserving original graph.Texc_infoNUnexpected error in graph transformation for %s/%s. Preserving original graph.zbGraph transformation returned invalid type %s for %s/%s. Expected list. Preserving original graph.F)rr)r=rjkeysr TypeErrorrAttributeError RuntimeErrorrwarning Exception exceptionrtyper rr) rrrr target_keysrrrrs &$$$ r_apply_graph_transformationrs76+6C RI$)T** KKtY((   .7)* $< % %  *)  U$5:~|D  KKS     MM`    s C$$4EE$EEc pV^8dQhR\R\R\\,R\R,RR/#rr)rs"rrr_sD       %    rcW3p\V!pVFBpRpV'dVPV4pV'd\WVV4RVnK;RVnKD R#)z Register all overrides in a single graph. Args: op_symbol: The operation symbol dispatch_key: The dispatch key graph: List of override nodes to register filter_state: Optional filter state for conditional registration TFN)rwr5r~r)rrrrrRryr-enables&&&$ rrr_sS  #C # &C !//5F  < 8DKDKrcV^8dQhRR/#)rr.Nr)rs"rrr~s     rcb\P4FwrVwr#\VVV4K R#)zG Perform all registration calls from previously-built override graphs. N)rjrkr)rRrrrs r_register_all_overridesr~s/mmo "% &    &rc4V^8dQhR\R\RR/#)rfnrr.N)rr)rs"rrrs( rc \VVR7R#)a Reorder override graphs using a user-provided ordering function. Args: fn: User-provided function that takes (op_symbol, dispatch_key, graph) and returns a reordered graph reregister_overrides: Whether to reregister graphs that have changed Note: This function uses the common graph transformation pattern and can serve as an example for other graph manipulation operations. rrNr)rrs&$r!reorder_graphs_from_user_functionrs" 1rclV^8dQhR\\\\.\3,R\RR/#)r filter_fnrr.N)rrr r)rs"rrrs955c=147855 5rc4aRV3Rllp\VVR7R#)a Apply a filter function to remove nodes from graphs. This is a convenience function that uses the graph transformation pattern to filter out unwanted nodes. Args: filter_fn: Function that takes (op_symbol, dispatch_key, node) and returns True to keep the node, False to remove it reregister_overrides: Whether to reregister modified graphs Example: # Remove all nodes with "deprecated" in the DSL name _apply_graph_filter( lambda op, dk, node: "deprecated" not in node.dsl_name, reregister_overrides=True ) Note: If filter_fn raises an exception for a specific graph, the original graph will be preserved and processing will continue. c tV^8dQhR\R\R\\,R\\,/#rrrrr.rrr )rs"rr)_apply_graph_filter..__annotate__s6&)26}2E m rc<VUu.uFpS!WV4'gKVNK up#uupi \\\\3d\P RTTRR7Tu#\ d\PRTT4Tu#i;i)z-Apply filter_fn to graph with error handling.rTrrrrrrrrrr)rrrr-rs&&& rfiltering_transformation5_apply_graph_filter..filtering_transformations %*WUTi QU.VDDUW WW:~|D  KKS   L  MM`  L  s++&&++6B #B ,B B rNr)rrrsf$ r_apply_graph_filterrs8, 21rcnV^8dQhR\\\.\3,R\R\RR/#)r condition_fn ordering_fnrr.N)rrrr)rs"rrrsDLLC:t+,LL L  Lrc8aaRVV3Rllp\VVR7R#)a Apply reordering only to graphs that match a condition. This allows for more targeted reordering operations. Args: condition_fn: Function that takes (op_symbol, dispatch_key) and returns True if the graph should be reordered ordering_fn: Ordering function to apply to matching graphs reregister_overrides: Whether to reregister modified graphs Example: # Only reorder CUDA operations _apply_selective_reordering( condition_fn=lambda op, dk: dk == "CUDA", ordering_fn=lambda op, dk, g: sorted(g, key=lambda n: n.dsl_name), reregister_overrides=True ) Note: If condition_fn or ordering_fn raises an exception for a specific graph, the original graph will be preserved and processing will continue. c tV^8dQhR\R\R\\,R\\,/#rr)rs"rr1_apply_selective_reordering..__annotate__s6)))&))26}2E) m )rc<S!W4pT'd S!YT4#T# \\\\3d\P RTTRR7Tu#\ d\PRTT4Tu#i;i \\\\3d\P RTTRR7Tu#\ d\PRTT4Tu#i;i)z=Apply ordering_fn conditionally based on condition_fn result.rTrrr)rrrshould_reorderrrs&&& rconditional_transformation?_apply_selective_reordering..conditional_transformations ))BN"  "9EBB" G:~|D  KKS   L  MM`  L  z><H  W !     d    s8B6A=A= A=<A=6C8CCCrNr)rrrrsff$ r_apply_selective_reorderingrs<))V 41rctV^8dQh/^\9d \;R&^\9d@\\\\3,\ P P3,;R&^\9d \;R&^\9d \;R&^\9d \;R&^\9d \;R&#)rrrvrjrdrfre) __conditional_annotations__r*rrCrrprqrr _GraphsTyper[)rs"rrrsH-,|,IN98tE#s(OU]]22 238OVW`*) )ab.-L-cd+**er)rrrrrr):rloggingcollections.abcrr dataclassesrrtypingrrr torch.libraryrp__all__ getLoggerr rr r DispatchKeySetrrrr rrrr*rrvrrCrrjr[rdrfrerErnrwr~r>rrrrrrrrrrrrrr)rs@rrsg.(22 ! cN CLU%9%91%<=q@A 1~  $  3T-%894 ;NNO TT Tp+n ,7985c?D$778 CeCHo../ (*)+--(**!HU&((&2/3215248 2j)04)26)59 )X$#LL/I%*/I$)/Id'# '0  )-  FN48N"' N )- Nb )-  >  "'.5"'5pL"' LLr