
    Αix                     r   S SK r S SKrS SKJr  SSKJrJr  S r " S S5      r\" 5       r	/ SQr
S	/rS
 r\" SSS9S 5       r\" SSS9S 5       r\" SSS9S 5       r\" SSS9S 5       r\" SSS9S 5       r\" SSS9S 5       r\" SSS9S 5       r " S S5      rS S!/r\" S"SS9S# 5       r\" S$S%S9S& 5       rg)'    N)product   )Token	TokenTypec                    ^ ^ U U4S jnU$ )Nc                 6   > [         R                  TU T5        U $ N)macro_registryregister_macro)funcnameprioritys    m/var/www/html/banglarbhumi/venv/lib/python3.13/site-packages/paddle/distributed/flex_checkpoint/aoa/macros.py	decoratormacro.<locals>.decorator   s    %%dD(;     )r   r   r   s   `` r   macror      s     r   c                   8   ^  \ rS rSrSrU 4S jrS rS rSrU =r	$ )MacroRegistry   Nc                 ^   > U R                   c  [        TU ]	  U 5      U l         U R                   $ r	   )	_instancesuper__new__)clsargskwargs	__class__s      r   r   MacroRegistry.__new__"   s'    == !GOC0CM}}r   c                 6    [        U S5      (       d  / U l        g g )Nmacros)hasattrr"   )selfs    r   __init__MacroRegistry.__init__'   s    tX&&DK 'r   c                    ^ [        U4S jU R                   5       5      (       a  [        ST S35      eU R                  R                  TX#S.5        U R                  R	                  S SS9  g )Nc              3   2   >#    U  H  oS    T:H  v   M     g7f)r   Nr   ).0r   r   s     r   	<genexpr>/MacroRegistry.register_macro.<locals>.<genexpr>,   s     >+V}$+s   zMacro 'z' is already registered.)r   r   r   c                     U S   $ )Nr   r   xs    r   <lambda>.MacroRegistry.register_macro.<locals>.<lambda>/   s    q}r   F)keyreverse)anyr"   
