'''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    lowerCamelCase__  (   __snake_case       ,				__snake_case       ,				__snake_case     ):



   @register_to_config
   def __init__(    self					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__					,      lowerCAmelCase__ = False					,      )				->     Dict:
        """simple docstring"""





        super().__init__()

        _UpperCamelCase  :List[Any]       =nn.Linear(lowerCAmelCase__					,      lowerCAmelCase__					,      bias=lowerCAmelCase__     )

        _UpperCamelCase  :List[str]       =nn.Embedding(lowerCAmelCase__					,      lowerCAmelCase__     )
        _UpperCamelCase  :Any       =False

        _UpperCamelCase  :Optional[Any]       =nn.Dropout(p=lowerCAmelCase__     )

        _UpperCamelCase  :List[str]       =TaConfig(
            d_model=lowerCAmelCase__					,      num_heads=lowerCAmelCase__					,      d_kv=lowerCAmelCase__					,      d_ff=lowerCAmelCase__					,      feed_forward_proj=lowerCAmelCase__					,      dropout_rate=lowerCAmelCase__					,      is_decoder=lowerCAmelCase__					,      is_encoder_decoder=lowerCAmelCase__					,      )
        _UpperCamelCase  :Dict       =nn.ModuleList()
        for lyr_num in range(lowerCAmelCase__     ):
             _UpperCamelCase  :List[str]       =TaBlock(lowerCAmelCase__     )
             self.encoders.append(lowerCAmelCase__     )

        _UpperCamelCase  :str       =TaLayerNorm(lowerCAmelCase__     )
        _UpperCamelCase  :List[str]       =nn.Dropout(p=lowerCAmelCase__     )




   def  _UpperCamelCase		(    self					,      lowerCAmelCase__					,      lowerCAmelCase__     )				->     Any:
        """simple docstring"""





        _UpperCamelCase  :str       =self.input_proj(lowerCAmelCase__     )

        # terminal relative positional encodings
        _UpperCamelCase  :List[str]       =encoder_inputs.shape[1]
        _UpperCamelCase  :Union[str, Any]       =torch.arange(lowerCAmelCase__					,      device=encoder_inputs.device     )

        _UpperCamelCase  :Dict       =encoder_inputs_mask.sum(-1     )
        _UpperCamelCase  :Optional[Any]       =torch.roll(input_positions.unsqueeze(0     )					,      tuple(seq_lens.tolist()     )					,      dims=0     )
        x += self.position_encoding(lowerCAmelCase__     )

        _UpperCamelCase  :List[str]       =self.dropout_pre(lowerCAmelCase__     )

        # inverted the attention mask
        _UpperCamelCase  :Dict       =encoder_inputs.size()
        _UpperCamelCase  :Any       =self.get_extended_attention_mask(lowerCAmelCase__					,      lowerCAmelCase__     )

        for lyr in self.encoders:
             _UpperCamelCase  :Union[str, Any]       =lyr(lowerCAmelCase__					,      lowerCAmelCase__     )[0]
        _UpperCamelCase  :str       =self.layer_norm(lowerCAmelCase__     )

        return self.dropout_post(lowerCAmelCase__     ), encoder_inputs_mask
'''simple docstring'''

import copy
import inspect

from transformers.testing_utils import require_torch, torch_device
from transformers.utils.backbone_utils import BackboneType




@require_torch
class    lowerCamelCase__  :
   __UpperCAmelCase           =						()
   __UpperCAmelCase           =						True



   def  _UpperCamelCase		(    self     )				->     Tuple:
        """simple docstring"""





        _UpperCamelCase  :List[str]       =self.config_class

        # test default config
        _UpperCamelCase  :List[Any]       =config_class()
        self.assertIsNotNone(lowerCAmelCase__     )
        _UpperCamelCase  :List[str]       =["""stem"""] + [f'''stage{idx}''' for idx in range(1					,      len(config.depths     ) + 1     )]
        self.assertEqual(config.stage_names					,      lowerCAmelCase__     )
        self.assertTrue(set(config.out_features     ).issubset(set(config.stage_names     )     )     )

        # Test out_features and out_indices are correctly set
        # out_features and out_indices both None
        _UpperCamelCase  :Any       =config_class(out_features=lowerCAmelCase__					,      out_indices=lowerCAmelCase__     )
        self.assertEqual(config.out_features					,      [config.stage_names[-1]]     )
        self.assertEqual(config.out_indices					,      [len(config.stage_names     ) - 1]     )

        # out_features and out_indices both set
        _UpperCamelCase  :int       =config_class(out_features=["""stem""", """stage1"""]					,      out_indices=[0, 1]     )
        self.assertEqual(config.out_features					,      ["""stem""", """stage1"""]     )
        self.assertEqual(config.out_indices					,      [0, 1]     )

        # Only out_features set
        _UpperCamelCase  :Dict       =config_class(out_features=["""stage1""", """stage3"""]     )
        self.assertEqual(config.out_features					,      ["""stage1""", """stage3"""]     )
        self.assertEqual(config.out_indices					,      [1, 3]     )

        # Only out_indices set
        _UpperCamelCase  :Tuple       =config_class(out_indices=[0, 2]     )
        self.assertEqual(config.out_features					,      [config.stage_names[0], config.stage_names[2]]     )
        self.assertEqual(config.out_indices					,      [0, 2]     )

        # Error raised when out_indices do not correspond to out_features
        with self.assertRaises(lowerCAmelCase__     ):
             _UpperCamelCase  :List[Any]       =config_class(out_features=["""stage1""", """stage2"""]					,      out_indices=[0, 2]     )



   def  _UpperCamelCase		(    self     )				->     List[Any]:
        """simple docstring"""





        _UpperCamelCase       ,      _UpperCamelCase  :Optional[Any]       =self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             _UpperCamelCase  :Union[str, Any]       =model_class(lowerCAmelCase__     )
             _UpperCamelCase  :Union[str, Any]       =inspect.signature(model.forward     )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             _UpperCamelCase  :Tuple       =[*signature.parameters.keys()]
             _UpperCamelCase  :Any       =["""pixel_values"""]
             self.assertListEqual(arg_names[:1]					,      lowerCAmelCase__     )



   def  _UpperCamelCase		(    self     )				->     int:
        """simple docstring"""





        _UpperCamelCase       ,      _UpperCamelCase  :List[str]       =self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             _UpperCamelCase  :int       =model_class(lowerCAmelCase__     )
             self.assertEqual(len(model.channels     )					,      len(config.out_features     )     )
             _UpperCamelCase  :List[str]       =model.num_features
             _UpperCamelCase  :int       =[config.stage_names.index(lowerCAmelCase__     ) for feat in config.out_features]
             _UpperCamelCase  :List[str]       =[num_features[idx] for idx in out_indices]
             self.assertListEqual(model.channels					,      lowerCAmelCase__     )

             _UpperCamelCase  :Any       =copy.deepcopy(lowerCAmelCase__     )
             _UpperCamelCase  :List[Any]       =None
             _UpperCamelCase  :Optional[Any]       =model_class(lowerCAmelCase__     )
             self.assertEqual(len(model.channels     )					,      1     )
             self.assertListEqual(model.channels					,      [num_features[-1]]     )

             _UpperCamelCase  :Tuple       =copy.deepcopy(lowerCAmelCase__     )
             _UpperCamelCase  :int       =None
             _UpperCamelCase  :int       =model_class(lowerCAmelCase__     )
             self.assertEqual(len(model.channels     )					,      1     )
             self.assertListEqual(model.channels					,      [num_features[-1]]     )



   def  _UpperCamelCase		(    self     )				->     str:
        """simple docstring"""





        _UpperCamelCase       ,      _UpperCamelCase  :Any       =self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             _UpperCamelCase  :List[Any]       =model_class(lowerCAmelCase__     )
             model.to(lowerCAmelCase__     )
             model.eval()
             _UpperCamelCase  :str       =model(**lowerCAmelCase__     )

             self.assertEqual(len(result.feature_maps     )					,      len(config.out_features     )     )
             self.assertEqual(len(model.channels     )					,      len(config.out_features     )     )
             self.assertEqual(len(result.feature_maps     )					,      len(config.out_indices     )     )
             self.assertEqual(len(model.channels     )					,      len(config.out_indices     )     )

             # Check output of last stage is taken if out_features=None, out_indices=None
             _UpperCamelCase  :Any       =copy.deepcopy(lowerCAmelCase__     )
             _UpperCamelCase  :List[Any]       =None
             _UpperCamelCase  :Dict       =model_class(lowerCAmelCase__     )
             model.to(lowerCAmelCase__     )
             model.eval()
             _UpperCamelCase  :str       =model(**lowerCAmelCase__     )

             self.assertEqual(len(result.feature_maps     )					,      1     )
             self.assertEqual(len(model.channels     )					,      1     )

             _UpperCamelCase  :str       =copy.deepcopy(lowerCAmelCase__     )
             _UpperCamelCase  :Dict       =None
             _UpperCamelCase  :Tuple       =model_class(lowerCAmelCase__     )
             model.to(lowerCAmelCase__     )
             model.eval()
             _UpperCamelCase  :Dict       =model(**lowerCAmelCase__     )

             self.assertEqual(len(result.feature_maps     )					,      1     )
             self.assertEqual(len(model.channels     )					,      1     )

             # Check backbone can be initialized with fresh weights
             _UpperCamelCase  :Any       =copy.deepcopy(lowerCAmelCase__     )
             _UpperCamelCase  :Optional[int]       =False
             _UpperCamelCase  :Any       =model_class(lowerCAmelCase__     )
             model.to(lowerCAmelCase__     )
             model.eval()
             _UpperCamelCase  :Optional[int]       =model(**lowerCAmelCase__     )



   def  _UpperCamelCase		(    self     )				->     Optional[int]:
        """simple docstring"""





        _UpperCamelCase       ,      _UpperCamelCase  :Dict       =self.model_tester.prepare_config_and_inputs_for_common()

        for backbone_class in self.all_model_classes:
             _UpperCamelCase  :Optional[int]       =backbone_class(lowerCAmelCase__     )

             self.assertTrue(hasattr(lowerCAmelCase__					,      """backbone_type"""     )     )
             self.assertTrue(hasattr(lowerCAmelCase__					,      """stage_names"""     )     )
             self.assertTrue(hasattr(lowerCAmelCase__					,      """num_features"""     )     )
             self.assertTrue(hasattr(lowerCAmelCase__					,      """out_indices"""     )     )
             self.assertTrue(hasattr(lowerCAmelCase__					,      """out_features"""     )     )
             self.assertTrue(hasattr(lowerCAmelCase__					,      """out_feature_channels"""     )     )
             self.assertTrue(hasattr(lowerCAmelCase__					,      """channels"""     )     )

             self.assertIsInstance(backbone.backbone_type					,      lowerCAmelCase__     )
             # Verify num_features has been initialized in the backbone init
             self.assertIsNotNone(backbone.num_features     )
             self.assertTrue(len(backbone.channels     ) == len(backbone.out_indices     )     )
             self.assertTrue(len(backbone.stage_names     ) == len(backbone.num_features     )     )
             self.assertTrue(len(backbone.channels     ) <= len(backbone.num_features     )     )
             self.assertTrue(len(backbone.out_feature_channels     ) == len(backbone.stage_names     )     )




   def  _UpperCamelCase		(    self     )				->     Tuple:
        """simple docstring"""





        _UpperCamelCase       ,      _UpperCamelCase  :List[Any]       =self.model_tester.prepare_config_and_inputs_for_common()
        _UpperCamelCase  :Union[str, Any]       =inputs_dict["""pixel_values"""].shape[0]

        for backbone_class in self.all_model_classes:
             _UpperCamelCase  :str       =backbone_class(lowerCAmelCase__     )
             backbone.to(lowerCAmelCase__     )
             backbone.eval()

             _UpperCamelCase  :Any       =backbone(**lowerCAmelCase__     )

             # Test default outputs and verify feature maps
             self.assertIsInstance(outputs.feature_maps					,      lowerCAmelCase__     )
             self.assertTrue(len(outputs.feature_maps     ) == len(backbone.channels     )     )
             for feature_map, n_channels in zip(outputs.feature_maps					,      backbone.channels     ):
                  self.assertTrue(feature_map.shape[:2]					,      (batch_size, n_channels)     )
             self.assertIsNone(outputs.hidden_states     )
             self.assertIsNone(outputs.attentions     )

             # Test output_hidden_states=True
             _UpperCamelCase  :Any       =backbone(**lowerCAmelCase__					,      output_hidden_states=lowerCAmelCase__     )
             self.assertIsNotNone(outputs.hidden_states     )
             self.assertTrue(len(outputs.hidden_states     )					,      len(backbone.stage_names     )     )
             for hidden_state, n_channels in zip(outputs.hidden_states					,      backbone.channels     ):
                  self.assertTrue(hidden_state.shape[:2]					,      (batch_size, n_channels)     )

             # Test output_attentions=True
             if self.has_attentions:
                  _UpperCamelCase  :int       =backbone(**lowerCAmelCase__					,      output_attentions=lowerCAmelCase__     )
                  self.assertIsNotNone(outputs.attentions     )