Ë
    éÓ^jc  ã                   ó¨  — d Z ddlmZ ddlmZ ddlmZmZmZm	Z	 ddl
Z
ddlmZ ddlmc mZ ddlmZ ddlmZ  G d	„ d
ej&                  «      Zdededej&                  f   fd„Zdedej&                  f   fd„Zdedej&                  f   defd„Zdddej2                  ddfdedee   dededeeej&                        deej&                     fd„Zy)z¹
BlurPool layer inspired by
 - Kornia's Max_BlurPool2d
 - Making Convolutional Networks Shift-Invariant Again :cite:`zhang2019shiftinvar`

Hacked together by Chris Ha and Ross Wightman
é    )Úpartial)Úcomb)ÚCallableÚOptionalÚTypeÚUnionNé   )Úget_padding)Ú	LayerTypec                   óœ   ‡ — e Zd ZdZ	 	 	 	 	 	 ddee   dedededdf
ˆ fd„Zdd	„Zdd
„Z	de
j                  de
j                  fd„Zdd„Zˆ xZS )Ú
BlurPool2daÖ  Creates a module that computes blurs and downsample a given feature map.
    See :cite:`zhang2019shiftinvar` for more details.
    Corresponds to the Downsample class, which does blurring and subsampling

    Args:
        channels = Number of input channels
        filt_size (int): binomial filter size for blurring. currently supports 3 (default) and 5.
        stride (int): downsampling filter stride

    Returns:
        torch.Tensor: the transformed tensor.
    NÚchannelsÚ	filt_sizeÚstrideÚpad_modeÚreturnc                 ó  •— t         ‰| �  «        |dkD  sJ ‚|| _        || _        || _        || _        t        ||d¬«      gdz  | _        |xs dd||f}| j                  dt        j                  |||¬«      d¬«       | j                  «        y )Nr	   )Údilationé   Úfilt©ÚdeviceÚdtypeF)Ú
persistent)ÚsuperÚ__init__r   r   r   r   r
   ÚpaddingÚregister_bufferÚtorchÚemptyÚreset_parameters)	Úselfr   r   r   r   r   r   Ú
