



'''simple docstring'''







import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union




class 				__lowerCAmelCase(							enum.Enum      ):
			__snake_case  :		Any		=      'timm'
			__snake_case  :		str		=      'transformers'







def  SCREAMING_SNAKE_CASE__ (				SCREAMING_SNAKE_CASE ,		SCREAMING_SNAKE_CASE ,		SCREAMING_SNAKE_CASE     ):
		if stage_names is None:
				raise ValueError('Stage_names must be set for transformers backbones'     )

		if out_features is not None:
				if not isinstance(SCREAMING_SNAKE_CASE ,		(list,)     ):
						raise ValueError(F'out_features must be a list {type(SCREAMING_SNAKE_CASE     )}'     )
				if any(feat not in stage_names for feat in out_features     ):
						raise ValueError(F'out_features must be a subset of stage_names: {stage_names} got {out_features}'     )

		if out_indices is not None:
				if not isinstance(SCREAMING_SNAKE_CASE ,		(list, tuple)     ):
						raise ValueError(F'out_indices must be a list or tuple, got {type(SCREAMING_SNAKE_CASE     )}'     )
				if any(idx >= len(SCREAMING_SNAKE_CASE     ) for idx in out_indices     ):
						raise ValueError('out_indices must be valid indices for stage_names {stage_names}, got {out_indices}'     )

		if out_features is not None and out_indices is not None:
				if len(SCREAMING_SNAKE_CASE     ) != len(SCREAMING_SNAKE_CASE     ):
						raise ValueError('out_features and out_indices should have the same length if both are set'     )
				if out_features != [stage_names[idx] for idx in out_indices]:
						raise ValueError('out_features and out_indices should correspond to the same stages if both are set'     )







def  SCREAMING_SNAKE_CASE__ (				SCREAMING_SNAKE_CASE ,		SCREAMING_SNAKE_CASE ,		SCREAMING_SNAKE_CASE ,		):
		if out_indices is None and out_features is None:
				SCREAMING_SNAKE_CASE_						:Dict										=      [len(SCREAMING_SNAKE_CASE     ) - 1]
				SCREAMING_SNAKE_CASE_						:Tuple										=      [stage_names[-1]]
		elif out_indices is None and out_features is not None:
				SCREAMING_SNAKE_CASE_						:Tuple										=      [stage_names.index(SCREAMING_SNAKE_CASE     ) for layer in out_features]
		elif out_features is None and out_indices is not None:
				SCREAMING_SNAKE_CASE_						:List[str]										=      [stage_names[idx] for idx in out_indices]
		return out_features, out_indices







def  SCREAMING_SNAKE_CASE__ (				SCREAMING_SNAKE_CASE ,		SCREAMING_SNAKE_CASE ,		SCREAMING_SNAKE_CASE ,		):
		verify_out_features_out_indices(out_features=SCREAMING_SNAKE_CASE ,		out_indices=SCREAMING_SNAKE_CASE ,		stage_names=SCREAMING_SNAKE_CASE     )
		SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      _align_output_features_output_indices(
		    out_features=SCREAMING_SNAKE_CASE ,		out_indices=SCREAMING_SNAKE_CASE ,		stage_names=SCREAMING_SNAKE_CASE     )
		# Verify that the aligned out_features and out_indices are valid
		verify_out_features_out_indices(out_features=SCREAMING_SNAKE_CASE ,		out_indices=SCREAMING_SNAKE_CASE ,		stage_names=SCREAMING_SNAKE_CASE     )
		return output_features, output_indices




