






import copy
from typing import Dict

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING
from ..bit import BitConfig


__UpperCAmelCase										=		logging.get_logger(__name__)

__UpperCAmelCase										=		{
    'google/vit-hybrid-base-bit-384': 'https://huggingface.co/vit-hybrid-base-bit-384/resolve/main/config.json',
    # See all ViT hybrid models at https://huggingface.co/models?filter=vit
}



class 						lowerCamelCase   (_snake_case    ):







		'''simple docstring'''



		_snake_case					:					Dict 		=						'''vit-hybrid'''



		def __init__(							self       ,	_UpperCamelCase=None       ,	_UpperCamelCase=7_6_8       ,	_UpperCamelCase=1_2       ,	_UpperCamelCase=1_2       ,	_UpperCamelCase=3_0_7_2       ,	_UpperCamelCase="gelu"       ,	_UpperCamelCase=0.0       ,	_UpperCamelCase=0.0       ,	_UpperCamelCase=0.02       ,	_UpperCamelCase=1E-12       ,	_UpperCamelCase=2_2_4       ,	_UpperCamelCase=1       ,	_UpperCamelCase=3       ,	_UpperCamelCase=[1, 1_0_2_4, 2_4, 2_4]       ,	_UpperCamelCase=True       ,	**_UpperCamelCase       ,	)			->							List[Any]:
							super().__init__(**_UpperCamelCase     )

							if backbone_config is None:
												logger.info('`backbone_config` is `None`. Initializing the config with a `BiT` backbone.'     )
												UpperCAmelCase_					:   Union[str, Any]   =		{
												    'global_padding': 'same',
												    'layer_type': 'bottleneck',
												    'depths': [3, 4, 9],
												    'out_features': ['stage3'],
												    'embedding_dynamic_padding': True,
												}

							if isinstance(_UpperCamelCase       ,	_UpperCamelCase     ):
												if "model_type" in backbone_config:
																	UpperCAmelCase_					:   Tuple   =		CONFIG_MAPPING[backbone_config['model_type']]
												else:
																	logger.info(
																	    '`model_type` is not found in `backbone_config`. Use `Bit` as the backbone configuration class.'     )
																	UpperCAmelCase_					:   Optional[Any]   =		BitConfig
												UpperCAmelCase_					:   str   =		backbone_config_class(**_UpperCamelCase     )

							UpperCAmelCase_					:   Optional[int]   =		backbone_featmap_shape
							UpperCAmelCase_					:   Union[str, Any]   =		backbone_config
							UpperCAmelCase_					:   List[str]   =		hidden_size
							UpperCAmelCase_					:   Tuple   =		num_hidden_layers
							UpperCAmelCase_					:   str   =		num_attention_heads
							UpperCAmelCase_					:   int   =		intermediate_size
							UpperCAmelCase_					:   str   =		hidden_act
							UpperCAmelCase_					:   Optional[int]   =		hidden_dropout_prob
							UpperCAmelCase_					:   Tuple   =		attention_probs_dropout_prob
							UpperCAmelCase_					:   Any   =		initializer_range
							UpperCAmelCase_					:   Union[str, Any]   =		layer_norm_eps
							UpperCAmelCase_					:   int   =		image_size
							UpperCAmelCase_					:   int   =		patch_size
							UpperCAmelCase_					:   str   =		num_channels
							UpperCAmelCase_					:   List[Any]   =		qkv_bias



		def       __UpperCAmelCase    (							self     )			->							Dict[str, any]:
							UpperCAmelCase_					:   str   =		copy.deepcopy(self.__dict__     )
							UpperCAmelCase_					:   Optional[int]   =		self.backbone_config.to_dict()
							UpperCAmelCase_					:   List[str]   =		self.__class__.model_type
							return output














import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
		import transformer_engine.pytorch as te




def   lowercase__   (  __snake_case					:							List[str]  , __snake_case					:							List[Any]=True  , __snake_case					:							str=True  , __snake_case					:							Tuple=True						):


					'''simple docstring'''


					if not is_fpa_available():
										raise ImportError('Using `convert_model` requires transformer_engine to be installed.'						)
					for name, module in model.named_children():
										if isinstance(__snake_case  , nn.Linear						) and to_transformer_engine and _convert_linear:
															# Return early if the linear layer weights are not multiples of 16
															if any(p % 16 != 0 for p in module.weight.shape						):
																				return
															UpperCAmelCase_					:   Optional[int]   =		module.bias is not None
															UpperCAmelCase_					:   Optional[int]   =		te.Linear(
															    module.in_features  , module.out_features  , bias=__snake_case  , params_dtype=module.weight.dtype						)
															UpperCAmelCase_					:   List[str]   =		module.weight.data.clone()
															if has_bias:
																				UpperCAmelCase_					:   List[Any]   =		module.bias.data.clone()

															setattr(__snake_case  , __snake_case  , __snake_case						)
										elif isinstance(__snake_case  , nn.LayerNorm						) and to_transformer_engine and _convert_ln:
															UpperCAmelCase_					:   int   =		te.LayerNorm(module.normalized_shape[0]  , eps=module.eps  , params_dtype=module.weight.dtype						)
															UpperCAmelCase_					:   Any   =		module.weight.data.clone()
															UpperCAmelCase_					:   Any   =		module.bias.data.clone()

															setattr(__snake_case  , __snake_case  , __snake_case						)
										elif isinstance(__snake_case  , te.Linear						) and not to_transformer_engine and _convert_linear:
															UpperCAmelCase_					:   int   =		module.bias is not None
															UpperCAmelCase_					:   Tuple   =		nn.Linear(
															    module.in_features  , module.out_features  , bias=__snake_case  , params_dtype=module.weight.dtype						)
															UpperCAmelCase_					:   str   =		module.weight.data.clone()
															if has_bias:
																				UpperCAmelCase_					:   Dict   =		module.bias.data.clone()

															setattr(__snake_case  , __snake_case  , __snake_case						)
										elif isinstance(__snake_case  , te.LayerNorm						) and not to_transformer_engine and _convert_ln:
															UpperCAmelCase_					:   List[Any]   =		nn.LayerNorm(module.normalized_shape[0]  , eps=module.eps  , params_dtype=module.weight.dtype						)
															UpperCAmelCase_					:   List[str]   =		module.weight.data.clone()
															UpperCAmelCase_					:   Tuple   =		module.bias.data.clone()

															setattr(__snake_case  , __snake_case  , __snake_case						)
										else:
															convert_model(
															    __snake_case  , to_transformer_engine=__snake_case  , _convert_linear=_convert_linear  , _convert_ln=_convert_ln  , )






def   lowercase__   (  __snake_case					:							List[Any]						):


					'''simple docstring'''


					if not is_fpa_available():
										raise ImportError('Using `has_transformer_engine_layers` requires transformer_engine to be installed.'						)
					for m in model.modules():
										if isinstance(__snake_case  , (te.LayerNorm, te.Linear)						):
															return True
					return False






