
    hM              
          S SK JrJrJrJrJrJrJr  SSKJ	r	  SSK
Jr  SSKJrJrJrJrJrJr  \" S\\\\\   \\   5      r\	R(                  " S5      S\\\4   S	\\\4   4S
 j5       rS\\\4   S\S\S	\\\4   4S jr SS\\\4   S\\   S\\   S	S4S jjrS rS\\\4   S\S\S	\\\4   4S jrS\\\4   S\S\S	\\\4   4S jrg)    )CallableListOptionalTupleTypeVarUnioncast   )registry)Model)Array2dFloats2dInts2dList2dPaddedRaggedSeqTzwith_list.v1layerreturnc                     [        SU R                   S3[        [        U /U R                   Vs0 s H  oU R                  U5      _M     snS9$ s  snf )Nz
with_list())initlayersdims)r   nameforwardr   	dim_namesmaybe_get_dim)r   r   s     P/home/james-whalen/.local/lib/python3.13/site-packages/thinc/layers/with_list.py	with_listr    
   sT    
UZZL"w:?//J/$E''--/J 
 Ks   AmodelXseqis_trainc           	         U R                   S   n[        U[        5      (       a  [        X1U5      $ [        U[        5      (       a  [        X1U5      $ [        [        [        [        4   U" [        [        U5      U5      5      $ )Nr   )r   
isinstancer   _padded_forwardr   _ragged_forwardr	   r   r   r   r   )r!   r"   r#   r   s       r   r   r      si     $)<<?E$uH55	D&	!	!uH55E$.)5fd1CX+NOO    NXYc                 v    U R                   S   R                  Ub  [        X5      OS Ub  [        X5      OS S9  g )Nr   )r)   r*   )r   
initialize	_get_list)r!   r)   r*   s      r   r   r   !   s;     
LLO!")E
D!")E
D  r(   c                     [        U[        5      (       a  U R                  R                  U5      $ [        U[        5      (       a0  U R                  R                  UR                  UR                  5      $ U$ N)r%   r   opspadded2listr   	unflattendatalengths)r!   seqs     r   r-   r-   *   sU    #vyy$$S))	C	 	 yy""388S[[99
r(   Xrc                 
  ^^^ U R                   R                  mU R                   R                  mU " T" UR                  UR                  5      U5      u  nmS[
        4UUU4S jjn[        T" U5      UR                  5      U4$ )NdYrc           
      ~   > [        T" T" T" U R                  U R                  5      5      5      U R                  5      $ r/   )r   r3   r4   )r8   flattenget_dXsr2   s    r   backprop!_ragged_forward.<locals>.backprop?   s3    GIchh<=>KK
 	
r(   )r0   r2   r:   r3   r4   r   )r   r6   r#   Ysr<   r:   r;   r2   s        @@@r   r'   r'   3   sq     		##IiiG
 	"''2::6AKB
f 
 
 '"+rzz*H44r(   Xpc                    ^^^ U R                   R                  mU R                   R                  mU " T" U5      U5      u  nmUUU4S jnT" U5      U4$ )Nc                 ,   > T" T" T" U 5      5      5      $ r/    )dYpr;   list2paddedr1   s    r   r<   !_padded_forward.<locals>.backpropT   s    7;s#3455r(   )r0   r1   rD   )r   r?   r#   r>   r<   r;   rD   r1   s        @@@r   r&   r&   H   sO     ))''K))''K
 B2KB6 r?H$$r(   )NN)typingr   r   r   r   r   r   r	   configr   r!   r   typesr   r   r   r   r   r   r   r   r    boolr   r   r-   r'   r&   rB   r(   r   <module>rJ      sa   H H H   E EvvvvtH~tF|L 
 U66>* uT4Z/@  !	Pt	P$(	P48	P
4>	P MQt!)$;CD>	5 5&,58<5
685*% %&,%8<%
68%r(   