U
    ¯Ê¹j%  ã                   @   sš   d Z ddlZddlZddlmZ ddlmZ ddlmZm	Z	m
Z
mZmZ ddlmZ eejejejee	e
eeed�	Zdd	„ e ¡ D ƒZd
d„ Zdd„ ZdS )z Norm Layer Factory

Create norm modules by string (to mirror create_act and creat_norm-act fns)

Copyright 2022 Ross Wightman
é    N)ÚTypeé   )Ú	GroupNormÚ
GroupNorm1Ú	LayerNormÚLayerNorm2dÚRmsNorm)ÚFrozenBatchNorm2d)	Z	batchnormZbatchnorm2dZbatchnorm1dZ	groupnormZ
groupnorm1Z	layernormZlayernorm2dZrmsnormZfrozenbatchnorm2dc                 C   s   h | ]\}}|’qS © r
   )Ú.0ÚnÚmr
   r
   úV/var/www/html/bwave/python/venv/lib/python3.8/site-packages/timm/layers/create_norm.pyÚ	<setcomp>   s     r   c                 K   s   t | ƒ}||f|Ž}|S )N)Úget_norm_layer)Ú
layer_nameZnum_featuresÚkwargsÚlayerZlayer_instancer
   r
   r   Úcreate_norm_layer   s    r   c                 C   sŠ   | d krd S t | tttjtjfƒs&t‚i }t | tjƒrH| | j	¡ | j
} t | tƒrp| sZd S |  dd¡}t| } n| } |r†tj| f|Ž} | S )NÚ_Ú )Ú
isinstanceÚtypeÚstrÚtypesÚFunctionTypeÚ	functoolsÚpartialÚAssertionErrorÚupdateÚkeywordsÚfuncÚreplaceÚ	_NORM_MAP)Z
norm_layerZnorm_kwargsr   r
   r
   r   r   $   s     

r   )Ú__doc__r   r   Útypingr   Ztorch.nnÚnnZnormr   r   r   r   r   Ztorchvision.ops.miscr	   ÚdictZBatchNorm2dZBatchNorm1dr#   ÚitemsZ_NORM_TYPESr   r   r
   r
   r
   r   Ú<module>   s(   ÷