
    hM                        S SK r S SKrS SKJr  S SKrS SKrS SKJrJrJ	r	J
r
JrJrJrJrJrJrJrJrJrJrJr  S SKJr  SSKJr  \R4                  S 5       rS rS	 rS
 rS rS r S r!S r"S r#S r$S r%S r&S r'S r(S r)S r*S r+S r,S r-S r.S r/S r0\Rb                  Re                  SSRg                  5       5      S 5       r4S  r5\Rb                  Rm                  \(       + S!S"9S# 5       r7S$ r8S% r9S& r:S' r;S( r<S) r=S* r>S+ r?S, r@g)-    N)Counter)AdamCupyOpsDropoutLinearModelReluShimSoftmaxchainchange_attr_valuesconcatenateset_dropout_rateuse_ops
with_debugwrap_model_recursive)has_cupy_gpu   )make_tempdirc                      [        5       $ N)r        V/home/james-whalen/.local/lib/python3.13/site-packages/thinc/tests/model/test_model.pymodel_with_no_argsr      s	    8Or   c                     [        U S 5      $ )Nc                     U S 4$ )Nc                     U $ r   r   dYs    r   <lambda>0create_model.<locals>.<lambda>.<locals>.<lambda>$   s    r   r   Xs    r   r!   create_model.<locals>.<lambda>$   s	    !]!3r   r   names    r   create_modelr)   #   s    344r   c                 F    [        U R                  [        5      (       a   eg r   )
isinstanceopsr   )r   s    r   test_model_defaults_to_cpur-   '   s    ,00':::::r   c                 d    [        5       n[        5       nUR                  UR                  :w  d   eg r   )r   id)r   model1model2s      r   test_models_get_different_idsr2   +   s&    XFXF99		!!!r   c                      " S S[         5      n [        S5      n[        SS SS S.[        R                  " S5      S S	.US S
.SS0U " S 5      /X/S9nUR                  S5      (       d   eUR                  S5      R                  S:X  d   eUR                  S5      b   e[        R                  " [        5         UR                  S5        S S S 5        [        R                  " [        5         UR                  S5        S S S 5        UR                  S[        R                  " S5      5        UR                  S5      (       d   eUR                  S5      R                  S:X  d   eUR                  UR                  S4[        R                  " S5      05         [        R                  " UR                  S5      [        R                  " S5      5      (       d   e S S S 5        [        R                  " UR                  S5      [        R                  " S5      5      (       d   eUR!                  S5      (       a   eUR!                  S5      (       a   e[        R                  " [        5         UR#                  S5        S S S 5        UR                  SUR$                  R'                  S5      5        UR)                  SUR$                  R'                  S5      5        [        R                  " [*        5         UR-                  S[        R                  " S5      5        S S S 5        UR/                  S5      (       d   eUR1                  S5      S:X  d   e[        R                  " [        5         UR1                  S5        S S S 5        [        R                  " [*        5         UR1                  S5        S S S 5        UR3                  S5      (       d   eUR5                  S5      R6                  S:X  d   eUR3                  S5      (       a   e[        R                  " [        5         UR5                  S5        S S S 5        UR3                  S5      b   e[        R                  " [*        5         UR5                  S5        S S S 5        UR9                  SU5        UR3                  S5      (       d   eUR5                  S5      R6                  S:X  d   e[        R                  " [*        5         UR9                  S[        S5      5        S S S 5        SUR:                  ;   d   eSUR:                  ;  d   eUR:                  S   S:X  d   e[        R                  " [        5         UR:                  S     S S S 5        SUR:                  S'   UR=                  5       nUR6                  S:X  d   eg ! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GNn= f! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GNn= f! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GNN= f! , (       d  f       N= f)Nc                       \ rS rSrSrSrg)test_model_init.<locals>.MyShim2   testshimr   N__name__
__module____qualname____firstlineno__r(   __static_attributes__r   r   r   MyShimr5   2       r   r>   atestc                     U S 4$ )Nc                     U $ r   r   r   s    r   r!   3test_model_init.<locals>.<lambda>.<locals>.<lambda>8       r   r   r#   s    r   r!   !test_model_init.<locals>.<lambda>8   	    1m$r   
   nInOrH   Wbr@   rO   foobardimsparamsrefsattrsshimslayersrN   rO   r$   xyz)   r   rJ   rK   cbaz)r
   r)   r   numpyzeros	has_param	get_paramshapepytestraisesKeyError	set_param
use_paramsr/   onesarray_equalhas_gradget_gradr,   alloc1fset_grad
ValueErrorinc_gradhas_dimget_dimhas_refget_refr(   set_refrW   copy)r>   model_amodel
model_copys       r   test_model_initry   1   s     3G$d#[['d3&end|n!	E ??3??3%%...??3'''	x	  
!	x	  
!	OOCU+,??3??3%%...			EHHc?EJJu,=>	?  !5uzz%7HIIII 
@U__S15;;u3EFFFF~~c""""~~e$$$$	x	 s 
!	OOC**2./	NN3		))"-.	z	"sEKK/0 
#===="$$$	x	 e 
!	z	"d 
#====""c)))}}U####	x	 e 
!==%%%	z	"c 
#	MM#w====""c)))	z	"c<,- 
#EKK###;;u&&&	x	 E 
!EKKJ??f$$$Y 
!	 	 	 
 
@	?
 
!	  
#	" 
!	 	"	"
 
!	  
#	"
 
#	"
 
!	 s   :V3.WAW0W):'W;2X&X+X13Y2Y)Y'3
W
W
W&)
W8;
X

X
X.1
Y 
Y
Y$'
Y5c                      " S S[         5      n [        S5      n[        SS SS S.S S S.US S	.S
S0U " S 5      /X/S9n[        R                  " [
        5         UR                  SS5        S S S 5        UR                  SSSS9  UR                  SUR                  R                  S5      5        UR                  SUR                  R                  S5      5        UR                  S5      (       d   eUR                  S5      S:X  d   e[        R                  " [        5         UR                  SS5        S S S 5        [        R                  " [
        5         UR                  SS5        S S S 5        [        R                  " [
        5         UR                  SSSS9  S S S 5        g ! , (       d  f       GND= f! , (       d  f       N= f! , (       d  f       Nf= f! , (       d  f       g = f)Nc                       \ rS rSrSrSrg)"test_model_set_dim.<locals>.MyShims   r7   r   Nr8   r   r   r   r>   r|   s   r?   r   r>   r@   rA   c                     U S 4$ )Nc                     U $ r   r   r   s    r   r!   6test_model_set_dim.<locals>.<lambda>.<locals>.<lambda>y   rE   r   r   r#   s    r   r!   $test_model_set_dim.<locals>.<lambda>y   rG   r   r[   rI   rM   rP   rQ   rR   rS   rJ   rH   T)forcerN   rZ      )r
   r)   r   rc   rd   rn   set_dimrf   r,   rl   rm   rp   rq   re   )r>   rv   rw   s      r   test_model_set_dimr   r   s{     3G$T"%&end|n!	E 
z	"dB 
# 
MM$$M'	OOC**2./	NN3		))"-.===="$$$	x	 eR  
!	z	"dB 
# 
z	"dBd+ 
#	" 
#	" 
!	 	"	" 
#	"s0   F)$F;GG)
F8;
G	
G
G+c                      [        S5      n U R                  [        5       :X  d   eU R                  SS 5        U R                  S:X  d   eU R                  SS 5        U R                  S:X  d   eg )Ntmpparam1r   param2)r   r   )r)   param_namestuplerf   rw   s    r   test_param_namesr      sg    E'''	OOHd#+++	OOHd# 4444r   c                  (   [        S5      n U R                  [        5       :X  d   eU R                  SU R                  R                  SS5      5        U R                  SU R                  R                  SS5      S-   5        U R                  S:X  d   eg )Nr   r         r   )r)   
grad_namesr   rf   r,   alloc2frm   r   s    r   test_grad_namesr      sy    Euw&&&	OOHeii//156	NN8UYY..q!4q89{***r   c                  @    [        SS5      n U R                  S:X  d   eg )Nr[      )rK   rJ   )r   	dim_namesr   s    r   test_dim_namesr      s    1aLE??l***r   c                     [        S5      n [        S5      n[        S5      nU R                  R                  U5        U R                  [	        5       :X  d   eU R                  SU5        U R                  S:X  d   eU R                  S5      UL d   eUR                  R                  U5        [        R                  " [        5         U R                  S5        S S S 5        U R                  SU5        U R                  S5      UL d   eU R                  U5        X!R                  ;  d   eU R                  S5      (       a   eg ! , (       d  f       Nq= f)Nparentchildkid)r   grandkid	grandkind)r)   rY   append	ref_namesr   rt   rs   rc   rd   re   remove_noderr   )r   r   
grandchilds      r   test_model_set_referencer      s   (#F!Eg&J
MMuw&&&
NN5% x'''>>% E)))	LL
#	x	 z" 
!
NN:z*>>*%333
z"\\)))~~k***** 
!	 s   E
Ec                  ,   [        S5      n U R                  S5      b   eU R                  SS5        U R                  S5      S:X  d   eU R                  S5      b   eU R	                  S5      b   eU R                  5         U R	                  S5      c   eg )Nr[   rJ   r   boorN   )r   maybe_get_dimr   maybe_get_refmaybe_get_param
initializer   s    r   test_maybe_methodsr      s    1IEt$,,,	MM$t$)))u%---  %---	  %111r   c                 t    [        5        nU R                  US-  5        S S S 5        g ! , (       d  f       g = fNthinc_model)r   to_disk)r   paths     r   test_model_can_save_to_diskr      s&    	4""4-#78 
s   )
7c                     [        5        nU R                  US-  5        U R                  US-  5      nS S S 5        U R                  5       WR                  5       :X  d   eg ! , (       d  f       N3= fr   )r   r   	from_diskto_bytes)r   r   m2s      r   test_model_can_load_from_diskr      s]    	4""4-#78))$*>? 
 &&(BKKM999 
s   )A!!
A/c                     US-  nU R                  U5        U R                  U5      nU R                  5       UR                  5       :X  d   eg r   )r   r   r   )r   pathy_fixturer   r   s       r   +test_model_can_roundtrip_with_path_subclassr      sG    =(Dt$		%	%d	+B&&(BKKM999r   c                     U nSUl         SUR                  S'   [        USSSS.05        UR                  S   SL d   eSUR                  ;  d   eg )NtargetFhas_varT)r   errorr   )r(   rW   r   r   rw   s     r   test_change_attr_valuesr      sY    EEJ"EKK	uxTD)IJK;;y!T)))%++%%%r   c                      [        5       n U R                  S   S:X  d   e[        U S5        U R                  S   S:X  d   eg )Ndropout_rate        皙?)r   rW   r   r   s    r   test_set_dropoutr      sA    IE;;~&#---UC ;;~&#---r   c                 v    U nSUl         SUR                  S'   [        US5        UR                  S   S:X  d   eg )Ndropoutr   r   r   )r(   rW   r   r   s     r   test_set_dropout_2r      s=    EEJ"%EKKUC ;;~&#---r   c                      [         R                  " SS 05         [        SS9[        SS9-   n U S:X  d   e S S S 5        g ! , (       d  f       g = f)N+c                 2    U R                   UR                   4$ r   r'   rP   s     r   r!    test_bind_plus.<locals>.<lambda>   s    AFFAFF3Cr   r@   r'   rO   rP   )r   define_operatorsr)   ms    r   test_bind_plusr      sC    			&C D	Ec"\s%;;J 
F	E	Es   A  
Ac                      [         R                  " SS 05         [        SS9[        SS9-   [        SS9-   [        SS9-   n U R                  S:X  d   e S S S 5        g ! , (       d  f       g = f)Nr   c                     U $ r   r   rP   s     r   r!   !test_plus_chain.<locals>.<lambda>   s    1r   r@   r'   rO   r\   d)r   r   r)   r(   r   s    r   test_plus_chainr      si    			n 5	6c"$%$% $% 	
 vv}} 
7	6	6s   :A
A,c                     [         R                  " [        SS9n [         R                  " [        SS9nUR                  5         U R                  5         U R	                  5         UR	                  5         [         R                  " [        SS9n [         R                  " [        SS9nUR                  5         U R                  5         U R	                  5         UR	                  5         g)zGTest we can create a model in a child thread with overloaded operators.)r   r   r   args)*r   )r   r   )r   r   N)	threadingThread_overload_plusstartjoin)worker1worker2s     r   $test_overload_operators_in_subthreadr     s    
 n8DGn8DGMMOMMOLLNLLNn8DGn8DGMMOMMOLLNLLNr   c                 0   [        SS9n[        SS9n[        R                  " U S 05         [        R                  " U5        U S:X  a  X#-   nOX#-  nS S S 5        WS:X  d   e[        R
                  R                  5       0 :X  d   eg ! , (       d  f       N;= f)Nr@   r'   rO   c                 4    U R                   UR                   -   $ r   r'   rP   s     r   r!    _overload_plus.<locals>.<lambda>  s    r   r   ab)r)   r   r   timesleep_context_operatorsget)operatorr   m1r   values        r   r   r     s    	3	B	3	B			+G H	I

5s?GEGE 
J D==##'')R/// 
J	Is   &B
Bc            	      V   [        SS9n [        SS9n[        R                  R                  5       0 :X  d   e[        R                  " SS 05         X-   n[
        R                  " [        5         X-  nS S S 5        [        R                  " SS 05         [
        R                  " [        5         X-   nS S S 5        X-  n[        R                  " SS	 05         [
        R                  " [        5         X-   nS S S 5        X-
  nS S S 5        [
        R                  " [        5         X-   nS S S 5        X-  nS S S 5        X-   n[
        R                  " [        5         X-  nS S S 5        S S S 5        WS
:X  d   e[        R                  R                  5       0 :X  d   eg ! , (       d  f       GN4= f! , (       d  f       GN= f! , (       d  f       N= f! , (       d  f       N= f! , (       d  f       N= f! , (       d  f       N= f! , (       d  f       N= f! , (       d  f       N= f)Nr@   r'   rO   r   c                 4    U R                   UR                   -   $ r   r'   rP   s     r   r!   /test_nested_operator_contexts.<locals>.<lambda>'  s    166AFF?r   r   c                 4    U R                   UR                   -   $ r   r'   rP   s     r   r!   r   +  s    qvvr   -c                 4    U R                   UR                   -   $ r   r'   rP   s     r   r!   r   /  s    166AFF?r   r   )r)   r   r   r   r   rc   rd   	TypeError)r   r   r   s      r   test_nested_operator_contextsr   #  s   	3	B	3	B##'')R///			&B C	D]]9%GE &##S*F$GHy) *GE''.J(KL]]9-GE . M y) *GE I ]]9%GE &! 
E$ D==##'')R///# &% *) .- ML *) IH &%! 
E	Ds   H/F!4!HG80F35%G8G5G	:G"G8(G'-G89&HH	$H!
F0	+H3
G=G8
GG
G$ G8'
G51G88
H	H	
H	H
H(opz+ - * @ / // % ** << >> & ^ |c                    [        5       n[        5       n[        R                  " U S 05         U S:X  a  X-   nO'[        R                  " [
        5         X-   nS S S 5        U S:X  a  X-
  nO'[        R                  " [
        5         X-
  nS S S 5        U S:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S:X  a  UR                  U5      nO4[        R                  " [
        5         UR                  U5      nS S S 5        U S:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S	:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S
:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S:X  a  X-	  nO'[        R                  " [
        5         X-	  nS S S 5        U S:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        U S:X  a  X-  nO'[        R                  " [
        5         X-  nS S S 5        S S S 5        [        R                  R                  5       0 :X  d   eg ! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GNf= f! , (       d  f       GNF= f! , (       d  f       GN&= f! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GN= f! , (       d  f       GNf= f! , (       d  f       GNF= f! , (       d  f       GN&= f! , (       d  f       GN0= f)Nc                 4    U R                   UR                   -   $ r   r'   rP   s     r   r!   $test_all_operators.<locals>.<lambda>A  s    !&&166/r   r   r   r   @/z//^%z**z<<z>>&|)	r   r   r   rc   rd   r   
__matmul__r   r   )r   r   r   r   s       r   test_all_operatorsr   =  s   	B	B			%A B	C9GEy) *9GEy) * 9GEy) * 9MM"%Ey)b) * 9GEy) * :HEy) *9GEy) *9GEy) *:FEy) *:HEy) *:HEy) *9GEy) *9GEy) *9GEy) *S 
DV ##'')R///O *)
 *) *) *) *) *)
 *)
 *)
 *)
 *)
 *)
 *)
 *)
 *)S 
D	Cs[  &P.L2-P.M-P.8M=:P.7M(	-P.6M:;-P.(N--P.N-P.N0-P.>O-P.0O5-P."O&'-P.O8-P.P
-P.8P=P.2
M	<P.
M	P.
M%	 P.(
M7	2P.:
N		P.
N	P.
N-	(P.0
N?	:P.
O	P.
O#	P.&
O5	0P.8
P	P.

P	P.
P+	&P..
P=c                     / n S nSn[        U 5      S:  a  / n[        S5       H/  n[        R                  " XU4S9nUR	                  U5        US-  nM1     U H  nUR                  5         M     U H  nUR                  5         M     [        U 5      S:  a  M  [        U 5      [        [        [        U 5      5      5      :X  d   eg)z<Create a bunch of threads and assert they all get unique IDsc                 N    U R                  [        SU 3S9R                  5        g )Nworkerr'   )r   r)   r/   )id_listindexs     r   get_model_id3test_unique_id_multithreading.<locals>.get_model_id  s     |6%)9:==>r   r   i  r6   r   r   N)	lenranger   r   r   r   r   listset)list_of_idsr  counterworkersiws         r   test_unique_id_multithreadingr    s     K? G
k
T
!rA  ;QRANN1qLG  AGGI AFFH  k
T
! {s4K(8#9::::r   zneeds CuPy GPU)reasonc            	         [         R                  " S5        SS Kn [        S5         SnSnU R	                  5       u  u  p4u  pV[        [        XS9[        XS9[        5       5      nUR                  R                  U5      nUR                  R                  U5      nUR                  R                  U5      nUR                  R                  U5      nUR                  US S US S S9  [        S	5      nS
n	[        S5       H  n
UR                  R                  XUSS9nU H3  u  pUR                  U5      u  pU" X-
  5        UR                  U5        M5     SnSnUR                  R                  XU5       HX  u  pUR!                  U5      nUUR#                  SS9UR#                  SS9:H  R%                  5       -  nUUR&                  S   -  nMZ     M     S S S 5        g ! , (       d  f       g = f)Nml_datasetsr   cupy    r   )rK   r   r[   r$   YgMbP?   r   T)shuffler   )axis)rc   importorskipr  r   mnistr   r	   r   r,   asarrayr   r   r  
multibatchbegin_updatefinish_updatepredictargmaxsumrb   )r  n_hiddenr   train_Xtrain_Ydev_Xdev_Yrw   	optimizer
batch_sizer  batchesr$   r  Yhbackpropcorrecttotals                     r   test_model_gpur/    s   
&	-8->->-@*NUH.H.I
 ))##G,))##G,		!!%(		!!%(72A;'"1+6K	
qAii**:QU*VG$11!4 ##I.  
 GE		,,ZF]]1%BII1I-q1AAFFHH!$ G ' 
s   F0G
G-c            	      \   [        S5      n [        S5      n[        X5      n[        U 5      n[        SS XX/XU S.S9nUR	                  X5        UR
                  S   U:X  d   eUR
                  S   U:X  d   eUR
                  S   U:X  d   eUR                  S	5      U:X  d   eUR                  S
5      U:X  d   eUR                  S5      U:X  d   eUR
                  S   U:X  d   eUR
                  S   R
                  S   U:X  d   eUR
                  S   R
                  S   U:X  d   eg )Nr[   rA   c                     U S 4$ )Nc                     U $ r   r   r   s    r   r!   5test_replace_node.<locals>.<lambda>.<locals>.<lambda>  rE   r   r   r#   s    r   r!   #test_replace_node.<locals>.<lambda>  rG   r   )relu1relu2relu3)rY   rV   r   r   r   r5  r6  r7  r   )r	   r   r   r   replace_noderY   rs   )r5  r6  
relu_chainrelu1_debugdebugs        r   test_replace_noder<    s9   GEGEu$JU#K$e0u=	E 
u*<<?k)))<<?e###<<?k)))==![000==!U***==![000 <<?j(((<<?!!!$333<<?!!!$---r   c                     S n U " S/ 5      nU " SU/5      n[        U5      nU " SU/5      nUR                  SU5        [        X$5      nSUl        UR	                  X5        UR
                  S   R
                  S   U:X  d   eUR
                  S   R
                  S   R
                  S   U:X  d   eUR
                  S   R                  S5      U:X  d   eg )Nc                     [        U S US9$ )Nc                     gN.r   rw   r$   is_trains      r   r!   Otest_replace_node_with_indirect_node_ref.<locals>.dummy_model.<locals>.<lambda>      cr   rY   r&   r(   rY   s     r   dummy_model=test_replace_node_with_indirect_node_ref.<locals>.dummy_model      T9&IIr   yxrO   r@   r   r   )r   rt   r   r(   r8  rY   rs   )rG  rJ  rK  y_debugrO   r@   s         r   (test_replace_node_with_indirect_node_refrM    s    J 	CAC!AmGC!AIIc1aAAFNN188A;a G+++88A;a ''*g55588A;s#w...r   c                  d  ^ [         R                  " S5        SS Kn U R                  5       u  u  pu  p4[	        5       mU4S jnU4S jnU4S jn[        5       n[        [        5       XVUS9n	[        XU	5      n
U
R                  US S US S S9  U
" US S S	S
9u  pU" US S 5        TSSSS.:X  d   eg )Nr  r   c                      > TS==   S-  ss'   g )Ninitr   r   _countss    r   on_init test_with_debug.<locals>.on_init  s    v!r   c                      > TS==   S-  ss'   g )Nforwardr   r   rQ  s    r   
on_forward#test_with_debug.<locals>.on_forward  s    yQr   c                      > TS==   S-  ss'   g )Nr,  r   r   rQ  s    r   on_backprop$test_with_debug.<locals>.on_backprop  s    zar   )rT  rX  r[  r[   r  F)r$   rB  r   r   )rP  rW  r,  )	rc   r  r  r  r   r	   r   r   r   )r  r$  r%  r&  r'  rT  rX  r[  relur6  chainedrR  r,  rS  s                @r   test_with_debugr_    s    
&)4):):)<&WYF  6DKE D'G!4GBQK%8KA WRa[ aA1====r   c                    ^ S m[        S5      n [        X 5      n[        UU4S j5      nUR                  S:X  d   eUR                  S   UL d   eUR                  S   R                  S   R                  S:X  d   eUR                  S   R                  S   R                  S   U L d   eUR                  S   R                  S   R                  S:X  d   eUR                  S   R                  S   R                  S   U L d   eg )Nc                     [        U S US9$ )Nc                     gr@  r   rA  s      r   r!   :test_recursive_wrap.<locals>.dummy_model.<locals>.<lambda>'  rD  r   rE  r&   rF  s     r   rG  (test_recursive_wrap.<locals>.dummy_model&  rI  r   r[   c                 4   > T" SU R                    S3U /5      $ Nzdummy()r'   rw   rG  s    r   r!   %test_recursive_wrap.<locals>.<lambda>1  s    {VEJJ<q+AE7Kr   zdummy(relu>>relu)r   dummy(relu)r   )r	   r   r   r(   rY   )r]  r^  chained_debugrG  s      @r   test_recursive_wraprl  %  s   J 7DDG(KM !4444"g---"))!,11]BBB"))!,33A6$>>>"))!,11]BBB"))!,33A6$>>>r   c                    ^ S m[        S5      n [        X 5      n[        XU 5      n[        UU4S j5      nSnUR	                  5        H*  nUR
                  R                  S5      (       d  M%  US-  nM,     US:X  d   eUR                  S   R                  S   R                  S   R                  S   R
                  S:X  d   eUR                  S   R                  S   R                  S   R                  S   R
                  S:X  d   eUR                  S   R                  S   R                  S   R                  S   R
                  S:X  d   eUR                  S   R                  S   R                  S   R                  S   R
                  S:X  d   eUR                  S   R                  S	   R
                  S:X  d   eg )
Nc                     [        U S US9$ )Nc                     gr@  r   rA  s      r   r!   Atest_recursive_double_wrap.<locals>.dummy_model.<locals>.<lambda>>  rD  r   rE  r&   rF  s     r   rG  /test_recursive_double_wrap.<locals>.dummy_model=  rI  r   r[   c                 4   > T" SU R                    S3U /5      $ rf  r'   rh  s    r   r!   ,test_recursive_double_wrap.<locals>.<lambda>D  s    kF5::,a*@5'Jr   r   dummyr   r   rj  r   )r	   r   r   r   walkr(   
startswithrY   )r]  r^  concatconcat_wrappedn_debugrw   rG  s         @r   test_recursive_double_wraprz  <  s   J 7DDG40F)JN G$$&::  ))qLG ' a<<  #**1-44Q7>>qAFF-WWW  #**1-44Q7>>qAFF-WWW  #**1-44Q7>>qAFF-WWW  #**1-44Q7>>qAFF-WWW  #**1-22mCCCr   c                      [        S5      n [        S5      n[        X 5      n[        X5      nUR                  SU 5        [        U[        5        g )Nr[   r]  )r	   r   rt   r   r   )r]  r6  r^  chained2s       r   test_wrap_non_child_referencesr}  Y  sB    7DGEDGU$HVT":.r   c                      [        S5      n [        S5      n[        X5      n[        X"5      n[        UR                  SS95      X2X/:X  d   e[        UR                  SS95      U UUU/:X  d   eg )Nr[   dfs_preorderdfs_post)r	   r   r  ru  )r]  r6  inner_chainr^  s       r   test_walk_dfsr  c  s~    7DGE$KK-G9-.72UUUU:./	4   r   c                      [        S5      n [        R                  " [        SS9   U R	                  SS9  S S S 5        g ! , (       d  f       g = f)Nr[   zInvalid order)matchdfs_post_orderr  )r	   rc   rd   rn   ru  )r]  s    r   test_walk_bfs_post_order_failsr  q  s4    7D	z	9		(	) 
:	9	9s	   >
Ac                      " S S[         5      n [        S5      nU " S 5      n[        SS SSS.[        R                  " S5      S S	.US S
.SS0U/X/S9n[        SS SSS.[        R                  " S5      S S	.US S
.SS0U/X/S9n[        S5      n[        X5X45      nUR                  5       nUR                  UR                  5       5        UR                  S:X  d   eUR                  S   UR                  S   :X  d   e[        UR                  S   R                  S   5      [        UR                  S   R                  S   5      :X  d   eg )Nc                   $    \ rS rSrSrS rS rSrg))test_model_copy_with_loop.<locals>.MyShimix  r7   c                     [         $ r   )rM  )selfs    r   r   2test_model_copy_with_loop.<locals>.MyShim.to_bytes{  s    ;;r   c                     g r   r   )r  bytess     r   
from_bytes4test_model_copy_with_loop.<locals>.MyShim.from_bytes~  s    r   r   N)r9   r:   r;   r<   r(   r   r  r=   r   r   r   r>   r  x  s    	<	r   r>   r@   rA   c                     U S 4$ )Nc                     U $ r   r   r   s    r   r!   =test_model_copy_with_loop.<locals>.<lambda>.<locals>.<lambda>  rE   r   r   r#   s    r   r!   +test_model_copy_with_loop.<locals>.<lambda>  rG   r   r[   rI   rL   rM   rP   rQ   rR   rS   test2c                     U S 4$ )Nc                     U $ r   r   r   s    r   r!   r    rE   r   r   r#   s    r   r!   r    rG   r   ztest>>relu>>test>>test2r   r   r   )r
   r)   r   r^   r_   r	   r   ru   	from_dictto_dictr(   rY   r/   rX   )r>   rv   working_shimlayerlayer2r]  rw   r1   s           r   test_model_copy_with_loopr  w  sS     3G$<L$Q[['d3&enn!	E $Q[['d3&enn!	F 7D%u-EZZ\F	OOFNN$%;;3333==v}}Q////fmmA$$Q'(Bv}}Q/?/E/Ea/H,IIIIr   )Ar   r   collectionsr   r^   rc   	thinc.apir   r   r   r   r   r	   r
   r   r   r   r   r   r   r   r   thinc.compatr   utilr   fixturer   r)   r-   r2   ry   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   markparametrizesplitr   r  skipifr/  r<  rM  r_  rl  rz  r}  r  r  r  r   r   r   <module>r     sM           " &   5;">%B,@5++
+&29
::&..(
004 >DDFGN0 HN0b;. $-=>#% ?#%L.2/<>B?.D:/*&Jr   