Ë
    ›¾jA  ã                   óx  — d dl mZmZmZmZ ddlmZ ddlmZm	Z	 ddl
mZ ddlmZmZ ddlmZmZ dd	lmZ dd
lmZ ddlmZ  ej.                  d«      	 	 ddddddœdee   dee   dee   dee   dee   dedeeef   fd„«       Zdeeef   dededeeef   fd„Z	 	 ddededeeef   dee   dee   ddfd„Zy)é    )ÚCallableÚOptionalÚTupleÚcasté   )Úregistry)Úhe_normal_initÚ	zero_init)ÚModel)ÚFloats1dÚFloats2d)Ú	get_widthÚpartialé   )Úchain)ÚDropout)Ú	LayerNormzDish.v1NF)Úinit_WÚinit_bÚdropoutÚ	normalizeÚnOÚnIr   r   r   r   Úreturnc          	      ó  — |€t         }|€t        }t        dt        t	        t
        ||«      | |dœd d dœ¬«      }|rt        |t        | ¬«      «      }|�0t        |t        t        t        t        f   t        |«      «      «      }|S )NÚdish)r   r   )ÚWÚb)ÚinitÚdimsÚparams)r   )r	   r
   r   Úforwardr   r   r   r   r   r   r   )r   r   r   r   r   r   Úmodels          ú[/home/spazio/kingdom_calls/work5/transenv/lib/python3.12/site-packages/thinc/layers/dish.pyÚDishr%      s‡   € ð €~ÜˆØ€~ÜˆÜ',ØÜÜ”T˜6 6Ó*Ø˜bÑ!Ø Ñ%ô(€Eñ Ü�eœY¨"Ô-Ó.ˆØÐÜ�eœT¤%¬´(Ð(:Ñ";¼WÀWÓ=MÓNÓOˆØ€Ló    r#   ÚXÚis_trainc                 ó$  ‡ ‡‡— t        t        ‰ j                  d«      «      Št        t        ‰ j                  d«      «      }‰ j                  j                  ‰‰|«      }‰ j                  j                  |«      }dt        dt        fˆˆˆ fd„}||fS )Nr   r   ÚdYr   c                 ó  •— ‰j                   j                  | ‰d¬«      } ‰j                  d| j                  d¬«      «       ‰j                  d‰j                   j	                  | ‰d¬«      «       ‰j                   j	                  | ‰«      S )	NF)Úinplacer   r   )Úaxisr   T)Útrans1)ÚopsÚbackprop_dishÚinc_gradÚsumÚgemm)r*   r   r'   r#   s    €€€r$   Úbackpropzforward.<locals>.backprop1   sl   ø€ Ø�Y‰Y×$Ñ$ R¨°EÐ$Ó:ˆØ�‰�s˜BŸF™F¨˜F›NÔ+Ø�‰�s˜EŸI™IŸN™N¨2¨q¸˜NÓ>Ô?Ø�y‰y�~‰~˜b !Ó$Ð$r&   )r   r   Ú	get_paramr   r/   Úaffiner   )r#   r'   r(   r   ÚY_preactÚYr4   r   s   ``     @r$   r"   r"   )   sx   ú€ ô 	ŒX�u—‘ sÓ+Ó,€AÜŒX�u—‘ sÓ+Ó,€AØ�y‰y×Ñ  1 aÓ(€HØ�	‰	�‰�xÓ €Að%”Xð %¤(÷ %ð ˆhˆ;Ðr&   r8   c           
      ód  — |�|j                  dt        |«      «       |�|j                  dt        |«      «       |j                  d | |j                  |j	                  d«      |j	                  d«      f«      «       |j                  d ||j                  |j	                  d«      f«      «       y )Nr   r   r   r   )Úset_dimr   Ú	set_paramr/   Úget_dim)r   r   r#   r'   r8   s        r$   r   r   :   s‡   € ð 	€}Ø�‰�dœI a›LÔ)Ø€}Ø�‰�dœI a›LÔ)Ø	‡O�O�C™ §	¡	¨E¯M©M¸$Ó,?ÀÇÁÈtÓATÐ+UÓVÔWØ	‡O�O�C™ §	¡	¨E¯M©M¸$Ó,?Ð+AÓBÕCr&   )NN)Útypingr   r   r   r   Úconfigr   Úinitializersr	   r
   r#   r   Útypesr   r   Úutilr   r   r   r   r   Ú	layernormr   ÚlayersÚintÚfloatÚboolr%   r"   r   © r&   r$   ú<module>rH      sV  ðß 2Ó 2å ß 4Ý ß &ß %Ý Ý Ý  ð €‡��ÓàØðð "&Ø!%Ø#ØòØ�‰ðà�‰ðð �XÑð	ð
 �XÑðð �e‰_ðð ðð ˆ8�XÐÑòó ðð6Ø�˜8Ð#Ñ$ðØ)1ðØ=Aðà
ˆ8�XÐÑóð* !Ø ñDØðDàðDð �˜8Ð#Ñ$ðDð �Ñð	Dð
 �ÑðDð 
ôDr&   