class 				__lowerCAmelCase:
			__snake_case  :		Optional[BackboneType]		=      None






			def 				_lowercase	(   self  :							List[Any]      ,       SCREAMING_SNAKE_CASE  :							Union[str, Any]	):




					"""simple docstring"""







					if getattr(self      ,       '_backbone'      ,       SCREAMING_SNAKE_CASE	) is None:
							raise ValueError('self._backbone must be set before calling _init_timm_backbone'	)

					# These will diagree with the defaults for the transformers models e.g. for resnet50
					# the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4']
					# the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4']
					SCREAMING_SNAKE_CASE_						:Tuple										=      [stage['module'] for stage in self._backbone.feature_info.info]
					SCREAMING_SNAKE_CASE_						:List[Any]										=      [stage['num_chs'] for stage in self._backbone.feature_info.info]
					SCREAMING_SNAKE_CASE_						:Dict										=      self._backbone.feature_info.out_indices
					SCREAMING_SNAKE_CASE_						:Tuple										=      self._backbone.feature_info.module_name()

					# We verify the out indices and out features are valid
					verify_out_features_out_indices(
					    out_features=SCREAMING_SNAKE_CASE      ,       out_indices=SCREAMING_SNAKE_CASE      ,       stage_names=self.stage_names	)
					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:str										=      out_features, out_indices






			def 				_lowercase	(   self  :							Optional[int]      ,       SCREAMING_SNAKE_CASE  :							Tuple	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:List[Any]										=      getattr(SCREAMING_SNAKE_CASE      ,       'stage_names'	)
					SCREAMING_SNAKE_CASE_						:Tuple										=      getattr(SCREAMING_SNAKE_CASE      ,       'out_features'      ,       SCREAMING_SNAKE_CASE	)
					SCREAMING_SNAKE_CASE_						:List[Any]										=      getattr(SCREAMING_SNAKE_CASE      ,       'out_indices'      ,       SCREAMING_SNAKE_CASE	)

					SCREAMING_SNAKE_CASE_						:Tuple										=      stage_names
					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:Tuple										=      get_aligned_output_features_output_indices(
					    out_features=SCREAMING_SNAKE_CASE      ,       out_indices=SCREAMING_SNAKE_CASE      ,       stage_names=SCREAMING_SNAKE_CASE	)
					# Number of channels for each stage. This is set in the transformer backbone model init
					SCREAMING_SNAKE_CASE_						:Dict										=      None






			def 				_lowercase	(   self  :							List[str]      ,       SCREAMING_SNAKE_CASE  :							Tuple	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Dict										=      config

					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      getattr(SCREAMING_SNAKE_CASE      ,       'use_timm_backbone'      ,       SCREAMING_SNAKE_CASE	)
					SCREAMING_SNAKE_CASE_						:Any										=      BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

					if self.backbone_type == BackboneType.TIMM:
							self._init_timm_backbone(SCREAMING_SNAKE_CASE	)
					elif self.backbone_type == BackboneType.TRANSFORMERS:
							self._init_transformers_backbone(SCREAMING_SNAKE_CASE	)
					else:
							raise ValueError(f'backbone_type {self.backbone_type} not supported.'	)






			@property
			def 				_lowercase	(   self  :							Any	):




					"""simple docstring"""







					return self._out_features






			@out_features.setter
			def 				_lowercase	(   self  :							int      ,       SCREAMING_SNAKE_CASE  :							List[str]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:List[Any]										=      get_aligned_output_features_output_indices(
					    out_features=SCREAMING_SNAKE_CASE      ,       out_indices=SCREAMING_SNAKE_CASE      ,       stage_names=self.stage_names	)






			@property
			def 				_lowercase	(   self  :							Tuple	):




					"""simple docstring"""







					return self._out_indices






			@out_indices.setter
			def 				_lowercase	(   self  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							Union[Tuple[int], List[int]]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:str										=      get_aligned_output_features_output_indices(
					    out_features=SCREAMING_SNAKE_CASE      ,       out_indices=SCREAMING_SNAKE_CASE      ,       stage_names=self.stage_names	)






			@property
			def 				_lowercase	(   self  :							Optional[int]	):




					"""simple docstring"""







					return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names	)}






			@property
			def 				_lowercase	(   self  :							Optional[Any]	):




					"""simple docstring"""







					return [self.out_feature_channels[name] for name in self.out_features]






			def 				_lowercase	(   self  :							List[str]      ,       *SCREAMING_SNAKE_CASE  :							Tuple      ,       **SCREAMING_SNAKE_CASE  :							str	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Tuple										=      dict(inspect.signature(self.forward	).parameters	)
					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      {k: v for k, v in kwargs.items() if k in signature}
					return self(*SCREAMING_SNAKE_CASE      ,       **SCREAMING_SNAKE_CASE	)






			def 				_lowercase	(   self  :							Optional[Any]      ,       SCREAMING_SNAKE_CASE  :							Any      ,       SCREAMING_SNAKE_CASE  :							Optional[bool] = None      ,       SCREAMING_SNAKE_CASE  :							Optional[bool] = None      ,       SCREAMING_SNAKE_CASE  :							Optional[bool] = None      ,       ):




					"""simple docstring"""







					raise NotImplementedError('This method should be implemented by the derived class.'	)






			def 				_lowercase	(   self  :							Any	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Dict										=      super().to_dict()
					SCREAMING_SNAKE_CASE_						:List[str]										=      output.pop('_out_features'	)
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      output.pop('_out_indices'	)
					return output




class 				__lowerCAmelCase:






			@property
			def 				_lowercase	(   self  :							Union[str, Any]	):




					"""simple docstring"""







					return self._out_features






			@out_features.setter
			def 				_lowercase	(   self  :							List[str]      ,       SCREAMING_SNAKE_CASE  :							List[str]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:Optional[int]										=      get_aligned_output_features_output_indices(
					    out_features=SCREAMING_SNAKE_CASE      ,       out_indices=SCREAMING_SNAKE_CASE      ,       stage_names=self.stage_names	)






			@property
			def 				_lowercase	(   self  :							List[Any]	):




					"""simple docstring"""







					return self._out_indices






			@out_indices.setter
			def 				_lowercase	(   self  :							Optional[int]      ,       SCREAMING_SNAKE_CASE  :							Union[Tuple[int], List[int]]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:List[Any]										=      get_aligned_output_features_output_indices(
					    out_features=SCREAMING_SNAKE_CASE      ,       out_indices=SCREAMING_SNAKE_CASE      ,       stage_names=self.stage_names	)






			def 				_lowercase	(   self  :							str	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:str										=      super().to_dict()
					SCREAMING_SNAKE_CASE_						:Tuple										=      output.pop('_out_features'	)
					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      output.pop('_out_indices'	)
					return output




"""simple docstring"""







import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import (
    TaBlock,
    TaConfig,
    TaLayerNorm,
)

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin






class  __A							(			_SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE     ):






 """simple docstring"""




 @register_to_config
 def __init__(						self					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A = False					,       )   ->	Optional[Any]:
     super().__init__()

     a											=nn.Linear(__A					,       __A					,       bias=__A		)

     a											=nn.Embedding(__A					,       __A		)
     a											=False

     a											=nn.Dropout(p=__A		)

     a											=TaConfig(
         d_model=__A					,       num_heads=__A					,       d_kv=__A					,       d_ff=__A					,       feed_forward_proj=__A					,       dropout_rate=__A					,       is_decoder=__A					,       is_encoder_decoder=__A					,       )
     a											=nn.ModuleList()
     for lyr_num in range(__A		):
         a											=TaBlock(__A		)
         self.encoders.append(__A		)

     a											=TaLayerNorm(__A		)
     a											=nn.Dropout(p=__A		)





 def  SCREAMING_SNAKE_CASE				(						self					,       __A					,       __A		)   ->	str:
     a											=self.input_proj(__A		)

     # terminal relative positional encodings
     a											=encoder_inputs.shape[1]
     a											=torch.arange(__A					,       device=encoder_inputs.device		)

     a											=encoder_inputs_mask.sum(-1		)
     a											=torch.roll(input_positions.unsqueeze(0		)					,       tuple(seq_lens.tolist()		)					,       dims=0		)
     x += self.position_encoding(__A		)

     a											=self.dropout_pre(__A		)

     # inverted the attention mask
     a											=encoder_inputs.size()
     a											=self.get_extended_attention_mask(__A					,       __A		)

     for lyr in self.encoders:
         a											=lyr(__A					,       __A		)[0]
     a											=self.layer_norm(__A		)

     return self.dropout_post(__A		), encoder_inputs_mask