filt_shapeÚ	__class__s	           €ú`/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/timm/layers/blur_pool.pyr   zBlurPool2d.__init__!   s–   ø€ ô 	‰ÑÔØ˜1Š}Ðˆ}Ø ˆŒØ"ˆŒØˆŒØ ˆŒÜ# I¨vÀÔBÐCÀaÑGˆŒð ’m ! Q¨	°9Ð=ˆ
Ø×Ñ˜V¤U§[¡[°ÀFÐRWÔ%XÐejÐÔkð 	×ÑÕó    c                 ó$   — | j                  «        y)zInitialize buffers.N©Ú_init_buffers©r"   s    r%   r!   zBlurPool2d.reset_parameters9   ó   € à×ÑÕr&   c           	      ó¶  — t        j                  t        | j                  «      D �cg c]  }t	        | j                  dz
  |«      ‘Œ c}dt         j
                  ¬«      d| j                  dz
  z  z  }|dd…df   |ddd…f   z  dddd…dd…f   }| j                  �|j                  | j                  ddd«      }| j                  j                  |«       yc c}w )z.Compute and fill non-persistent buffer values.r	   Úcpur   é   N)
r   ÚtensorÚranger   r   Úfloat32r   Úrepeatr   Úcopy_)r"   ÚkÚcoeffsÚblur_filters       r%   r)   zBlurPool2d._init_buffers=   s¿   € ô —‘Ü27¸¿¹Ó2GÖH¨QŒT�$—.‘. 1Ñ$ aÕ(ÒHØÜ—-‘-ô
ð �4—>‘> AÑ%Ñ&ñ	(ˆð
 ša ˜g‘¨°²a°©Ñ8¸$ÀÂaÊÐ:JÑKˆØ�=‰=Ð$Ø%×,Ñ,¨T¯]©]¸A¸qÀ!ÓDˆKØ�	‰	�‰˜Õ$ùò Is   § CÚxc                 ól  — t        j                  || j                  | j                  ¬«      }| j                  €B|j
                  d   }| j                  j                  |d| j                  | j                  «      }n| j                  }| j                  }t        j                  ||| j                  |¬«      S )N)Úmoder	   )r   Úgroups)ÚFÚpadr   r   r   Úshaper   Úexpandr   Úconv2dr   )r"   r7   r   Úweights       r%   ÚforwardzBlurPool2d.forwardJ   s€   € Ü�E‰E�!�T—\‘\¨¯©Ô6ˆØ�=‰=Ð Ø—w‘w˜q‘zˆHØ—Y‘Y×%Ñ% h°°4·>±>À4Ç>Á>ÓR‰Fà—}‘}ˆHØ—Y‘YˆFÜ�x‰x˜˜6¨$¯+©+¸hÔGÐGr&   c                 ó$   — | j                  «        y)z"Initialize non-persistent buffers.Nr(   r*   s    r%   Úinit_non_persistent_buffersz&BlurPool2d.init_non_persistent_buffersT   r+   r&   )Né   r.   ÚreflectNN)r   N)Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   ÚintÚstrr   r!   r)   r   ÚTensorrA   rC   Ú__classcell__)r$   s   @r%   r   r      s‚   ø„ ñð '+ØØØ%ØØñ à˜s‘mð ð ð ð ð	 ð
 ð ð 
õ ó0ó%ðH˜Ÿ™ð H¨%¯,©,ó H÷r&   r   Úaa_layerr   .c                 ó  — t        | t        «      rs| j                  «       j                  dd«      j                  dd«      }|dv rt        j
                  S |dv rt        S |dk(  rt        t        d¬«      S t        d	| › d
�«      ‚| S )z6Map string shorthands to callables (class or partial).Ú_Ú ú-)ÚavgÚavgpool)ÚblurÚblurpoolÚblurpcÚconstant)r   zUnknown anti-aliasing layer (z).)	Ú
isinstancerK   ÚlowerÚreplaceÚnnÚ	AvgPool2dr   r   ÚAssertionError)rN   Úkeys     r%   Ú_normalize_aa_layerr`   Y   s‚   € ä�(œCÔ Ø�n‰nÓ×&Ñ& s¨BÓ/×7Ñ7¸¸RÓ@ˆØÐ$Ñ$Ü—<‘<ÐØÐ&Ñ&ÜÐØ�(Š?äœ:°
Ô;Ð;ÜÐ<¸X¸JÀbÐIÓJÐJØ€Or&   Úlayer_callablec                 ób   — t        | t        «      r| j                  S t        | t        «      r| S dS )z?Return the class behind a callable (unwrap partial), else None.N)rY   r   ÚfuncÚtype)ra   s    r%   Ú_underlying_clsre   h   s,   € ä�.¤'Ô*Ø×"Ñ"Ð"Ü'¨¼Ô=ˆ>ÐGÀ4ÐGr&   c                 ól   — t        | «      }	 t        |t        «      S # t        $ r Y yt        $ r Y yw xY w)z2True if callable is BlurPool2d or a partial of it.F)re   Ú
issubclassr   Ú	TypeErrorÚ	Exception)ra   Úclss     r%   Ú_is_blurpoolrk   o   s;   € ä
˜.Ó
)€CðÜ˜#œzÓ*Ð*øÜò ÙÜò Ùðús   � �	3¨3²3r.   Tr   r   ÚenableÚnoopc                 ó  — | r|s|� |«       S dS t        | «      } ||dœ}t        | «      r@t        | t        «      r| j                  r| j                  ni }d|vr|�||d<   d|vr|�||d<   	  | di |¤ŽS # t
        $ r  | |«      cY S w xY w)zE Anti-aliasing factory that supports strings, classes, and partials. N)r   r   r   r   © )r`   rk   rY   r   Úkeywordsrh   )	rN   r   r   rl   rm   r   r   Úcall_kwargsÚexisting_kws	            r%   Ú	create_aars   z   s·   € ñ ™6ØÐ)‰t‹vÐ3¨tÐ3ô # 8Ó,€Hð  (°6Ñ:€Kô �HÔä+5°hÄÔ+HÈX×M^ÒM^�h×'Ò'ÐdfˆØ˜;Ñ&¨6Ð+=Ø$*ˆK˜Ñ!Ø˜+Ñ%¨%Ð*;Ø#(ˆK˜Ñ ð ÙÑ&˜+Ñ&Ð&øÜò  á˜ÓÒð ús   Á,A4 Á4BÂB)rI   Ú	functoolsr   Úmathr   Útypingr   r   r   r   r   Útorch.nnr\   Útorch.nn.functionalÚ
functionalr;   r   r
   r   ÚModuler   r`   re   Úboolrk   ÚIdentityrJ   rs   ro   r&   r%   ú<module>r}      s  ðñõ Ý ß 2Ó 2ã Ý ß Ð å  Ý ôB�—‘ô BðJ )ð °¸¸b¿i¹i¸Ñ0Hó ðH H¨S°"·)±)¨^Ñ$<ó Hð ¨#¨r¯y©y¨.Ñ!9ð ¸dó ð #'ØØØ*,¯+©+ØØñ! Øð! à˜3‘-ð! ð ð! ð ð	! ð
 �t˜BŸI™I‘Ñ'ð! ð ˆb�i‰iÑô! r&   