ValueErrorappendsort)r$   r   r   r   s    `  r   r   MacroRegistry.register_macro+   s\    >$++>>>wtf,DEFFD$MN4eDr   )r"   )
__name__
__module____qualname____firstlineno__r   r   r%   r   __static_attributes____classcell__)r   s   @r   r   r      s    I
E Er   r   )
axis	fused_ffnfused_qkv_old	num_headsnum_key_value_groupspermutedtype	fused_qkv	src_dtype	dst_dtype^Tc                 V   Sn[        U 5       H  u  p#UR                  S:X  d  M  US-   [        U 5      :  d  M+  [        XS-      R                  5      nUS-   nU[        U 5      S-
  :  a.  X   R                  S:X  d   SX   R                   S35       eUS-  nU S U XS  -   n   X4$    X4$ )Nr   r>         ,zEThe different attributes must split by a comma, but now the token is .)	enumeratevaluelenint)tokensr>   idxtknend_idxs        r   extract_axis_and_clean_tokensrV   F   s    Df%9937S[#8vAg,,-DAgGVq(,,3 [\b\k\q\q[rrst3 1DS\F8$44F< & <r   
star_macrorK   )r   r   c           	      &   SnX1;  a  U$ S nSn/ nU  GHI  nUR                   [        R                  :X  a  SnUR                   [        R                  :X  a  X7R                  ;   a  UR                  R                  U5      u  pU(       d  UR                  5       OUR                  5       n
[        U
5      S:w  d   SU SU	 SU(       d  S	OS
 S35       e[        U" XU
5      5      nU H_  nUR                  [        [        R                  U5      5        XS   :w  d  M6  UR                  [        [        R                  S5      5        Ma     GM8  UR                  U5        GML     SR                  U Vs/ s H  owR                  PM     sn5      nU$ s  snf )N*c                 z   [         R                  " [         R                  " U 5       S[         R                  " U5       35      n/ nU HI  nUR                  U5      nU(       d  M  [	        UR                  S5      5      nUR                  XW45        MK     [        US S9nU VV	s/ s H  u  pYUPM	     sn	n$ s  sn	nf )Nz(\d+)r   c                     U S   $ )Nr   r   r-   s    r   r/   @star_macro.<locals>._sort_keys_by_numeric_part.<locals>.<lambda>e   s    !A$r   r1   )recompileescape	fullmatchrQ   groupr5   sorted)
prefixsuffixallkeyspatternfiltered_keysr1   matchnumsorted_keys_s
             r   _sort_keys_by_numeric_part.star_macro.<locals>._sort_keys_by_numeric_part]   s    **6!2 3569J8KLMC%%c*Eu%++a.)$$cZ0	 
 ]?"-.++...s   %B7TFr   zNo keys found with prefix 'z' and suffix 'z' in destination_state_shard_infosource_state_shard_infoz, please check!rL    )typer   RARROW
IDENTIFIERrO   splitget_all_dst_state_keysget_all_src_state_keysrP   listr5   r   COMMAjoin)rR   
expressioncontextSTAR_TAGrm   
pre_rarrow
new_tokenstokenrd   re   rf   keysr1   new_expressions                 r   rW   rW   W   sq   H!	/ JJ::)))J::---(kk2I"[[..x8NF " ..0335 
 w<1$ -fX^F859C1Ibccrt$ 267KLD!!%	(<(<c"BCr(?%%eIOOS&AB 
 e$) * WWzBzekkzBCN Cs   0Flayer_id_offset_macroc           	        ^ SmTU;  a  U$ [        U4S jU  5       S 5      nU(       d   S5       e[        U4S jU  5       5      (       d   ST S35       eUR                  UT5      n/ n[        U5      nU H  nSnSnU  H  n	U	R                  [
        R                  :X  a  S	nU(       a  Un
OUS
-
  n
U	R                  [
        R                  :X  ar  TU	R                  ;   a)  XyR                  R                  T[        U
5      5      -  nM  U	R                  [        ;  a  XyR                   SU
 3-  nM  XyR                  -  nM  XyR                  -  nM     UR                  U5        M     U$ )Nz$LAYER_ID_OFFSETc              3      >#    U  HC  nUR                   [        R                  :X  d  M#  TUR                  ;   d  M5  UR                  v   ME     g 7fr	   )rs   r   ru   rO   )r)   r   LAYER_ID_OFFSET_MACRO_TAGs     r   r*   (layer_id_offset_macro.<locals>.<genexpr>   sC      	
zzY111  *U[[8 EKKs   "AAAzENo $LAYER_ID_OFFSET found in NAME tokens.Please check the aoa_config.c              3      >#    U  HR  nUR                   [        R                  :g  =(       d)    TUR                  ;   =(       d    UR                  [        ;   v   MT     g 7fr	   )rs   r   ru   rO   GLOBAL_ATTRIBUTE_KEYWORDS)r)   tr   s     r   r*   r      sS       A 
9''	' 	2%0	2GG00	2 s   AAz#All IDENTIFIER tokens must contain zF when a NAME with it is present, except for GLOBAL_ATTRIBUTE_KEYWORDS.rr   TFr   z.layer.)nextallget_num_hidden_layersrc   rs   r   rt   ru   rO   replacestrr   r5   )rR   r|   r}   name_with_layer_id_offsetmatch_layer_id_offsetexpanded_expressionslayer_idexprbefore_rarrowr   cur_layer_idr   s              @r   r   r      s    2 
2 $	
	
 	! % O$   	    ..G-H  IO  	P  $99!#< "#89)EzzY--- %''!|zzY111,;KK//13|3D D [[(AA{{m7<.AADKK'D## $ 	##D)+ *,  r   array_macrorJ   c           
         SU;  a  U$ / nSnU[        U 5      :  Ga  X   R                  [        R                  :X  Ga  XS-
     R                  nXS-      R                  [        R
                  :X  ai  XS-      R                  [        R                  :X  aF  XS-      R                  [        R
                  :X  a#  XS-      R                  [        R                  :X  dK   SX   R                   XS-      R                   SXS-      R                   XS-      R                   S	35       eUR                  5         [        XS-      R                  5      n[        XS-      R                  5      n[        Xg5       Hn  nUR                  [        [        R                  US
-   [        U5      -   5      5        XS-
  :w  d  ME  UR                  [        [        R                  S5      5        Mp     US-  nOUR                  X   5        US-  nU[        U 5      :  a  GM  SR!                  U V	s/ s H  oR                  PM     sn	5      n
U
$ s  sn	f )N[r   r   rJ   rK      zbThe array macro format is incorrect which is must be like: NAME[START:END], but now the format is :rM   rl   rL      rr   )rP   rs   r   LBRACKETrO   NUMBERCOLONRBRACKETpoprQ   ranger5   r   ru   r   rz   r{   )rR   r|   r}   r   rS   r   startendir   r   s              r   r   r      s   
*J
C
F
;y111'?((DQw$$	(8(887O((IOO;7O((I,<,<<7O((I,>,>>	 uu{  vA  vG  vG  uH  IO  VW  PW  IX  I^  I^  H_  _`  ag  no  ho  ap  av  av  `w  x~  EF  F  xG  xM  xM  wN  NO  P? NNQw--.Ef1Wo++,C5&!!)..s
SV0CD a<%%eIOOS&AB ' 1HCfk*1HC/ F
0 WWzBzekkzBCN Cs   &Ifused_qkv_old_macro   c                 	  ^^  Sm[        U4S jU  5       5      (       d  U$ [        U 5      u  m n S nS nS nS nS n[        U 5       GH  u  pU	R                  [        R
                  :X  a  U	R                  S:X  a-  US-   [        U 5      :  a  [        XS-      R                  5      nOsU	R                  S:X  a-  US-   [        U 5      :  a  [        XS-      R                  5      nO6U	R                  T:X  a  UnO#U	R                  [        R                  :X  a  Uc  UnUb  M  U	R                  [        R
                  :X  d  M  U	R                  TSS1;   d  GM  US-   nGM     U(       a  US:  d   SU S	35       eU(       a  US:  d   S
U S	35       eUc   ST S35       eUc   S5       eX4-  S:X  d   SU SU S	35       e/ n
UnUS:X  Ga  U S   R                  nUS:  a  S nOU S   R                  nUR                  (       a(  UR                  U5      nUb  UR                  U5      OSnO'UR                  U5      nUb  UR                  U5      OSnX4X4/nSU4SU4SU4/nS n[        U5       H  u  nu  nn[        U5       VVVs/ s H  nU  H  u  nnU" UUUU5      PM     M     nnnnUS:X  a,  U SSR                  U5       ST  3nU
R                  U5        Mp  Uc  Mu  SR                  U5       SU ST  3nU
R                  U5        M     US:  a~  U 4S jnU S   R                  nU S   R                  nU S   R                  nU
R                  U" SUU5      5        U
R                  U" SUU5      5        U
R                  U" SUU5      5        U
$ US:X  Ga%  U S   R                  nU S   R                  nU S   R                  nU S   R                  nU SU SU S3nU
R                  U SU SU SU ST  3	5        UR                  U5      nSU4X4/nSU4SU4SU4/nS n[        U5       H  u  nu  nn[        U5       VVVs/ s H  nU  H  u  nnU" UUUU5      PM     M     nnnnUS:X  a  U SSR                  U5       ST  3nOSR                  U5       SU ST  3nU
R                  U5        M     U
$ [!        SU S35      es  snnnf s  snnnf )Nr@   c              3   @   >#    U  H  oR                   T:H  v   M     g 7fr	   rO   )r)   rT   FUSED_QKV_OLD_TAGs     r   r*   &fused_qkv_old_macro.<locals>.<genexpr>   s     @#yy--   rA   rJ   rB   r   r   z!num_heads must be positive.(got: ).z,num_key_value_groups must be positive.(got: z:No fused_qkv_old tag found in expression. The tag must be rM   No -> found in expression.num_heads (-) must be divisible by num_key_value_groups (r   QKVc                 h   ^ X!-  U -  nX-  nSR                  U4S j[        XDU-   5       5       5      $ )NrL   c              3   4   >#    U  H  nS T SU 3v   M     g7fzfused_qkv_old_tmp.rl   Nr   r)   r   comps     r   r*   8fused_qkv_old_macro.<locals>.gen_expr.<locals>.<genexpr>:  %      4A %TF!A3/4   r{   r   	tp_degreerA   tp_rankr   r   counts      `  r   gen_expr%fused_qkv_old_macro.<locals>.gen_expr7  ?    '94E*E88 uem4  r    -> rL   , axis=c                 b   >^  SR                  U 4S j[        U5       5       5      nU SU ST 3$ )NrL   c              3   4   >#    U  H  nS T SU 3v   M     g7fr   r   r)   r   rd   s     r   r*   >fused_qkv_old_macro.<locals>._generate_expr.<locals>.<genexpr>S  s"      $?K!(!5|r   r   r   r   rd   r   target_nameelementsr>   s   `   r   _generate_expr+fused_qkv_old_macro.<locals>._generate_exprR  s;    88 $?DU|$  #4}GD6BBr   r   r   .tmp  ->  c                 h   ^ X!-  U -  nX-  nSR                  U4S j[        XDU-   5       5       5      $ )NrL   c              3   4   >#    U  H  nS T SU 3v   M     g7fr   r   r   s     r   r*   r   {  r   r   r   r   s      `  r   r   r   x  r   r   z(Unsupported fused_qkv_old macro format: )r3   rV   rN   rs   r   ru   rO   rP   rQ   rt   aoa_config_reverseget_src_state_shard_numget_dst_state_shard_numr   r{   r5   r4   )!rR   r|   r}   attn_head_numrB   fused_qkv_old_pos
rarrow_posright_var_end_posrS   r   resultsnum_key_value_headssrc_qkv_weight_namedst_qkv_weight_namedst_state_shard_numsrc_state_shard_numconfigshead_configr   r   qkv_weight_namer   cn	qkv_partsmappingr   q_namek_namev_namefused_qkv_tmp_namer   r>   s!                                  @@r   r   r      s   '@@@@08LD&MJ'
::---{{k)cAgF.C #F7O$9$9 : 6637SF < (+6'?+@+@'A$ 11$'!ZZ9+++
0BJ%

i222!;0FGH !$a% (( ]Q. 
+M?"=.  $81$< 
67K6LBO< ( 
DEVDWWXY( !?#??!/14 
m_$QRfQggij4 G.Q$Qiooq "&"()//%%")"A"A##
 '2 //0CD   #*"A"A##
 '2 //0CD   !6 6
 - %&%&
	 2;71C-C-)_  %Y//G'DAq Aw2' 3/  
 ax&'tCHHY,?+@vN  w' ,xx	*+4/@vN  w' 2D" q C AY__FAY__FAY__FNN>#}fEFNN>#/BFKLNN>#/BFKLj Ni 
q$Qioo &xq&>haxq/A.B'$P	
 &==

 "# 6
 - %&%&
	 2;71C-C-)_  %Y//G'DAq Aw2' 3/  
 ax&'tCHHY,?+@vN 
 xx	*+4/@vN  NN7# 2D& N 6zl!D
 	
aBs   "S
"S
fused_ffn_macroc                 l  ^^ Sm[        U4S jU  5       5      (       d  U$ [        U 5      u  mn S nS n[        U 5       H^  u  pVUR                  [        R
                  :X  a  Uc  UnM*  UR                  [        R                  :X  d  MJ  UR                  T:X  d  M\  UnM`     Uc   S5       eUc   ST S35       e/ nUS:X  Ga  U S   R                  nUS:X  a  U S	   R                  n	OS n	UR                  (       a(  UR                  U	5      n
Ub  UR                  U5      OSnO'UR                  U5      nU	b  UR                  U	5      OSn
[        R                  " X5      nX4X4/nS
U4SU4/nS n[        U5       H  u  nu  nn[        U5       VVVs/ s H  nU  H  u  nnU" UUUU5      PM     M     nnnnUS:X  a*  UR                  U SSR                  U5       ST 35        Mn  Uc  Ms  UR                  SR                  U5       SU ST 35        M     US:  aV  U4S jnU S	   R                  nU S   R                  nUR                  U" S
UU5      5        UR                  U" SUU5      5        U$ US:X  Ga  U S   R                  nU S	   R                  nU S   R                  n	U SU S3nUR                  U SU SU ST 35        UR                  U	5      n
SU4X4/nS
U
4SU
4/nS n[        U5       H  u  nu  nn[        U5       VVVs/ s H  nU  H  u  nnU" UUUU5      PM     M     nnnnUS:X  a*  UR                  U SSR                  U5       ST 35        Mn  UR                  SR                  U5       SU ST 35        M     U$ [!        SU S35      es  snnnf s  snnnf )Nr?   c              3   @   >#    U  H  oR                   T:H  v   M     g 7fr	   r   )r)   rT   FUSED_FFN_TAGs     r   r*   "fused_ffn_macro.<locals>.<genexpr>       <VcyyM)Vr   r   z6No fused_ffn tag found in expression. The tag must be rM   r   r   r   rJ   GATEUPc                 ^   ^ ^^^ SR                  UUU U4S j[        TT -  5       5       5      $ )NrL   c              3   F   >#    U  H  nS T STT-  T-  U-    3v   M     g7fzfused_ffn_tmp.rl   Nr   r)   rS   r   splited_numr   r   s     r   r*   4fused_ffn_macro.<locals>.gen_expr.<locals>.<genexpr>  5      :C !a+(=(JS(P'QR:   !r   r   r   r   r   s   ````r   r   !fused_ffn_macro.<locals>.gen_expr  +    88  	!9:  r   z  -> rL   r   r   c                 b   >^  SR                  U 4S j[        U5       5       5      nU SU ST 3$ )NrL   c              3   4   >#    U  H  nS T SU 3v   M     g7fr   r   r   s     r   r*   :fused_ffn_macro.<locals>._generate_expr.<locals>.<genexpr>  s!      $;GanVHAaS1<r   r   r   r   r   s   `   r   r   'fused_ffn_macro.<locals>._generate_expr  s;    88 $;@<$  #4}GD6BBr   rK   r   r   c                 ^   ^ ^^^ SR                  UUU U4S j[        TT -  5       5       5      $ )NrL   c              3   F   >#    U  H  nS T STT-  T-  U-    3v   M     g7fr   r   r   s     r   r*   r     r   r   r   r   s   ````r   r   r     r   r   z$Unsupported fused_ffn macro format: )r3   rV   rN   rs   r   rt   ru   rO   r   r   r   mathlcmr   r5   r{   r4   )rR   r|   r}   r   fused_ffn_posrS   r   r   src_ffn_weight_namedst_ffn_weight_namer   r   r   r   split_configr   r   ffn_weight_namer   r   r   	ffn_partsr   	gate_nameup_namefused_gate_up_tmp_namer   r>   s                             @@r   r   r     sp   M<V<<<08LD&JM'
::)))j.@JJJ)...5;;-3OM ( !?#??!$ 
@qQ$ GQ$QiooA"()//"&%%")"A"A##
 '2 //0CD   #*"A"A##
 '2 //0CD  
 hh2H !6 6
  -k/BC	 2;71C-C-)_  %Y//G(DAq Aw2( 3/  
 ax&'uSXXi-@,AO !,xx	*+4/@vN 2D 1C q	IQiooGNN>&+yIJNN>$WEF^ N[ 
q1IOO	)//$Qioo$-;ay!=k7)6*@)AO	
 &==

 &' 6
 ()&'

	 2;71C-C-)_  %Y//G(DAq Aw2( 3/  
 ax&'uSXXi-@,AO xx	*+4/@vN 2D  N ?
|1MNNOrs   "N(
"N/
transpose_macror   c                    SnX1;  a  U$ [        5       nSnS n[        U 5       H'  u  pxUR                  [        R                  :X  d  M%  Un  O   Uc   S5       eXS-   S   HN  nUR                  [        R
                  :X  d  M#  UR                  R                  U5      (       d  ME  [        S5      e   U  H  nUR                  [        R
                  :X  aS  UR                  R                  U5      (       a3  UR                  S [        U5      *  n	UR                  U	5        XYS-   -  nMt  XXR                  -  nM     U V
s/ s H  n
U
 SU
 S3PM     nn
UR                  U5        U$ s  sn
f )	NrH   rr   r   r   zCannot assign to transpose (e.g., 'A -> B^T').
B^T is not a real variable, just a view.
Assign first:  A -> B
Then transpose: B^T -> B_transpose_tmpr   z_transpose_tmp, permute = "[]")setrN   rs   r   rt   ru   rO   endswithr4   rP   addr5   )rR   r|   r}   TRANSPOSE_TAGtranspose_varsr   r   rS   r   var_namevarr   s               r   r  r    sx   M&UNNJ'
::)))J (
 !?#??!Q()::---%++2F2F3
 3
 + 	 * ::---%++2F2F3
 3
 {{#8c-&8%89Hx()999Nkk)N  "!C %tC567!   NN>"Ns   =E$fused_qkv_macroc                   ^ Sm[        U4S jU  5       5      (       d  U$ [        U 5      u  p0S =pES nS nS n[        U 5       H  u  pU
R                  [        R
                  :X  a  U
R                  S:X  a.  U	S-   [        U 5      :  a  [        X	S-      R                  5      nMa  U
R                  S:X  a.  U	S-   [        U 5      :  a  [        X	S-      R                  5      nM  U
R                  T:X  a  U	nM  M  U
R                  [        R                  :X  d  M  Ub  M  U	nM     U(       a  US:  d   SU S35       eU(       a  US:  d   S	U S35       eUc   S
T S35       eUc   S5       eUS:X  d  US:X  d   S5       eXF-  S:X  d   SU SU S35       eXF-  nS n/ nUS:X  GaW  U S   R                  nXS-      R                  nXS-      R                  nXS-      R                  nU" X5      nU" UU5      nU" UU5      n/ n[        U5       HH  nUR                  UUU-  US-   U-   5        UR                  UU   5        UR                  UU   5        MJ     UR                  U SSR                  U5       SU 35        UR                  SR                  U5       SU SU 35        UR                  SR                  U5       SU SU 35        UR                  SR                  U5       SU SU 35        U$ US:X  GaS  U S   R                  nU S   R                  nU S   R                  nXS-      R                  nU" X5      nU" UU5      nU" UU5      nUR                  U SSR                  U5       SU 35        UR                  U SSR                  U5       SU 35        UR                  U SSR                  U5       SU 35        / n[        U5       HH  nUR                  UUU-  US-   U-   5        UR                  UU   5        UR                  UU   5        MJ     UR                  SR                  U5       SU SU 35        U$ U$ )NrE   c              3   @   >#    U  H  oR                   T:H  v   M     g 7fr	   r   )r)   rT   FUSED_QKV_TAGs     r   r*   "fused_qkv_macro.<locals>.<genexpr>O  r   r   rA   rJ   rB   r   z!num_heads must be positive (got: )z,num_key_value_groups must be positive (got: z6No fused_qkv tag found in expression. The tag must be rM   r   r   r   z>Only support q,k,v -> fused_qkv or fused_qkv -> q,k,v patternsr   r   r   c                 J    [        U5       Vs/ s H  o  U 3PM
     sn$ s  snf r	   )r   )baser   r   s      r   
make_names#fused_qkv_macro.<locals>.make_namesy  s$    &+Ah/h&h///s    rK   r   rL   r   r   )r3   rV   rN   rs   r   ru   rO   rP   rQ   rt   r   extendr5   r{   )rR   r|   r}   r>   r   rA   rB   fused_qkv_posr   rS   r   r   r#  r   fused_qkv_varq_vark_varv_varq_namesk_namesv_namesfused_qkv_ordergr  s                          @r   r  r  L  s   M<V<<<08LD $$MMJ'
::---{{k)cAgF.C #F7O$9$9 : 6637SF < (+6'?+@+@'A$- # .ZZ9+++
0BJ ( ]Q. 
+M?!<.  $81$< 
67K6LAN< $ 
@qQ$ !?#??!?jAo H- /14 
m_$QRfQggij4 (?0 GQq	A~&,,A~&,,A~&,,U2U$89U$89+,A""//1q5<O2OP ""71:.""71:. - 	oT#((?";!<GD6J	
 	#((7+,DwtfEF#((7+,DwtfEF#((7+,DwtfEF	qq	q	q	A~.44U2U$89U$89%SXXg%6$7wtfEF%SXXg%6$7wtfEF%SXXg%6$7wtfEF+,A""//1q5<O2OP ""71:.""71:. - 	xx()m_GD6J	
  r   c                       \ rS rSrS\\   S\\   S\\   4S jrS\S\4S jrS	\S\\	R                  \\   4   4S
 jrS	\S\\\4   S\4S jrS	\S\\\\   4   4S jrSrg)	IDMatcheri  source_keysextra_suffixesallowed_placeholdersc                     [        U5      U l        X0l        SR                  S U R                   5       5      n[        R
                  " SU S35      U l        [        US S9U l        g )N|c              3   N   #    U  H  n[         R                  " U5      v   M     g 7fr	   )r^   r`   )r)   phs     r   r*   %IDMatcher.__init__.<locals>.<genexpr>  s      '
$=bBIIbMM$=s   #%(r   c                     [        U 5      * U 4$ r	   )rP   r-   s    r   r/   $IDMatcher.__init__.<locals>.<lambda>  s    SVGQ<r   r]   )	r  r2  r4  r{   r^   r_   _placeholder_patternrc   r3  )r$   r2  r3  r4  placeholder_patterns        r   r%   IDMatcher.__init__  sg     {+$8!!hh '
$($=$='
 
 %'JJ3F2Gq/I$J!$^9OPr   r1   returnc                 |    U R                    H+  nUR                  U5      (       d  M  US [        U5      *  n  U$    U$ r	   )r3  r  rP   )r$   r1   sfxs      r   _remove_extra_suffixes IDMatcher._remove_extra_suffixes  sB    &&C||C  +SXI&
	 ' 
r   rg   c                 `   [        [        U R                  R                  U5      5      5      n[        R
                  " U5      nU H2  nUSS  nUR                  [        R
                  " U5      SU S35      nM4     [        R                  " SU S35      U Vs/ s H  oDSS  PM	     sn4$ s  snf )Nr   z(?P<z>\d+)^$)rc   r  r=  findallr^   r`   r   r_   )r$   rg   placeholders	regex_strr8  
group_names         r   _pattern_to_regexIDMatcher._pattern_to_regex  s    c$";";"C"CG"LMNIIg&	BABJ!))		"j\8I 
 zzAi[*+|-L|f|-LLL-Ls   B+id_dictc                 x    UnUR                  5        H#  u  pEUR                  SU 3[        U5      5      nM%     U$ )NrG  )itemsr   r   )r$   rg   rN  r1   r8  rO   s         r   _substitute_idsIDMatcher._substitute_ids  s8     IB++"hE
3C )
r   c                    U R                  U5      nU R                  U5      u  p#U Vs0 s H  oD[        5       _M     nnU R                   H`  nUR	                  U5      nU(       d  M  UR                  5       R                  5        H!  u  pXX   R                  [        U	5      5        M#     Mb     UR                  5        VV
s0 s H  u  pU[        U
5      _M     sn
n$ s  snf s  sn
nf r	   )
rC  rL  r  r2  ri   	groupdictrP  r  rQ   rc   )r$   rg   regexph_namesr8  	id_valuesr1   ri   kvvss              r   find_matchesIDMatcher.find_matches  s    --g6009)122Y	2##CKK$Eu!OO-335DAL$$SV, 6 $
 ,5??+<=+<%!6":+<== 3 >s   CC#)r=  r4  r3  r2  N)r8   r9   r:   r;   ry   r   r%   rC  tupler^   PatternrL  dictrQ   rQ  r[  r<   r   r   r   r1  r1    s    Q#YQ S	Q #3i	Q# # M Mrzz497L1M Ms T#s(^  	>C 	>Dd3i,@ 	>r   r1  z
$EXPERT_IDz	$LAYER_IDid_macroc           	        ^^^^ [         m[        U4S jT 5       5      nU(       d  T$ UR                  (       d  [        U4S jU  5       S 5      nO|SnU  Ht  mTR                  [
        R                  :X  a  SnTR                  [
        R                  :X  d  MC  [        U4S jT 5       5      (       d  M_  U(       d  Mh  TR                  n  O   Wc   S5       eUR                  5       n[        U[        T5      nUR                  U5      n[        UR                  5       5      n	/ n
U  HJ  mTR                  [        ;   a    O5TR                  [
        R                  :X  d  M9  U
R!                  T5        ML     U
 H5  m[#        U4S jU	 5       5      (       a  M   STR                   S	U	 S
35       e   S[$        [&        [        [(           4   4S jn/ nU" U5      n[+        U5      nU H  nSnU  H}  mTR                  nTR                  [
        R                  :X  aK  [        U4S jT 5       5      (       a1  U H$  u  nnUR-                  SU-   ['        U5      5      nM&     UU-  nMx  UU-  nM     UR!                  U5        M     U$ )Nc              3   ,   >#    U  H	  oT;   v   M     g 7fr	   r   )r)   r8  r|   s     r   r*   id.<locals>.<genexpr>  s      "#7Rj#7s   c              3      >^#    U  HM  mTR                   [        R                  :X  d  M#  [        U4S  jT 5       5      (       d  M?  TR                  v   MO     g7f)c              3   @   >#    U  H  oTR                   ;   v   M     g 7fr	   r   r)   r8  r   s     r   r*   id.<locals>.<genexpr>.<genexpr>  s     I4Hbekk)4Hr   N)rs   r   ru   r3   rO   )r)   r   r4  s    @r   r*   rc    sF      #E::!5!55  I4HII #s   "AAAFTc              3   @   >#    U  H  oTR                   ;   v   M     g 7fr	   r   rf  s     r   r*   rc    s      :,@bekk!,@r   zNo $ID found in NAME tokensc              3   @   >#    U  H  oTR                   ;   v   M     g 7fr	   r   )r)   rX  r   s     r   r*   rc    s     8Z#Zr   zThe token: z) must contain all of the following keys: zS.When use the id macro all IDENTIFIER tokens must contain the same ID placeholders.dc              3      #    [        U R                  5       5      nU Vs/ s H  o U   PM	     nn[        U6  H  n[        [	        X5      5      v   M     g s  snf 7fr	   )ry   r   r   r]  zip)rj  r   rX  value_listsprods        r   dict_cartesian_tuples!id.<locals>.dict_cartesian_tuples!  sM     AFFH~%)*TtT*[)DD(( * +s   AA-Arr   c              3   @   >#    U  H  oTR                   ;   v   M     g 7fr	   r   )r)   r8  rT   s     r   r*   rc  .  s      8*>Bcii*>r   rG  )_REGISTERED_PLACEHOLDERSr3   r   r   rs   r   rt   ru   rO   rx   r1  EXTRA_SUFFIXr[  ry   r   r   r5   r   r_  r   rQ   rc   r   )rR   r|   r}   has_allowed_placeholdername_with_idflag_right_varall_src_state_keys
id_matchervalid_id_combos
valid_keysIDENTIFIER_tokensro  r   id_combsid_combcur_statementtkn_valid_tagid_valr4  rT   r   s    `                 @@@r   idr    sm   3! "#7"  #%%# 
 EzzY---!%zzY111c :,@: 7 7 ">#(;;L  #B%BB# 779L*>J !--l;Oo**,-J;;33::---$$U+	  #8Z888 	
%++&OPZ|  \o  p	
8 #
)c49n!5 ) G$_5HhHCiiGxx9///C 8*>8 5 5 '.NFF%oocFlCKHG '.((  	}%  Nr   get_var_mapping_chain_macror   c                    Sn/ n/ nU  H  nUR                   [        ;   a    OUR                  [        R                  :X  a  SnUR                  [        R
                  :X  d  MY  UR                   n[         H  nUR                  U5      nM     U(       a  UR                  U5        M  UR                  U5        M     [        U5      S:X  d  [        U5      S:X  d   S5       e[        U5      S:X  a*  XRR                  US   '   U H  n	XBR                  U	'   M     U$ UUR                  US   '   U H  n
XRR                  U
'   M     U$ )NTFr   z|Left or right variable must have the only one element,the aoa_statements not support 'multiple var -> multiple var' pattern.r   )rO   r   rs   r   rt   ru   rs  removesuffixr5   rP   left_var_to_right_var_mappingright_var_from_left_var_mapping)rR   r|   r}   flag_left_varleft_var_listright_var_listrT   extra_suffix_removed_valuerB  	right_varleft_vars              r   r  r  ?  sH   MMN991188y'''!M88y+++),&#.;;C@ + $ $$%?@%%&@A  }"c.&9Q&>  	G> =QBP--mA.>?'IAN33I> ( 	  	//q0AB &H>L11(; &r   )r  r^   	itertoolsr   lexerr   r   r   r   r
   r   rs  rV   rW   r   r   r   r   r  r  r1  rr  r  r  r   r   r   <module>r     sg     	  #E E&   	
" L1%( &(V #a07  17 t MA& 'B !A.r /rj *D +DN ** +*Z *g +gT0> 0>h )+6  J#H $H^ )A6  7 r   