L iHddlmZddlZddlZddlZddlZddlZddlZddlZddl m Z m Z ddl m Z mZddlmZddlmZmZmZmZmZmZmZddlmZmZmZddlmZerdd lm Z dd l!m"Z"m#Z#m$Z$ee%jMjNZ(e(jNZ)Gd d eZ*d Z+d$dZ,edZ-edZ.d%dZ/d&dZ0ejbd'dZ2eredd(dZejfd)dZ4d*dZ5GddZ6 d+ d,dZ7 d- d.dZ8 d/ d0dZ9 d1dZ: d2dZ; d3dZ<Gd d!Z=Gd"d#ee-Z>y)4) annotationsN)fields is_dataclass)autoEnum)Path)AnyCallableGenericLiteralNoReturn TYPE_CHECKINGTypeVar) assert_never deprecatedSelf) CodeTemplate) Namespace)IterableIteratorSequencec`eZdZeZeZeZeZeZeZ y)TargetN) __name__ __module__ __qualname__r DEFINITION DECLARATION REGISTRATIONANONYMOUS_DEFINITIONNAMESPACED_DEFINITIONNAMESPACED_DECLARATIONT/mnt/ssd/data/python-lab/Trading/venv/lib/python3.12/site-packages/torchgen/utils.pyrr$s/J&K6L6 F!Vr$rz(^|\W){}($|\W)ctjd|}|td||j\}}}||j dfS)Nz(\w+)(\.\w+)?\((.*)\)zUnsupported function schema: , )rematch RuntimeErrorgroupssplit)schemamname_paramss r%split_name_paramsr27sP )62Ay:6(CDDhhjOD!V d# ##r$TSc#8K|D]}||}||ywNr#)funcxsxrs r%mapMayber;Gs(  G =Gsc#>K|D]}||Ed{y7wr6r#)r7r8r9s r% concatMapr=Os% 7s c#K dy#t$r]}|}tj|d}|jr|jdd|n|}|f|jddz|_d}~wwxYww)Nz r ) Exceptiontextwrapindentargs)msg_fnemsgs r%contextrHWso hooc4('(vv 2cU#3!&&*$ s%A2 A2 A/AA**A//A2z*Use typing_extensions.assert_never insteadcDtdt|j)NzUnhandled type: )AssertionErrortyper)r9s r%rrhs /Q0@0@/ABCCr$c,tj|Sr6)r from_file) template_fns r%_read_templaterOms  ! !+ ..r$ctj|jddj}tj |dS)Nlatin1F)usedforsecuritylittle) byteorder)hashlibsha1encodedigestint from_bytes)srVs r%string_stable_hashr\ss7 <<*E B I I KD >>$(> 33r$ceZdZ d dZeddZddZ ddZ ddZ ddZ dd dd Z dd dd Z dd Z dd Z y) FileManagercpt||_t||_t|_||_yr6)r install_dir template_dirsetfilesdry_run)selfr`rards r%__init__zFileManager.__init__|s.  , . #  r$cnt|jDchc]}|jc}Scc}wr6) frozensetrcas_posix)refiles r% filenameszFileManager.filenamess%djjAd$--/ABBAs2ct|}d} |jd}||k7r1|jj dd|j |dyy#t$rYBwxYw)Nzutf-8)encodingT)parentsexist_ok)r read_textOSErrorparentmkdir write_text)refilenamecontentsrj old_contentss r%_write_if_changedzFileManager._write_if_changedsoH~#'  >>7>;L | # KK  dT  : OOHwO 7 $   sA A$#A$ct|jr Jd||j|z }|}t|trd|vrwt dz } tt jdjxs|j} |jtj}i|ddd|d|zi}t#|}|j%|} t'j(dd | t&j* St|t,r|St/|y#ttf$r|j}YwxYw#t$r|j }YwxYw) Nztemplate_fn must be relative: generated_commentzgen.py__main__@z generated by z from z (""")\n+ # match triple quotes ( (\s*@.+\n)* # match decorators if any \s*(class|def) # match class/function definition ) z\g<1>\n\n\g<2>)flags)r is_absolutera isinstancedict TORCHGEN_ROOTsysmodules__file__absoluteKeyErrorAttributeError relative_to REPO_ROOTri ValueErrorr/rO substituter(subVERBOSEstrr) rerN env_callable template_pathenvgenerator_default generatorgenerator_pathtemplatesubstitute_outs r%substitute_with_templatez$FileManager.substitute_with_templatesx  $002 ,[M : 2))K7 n c4 "#-$1H$<!= $ J/88M'')S-H +SS)(*4::%E)>tf'EE% t||!::');N  " "D> " J r$c*|j|||yr6)r)rerurs r%writezFileManager.writes   8\Br$N)base_envc 6|j||||||||y)N)key_fnr num_shardsr sharded_keys)write_sharded_with_template)reruitemsrrrrrs r% write_shardedzFileManager.write_shardeds/ ((   %!% ) r$ct|} | jr Jd|ddi} t|D cgc] } dd| i } } | g| z} || D]j|D]C}| D]<|vr1t |t sJd|j |<8g|<>Ed fd }|jrd}|D]5}||}t||z}||}|| |||| |7| D]:d}|j| j| j||fd<|jj|j| j| jdz ycc} w) Nrshard_id Everythingr0z'sharded keys in base_env must be a listcl|jD] \}}|vs Jd|||xx|z cc<"y)Nzundeclared sharded key )r)intofrom_kvrs r% merge_envz:FileManager.write_sharded_with_template..merge_envsD  1L(G,CA3*GG(Q1  r$r#cSr6r#)shardsr%z9FileManager.write_sharded_with_template..5sr$)rdict[str, list[str]]rrreturnNone)rr~rangeupdaterlistcopyrdr\r with_stemstemrcdiscardr`)rerurNrrrrrrrj everythingishards all_shardskeyritemsidrrrs ` @r%rz'FileManager.write_sharded_with_templatesH~##%O)DXJ'OO%&0,%? +0+<( &'Z1QC !( ( !\F*  # ' X& '  $C# $%<%eCj$7A7"'s!2E#J!#E#J $ $  <<E 'D,C$S)J6Ct$C fSk3 ' j# &  '  EZ(H  $ $$))XJ78   4++dnn {*=U.VVW[( sFcdjd|gdt|jDd}|j||y)zSWrite a file containing the list of all outputs which are generated by this script.r?zset(c3DK|]}d|jdyw)z ""N)ri).0rjs r% z,FileManager.write_outputs..Bs LE$--/*!,Ls )N)joinsortedrcrx)re variable_namerucontents r% write_outputszFileManager.write_outputs;sP)) M 9KL       x1r$ctjj|jtjj t S)z This needs to be deterministic. The template dir is an absolute path that varies across builds. So, just use the path relative to this file, which will point to the codegen source but will be stable. )ospathrelpathradirnamerres r%template_dir_for_commentsz%FileManager.template_dir_for_commentsHs, wwt00"''//(2KLLr$)r` str | Pathrarrdboolrr)rzfrozenset[str])rurrvrrr)rNrr"Callable[[], str | dict[str, Any]]rr)rurrNrrrrr)rurrrrr)rurr Iterable[T]rCallable[[T], str]r#Callable[[T], dict[str, list[str]]]rrYrdict[str, Any] | Nonerset[str]rr)rurrNrrrrrrrrrYrrrrrr)rrrurrrrr)rrrrfpropertyrkrxrrrrrrrr#r$r%r^r^{s  !    CC 84494 4lKK K9 K  K$CC9C  C+/   #  :   (    >+/<X<X <X <X # <X:<X<X(<X<X <X| 2Mr$r^ctjj|jd}|r|n |j}t |||j S)N templates)r`rard)rrr source_pathr`r^rd)optionsr`ras r%make_file_managerrRsG77<< 3 3[AL!,+'2E2EK ! r$cdtjdk\rddlm}||||St |||S)N) r)pformat)rCwidth)r version_infopprintr_pformat)objrCrrs r%dataclass_reprr`s3  7""sFE** Ce 44r$ct|sJdt||jj}|t |dzz }t |Dcgc]1}|j s|jt||jf3}}g}|D]\}} |t |zdz} t| rt| ||| } nTt| trt| ||| } n5t| tttfrt!| ||| } n t | } |j#|d| |dz} d| j%|} |d| dScc}w)Nz%obj should be a dataclass, received: r@= z, (r)rrK __class__rlenrreprr/getattrrrr _format_dictrrbtuple _format_listappendr)rrCr curr_indent class_namef fields_list fields_strr/attr _curr_indentstr_repr indent_strbodys r%rrnsT  Q Ed3i[QQ ''J3z?Q&&K;A#;Qa!&&AFFGC01QKQJ! 0 d#SY.2  fe\BH d ##D&%FH tS%0 1#D&%FHDzHTF!H:./ 0s"J   " ": .D\4& "")Rs E %$E c ||dzz }g}|jD]V\}}t|}t|rt||||t |zn t|}|j |d|Xt ||||ddS)Nrz: {})rrrrrr_format) rrCrr dict_reprrrk_reprv_strs r%rrs 6A:KI /1aA Q{S['@ Aa  F82eW-./ 9fe[#s CCr$c ||dzz }|Dcgc]&}t|rt||||n t|(}}t|trdnd\}}t ||||||Scc}w)Nr@)[])rr)rrrrrr)rrCrrl list_reprstartends r%rrs| 6A:K 4@?FE;/QOI *$5:JE3 9fe[% EE s+A!cd\}}tt||k\rd}d|z}d|z}d||j|} ||| |S)N)rr?rr')rrr) rrCrrrr delimitercurr_indent_strrrs r%rrsk"(I 4 %  +vJ  {?+ , 1 1* =DWZLse ,,r$ceZdZdZ d d dZe d d dZed dZed dZ ed dZ dddZ y)NamespaceHelperaA helper for constructing the namespace open and close strings for a nested set of namespaces. e.g. for namespace_str torch::lazy, prologue: namespace torch { namespace lazy { epilogue: } // namespace lazy } // namespace torch cL|jd}t||ksJd|d|d||_dj|Dcgc]}d|d c}|_djt |Dcgc]}d| c}|_||_||_ycc}wcc}w) N::z"Codegen doesn't support more than z# level(s) of custom namespace. Got .r?z namespace z {z} // namespace ) r,rcpp_namespace_r prologue_reversed epilogue_ namespaces_ entity_name_)re namespace_str entity_name max_levelcpp_namespacesns r%rfzNamespaceHelper.__init__s',,T2>"i/ 0 ;^_l^mmn o /,#PAj3$7#PQ-5n-E Fs # F *' $Q Fs  B5 B!cr|jd}|d}dj|dd}t|||S)zq Generate helper from nested namespaces as long as class/function name. E.g.: "torch::lazy::add" rN)r!r"r#)r,rr)namespaced_entityr#namesr"r!s r%from_namespaced_entityz&NamespaceHelper.from_namespaced_entitysE"''-Bi  %*- '[I  r$c|jSr6)rrs r%prologuezNamespaceHelper.prologue ~~r$c|jSr6)rrs r%epiloguezNamespaceHelper.epiloguer-r$c|jSr6)r rs r%r"zNamespaceHelper.entity_names   r$c6|jr |jS|S)z Return the namespace string from joining all the namespaces by "::" (hence no leading "::"). Return default if namespace string is empty. )r)redefaults r%get_cpp_namespacez!NamespaceHelper.get_cpp_namespaces '+&9&9t""FwFr$N)r)r!rr"rr#rYrr)r4)r(rr#rYrrr)r)r2rrr) rrr__doc__rf staticmethodr*rr,r/r"r3r#r$r%rrs   ((( (  (&           !!Gr$rcteZdZUded<dddZddZddZddZddZdd Z e dd Z dd Z dd Z dd Zy) OrderedSetzdict[T, Literal[None]]storageNcL|i|_ytj||_yr6)r9rfromkeys)reiterables r%rfzOrderedSet.__init__s  DL==2DLr$c||jvSr6r9rers r% __contains__zOrderedSet.__contains__st||##r$cHt|jjSr6)iterr9keysrs r%__iter__zOrderedSet.__iter__sDLL%%'((r$cN|jj|jyr6)r9r)rers r%rzOrderedSet.updates EMM*r$c"d|j|<yr6r>r?s r%rzOrderedSet.adds! Tr$cXt}|jj|_|Sr6)r8r9r)rerets r%rzOrderedSet.copy!s"'\ll'')  r$cb|dj}|ddD]}|j||S)Nrr@)rr)rDrHr[s r%unionzOrderedSet.union&s61gllnab A JJqM  r$c.tj||Sr6)r8rJreothers r%__or__zOrderedSet.__or__-se,,r$c(|j||Sr6)rrLs r%__ior__zOrderedSet.__ior__0s E r$ct|tr|j|jk(St|jj |k(Sr6)rr8r9rbrCrLs r%__eq__zOrderedSet.__eq__4s< eZ (<<5==0 0t||((*+u4 4r$r6)r<zIterable[T] | Nonerr)rr3rr)rz Iterator[T])r OrderedSet[T]rr)rr3rr)rrS)rDrSrrS)rMrSrrS)rMrSrr)rMobjectrr)rrr__annotations__rfr@rDrrrr6rJrNrPrRr#r$r%r8r8 sI ##3 $)+"  -5r$r8)r-rrztuple[str, list[str]])r7zCallable[[T], S | None]r8rr Iterator[S])r7zCallable[[T], Sequence[S]]r8rrrV)rEzCallable[[], str]rzIterator[None])r9r rr )rNrrr)r[rrrYr6)rrr`zstr | Path | Nonerr^)rP)rr rCrYrrYrr)r) rr rCrYrrYrrYrr) rzdict[Any, Any]rCrYrrYrrYrr) rz&list[Any] | set[Any] | tuple[Any, ...]rCrYrrYrrYrr)rz list[str]rCrYrrYrrYrrrrrr)? __future__r contextlib functoolsrUrr(rrB dataclassesrrenumrrpathlibrtypingr r r r r rrtyping_extensionsrrrtorchgen.code_templaterargparsercollections.abcrrrrrrrrrr IDENT_REGEXr2r3r4r;r=contextmanagerrHcacherOr\r^rrrrrrrr8r#r$r%rfsn" ,TTT<</"<<X'')00   $T $ $ CL CL   <=D>D // 4SMSMr&*   "    5  5  5  5 5$ #  #  #  # #  #FD D D D D  D( F 0 F  F  F F  F-- - - -  - - -&BGBGJ,5,5r$