




'''simple docstring'''





import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       RemBertConfig,
       RemBertForCausalLM,
       RemBertForMaskedLM,
       RemBertForMultipleChoice,
       RemBertForQuestionAnswering,
       RemBertForSequenceClassification,
       RemBertForTokenClassification,
       RemBertModel,
   )
   from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST

class      _a    :






     def __init__(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE=13      ,_SCREAMING_SNAKE_CASE=7      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=99      ,_SCREAMING_SNAKE_CASE=32      ,_SCREAMING_SNAKE_CASE=18      ,_SCREAMING_SNAKE_CASE=43      ,_SCREAMING_SNAKE_CASE=5      ,_SCREAMING_SNAKE_CASE=4      ,_SCREAMING_SNAKE_CASE=37      ,_SCREAMING_SNAKE_CASE="gelu"      ,_SCREAMING_SNAKE_CASE=0.1      ,_SCREAMING_SNAKE_CASE=0.1      ,_SCREAMING_SNAKE_CASE=512      ,_SCREAMING_SNAKE_CASE=16      ,_SCREAMING_SNAKE_CASE=2      ,_SCREAMING_SNAKE_CASE=0.0_2      ,_SCREAMING_SNAKE_CASE=3      ,_SCREAMING_SNAKE_CASE=4      ,_SCREAMING_SNAKE_CASE=None      ,)							->    Any:
          _snake_case							=   parent
          _snake_case							=   batch_size
          _snake_case							=   seq_length
          _snake_case							=   is_training
          _snake_case							=   use_input_mask
          _snake_case							=   use_token_type_ids
          _snake_case							=   use_labels
          _snake_case							=   vocab_size
          _snake_case							=   hidden_size
          _snake_case							=   input_embedding_size
          _snake_case							=   output_embedding_size
          _snake_case							=   num_hidden_layers
          _snake_case							=   num_attention_heads
          _snake_case							=   intermediate_size
          _snake_case							=   hidden_act
          _snake_case							=   hidden_dropout_prob
          _snake_case							=   attention_probs_dropout_prob
          _snake_case							=   max_position_embeddings
          _snake_case							=   type_vocab_size
          _snake_case							=   type_sequence_label_size
          _snake_case							=   initializer_range
          _snake_case							=   num_labels
          _snake_case							=   num_choices
          _snake_case							=   scope






     def        _lowercase			(				self					)							->    Tuple:
          _snake_case							=   ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size					)

          _snake_case							=   None
          if self.use_input_mask:
               _snake_case							=   random_attention_mask([self.batch_size, self.seq_length]					)

          _snake_case							=   None
          if self.use_token_type_ids:
               _snake_case							=   ids_tensor([self.batch_size, self.seq_length]      ,self.type_vocab_size					)

          _snake_case							=   None
          _snake_case							=   None
          _snake_case							=   None
          if self.use_labels:
               _snake_case							=   ids_tensor([self.batch_size]      ,self.type_sequence_label_size					)
               _snake_case							=   ids_tensor([self.batch_size, self.seq_length]      ,self.num_labels					)
               _snake_case							=   ids_tensor([self.batch_size]      ,self.num_choices					)

          _snake_case							=   RemBertConfig(
              vocab_size=self.vocab_size      ,hidden_size=self.hidden_size      ,input_embedding_size=self.input_embedding_size      ,output_embedding_size=self.output_embedding_size      ,num_hidden_layers=self.num_hidden_layers      ,num_attention_heads=self.num_attention_heads      ,intermediate_size=self.intermediate_size      ,hidden_act=self.hidden_act      ,hidden_dropout_prob=self.hidden_dropout_prob      ,attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,max_position_embeddings=self.max_position_embeddings      ,type_vocab_size=self.type_vocab_size      ,is_decoder=_SCREAMING_SNAKE_CASE      ,initializer_range=self.initializer_range      ,)

          return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






     def        _lowercase			(				self					)							->    List[str]:
          (
              (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,            
          )							=   self.prepare_config_and_inputs()

          _snake_case							=   True
          _snake_case							=   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
          _snake_case							=   ids_tensor([self.batch_size, self.seq_length]      ,vocab_size=2					)

          return (
              config,
              input_ids,
              token_type_ids,
              input_mask,
              sequence_labels,
              token_labels,
              choice_labels,
              encoder_hidden_states,
              encoder_attention_mask,
          )






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    Any:
          _snake_case							=   RemBertModel(config=_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(_SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE					)
          _snake_case							=   model(_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE					)
          _snake_case							=   model(_SCREAMING_SNAKE_CASE					)
          self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, self.seq_length, self.hidden_size)					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,)							->    List[str]:
          _snake_case							=   True
          _snake_case							=   RemBertModel(_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(
              _SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,encoder_hidden_states=_SCREAMING_SNAKE_CASE      ,encoder_attention_mask=_SCREAMING_SNAKE_CASE      ,)
          _snake_case							=   model(
              _SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,encoder_hidden_states=_SCREAMING_SNAKE_CASE      ,)
          _snake_case							=   model(_SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE					)
          self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, self.seq_length, self.hidden_size)					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,)							->    Optional[int]:
          _snake_case							=   RemBertForCausalLM(config=_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(_SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,labels=_SCREAMING_SNAKE_CASE					)
          self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size)					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    Dict:
          _snake_case							=   RemBertForMaskedLM(config=_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(_SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,labels=_SCREAMING_SNAKE_CASE					)
          self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size)					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,)							->    int:
          _snake_case							=   True
          _snake_case							=   True
          _snake_case							=   RemBertForCausalLM(config=_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()

          # first forward pass
          _snake_case							=   model(
              _SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,encoder_hidden_states=_SCREAMING_SNAKE_CASE      ,encoder_attention_mask=_SCREAMING_SNAKE_CASE      ,use_cache=_SCREAMING_SNAKE_CASE      ,)
          _snake_case							=   outputs.past_key_values

          # create hypothetical multiple next token and extent to next_input_ids
          _snake_case							=   ids_tensor((self.batch_size, 3)      ,config.vocab_size					)
          _snake_case							=   ids_tensor((self.batch_size, 3)      ,vocab_size=2					)

          # append to next input_ids and
          _snake_case							=   torch.cat([input_ids, next_tokens]      ,dim=-1					)
          _snake_case							=   torch.cat([input_mask, next_mask]      ,dim=-1					)

          _snake_case							=   model(
              _SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,encoder_hidden_states=_SCREAMING_SNAKE_CASE      ,encoder_attention_mask=_SCREAMING_SNAKE_CASE      ,output_hidden_states=_SCREAMING_SNAKE_CASE      ,)["hidden_states"][0]
          _snake_case							=   model(
              _SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,encoder_hidden_states=_SCREAMING_SNAKE_CASE      ,encoder_attention_mask=_SCREAMING_SNAKE_CASE      ,past_key_values=_SCREAMING_SNAKE_CASE      ,output_hidden_states=_SCREAMING_SNAKE_CASE      ,)["hidden_states"][0]

          # select random slice
          _snake_case							=   ids_tensor((1,)      ,output_from_past.shape[-1]					).item()
          _snake_case							=   output_from_no_past[:, -3:, random_slice_idx].detach()
          _snake_case							=   output_from_past[:, :, random_slice_idx].detach()

          self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]					)

          # test that outputs are equal for slice
          self.parent.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,atol=1e-3					)					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    Any:
          _snake_case							=   RemBertForQuestionAnswering(config=_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(
              _SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,start_positions=_SCREAMING_SNAKE_CASE      ,end_positions=_SCREAMING_SNAKE_CASE      ,)
          self.parent.assertEqual(result.start_logits.shape      ,(self.batch_size, self.seq_length)					)
          self.parent.assertEqual(result.end_logits.shape      ,(self.batch_size, self.seq_length)					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    Tuple:
          _snake_case							=   self.num_labels
          _snake_case							=   RemBertForSequenceClassification(_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(_SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,labels=_SCREAMING_SNAKE_CASE					)
          self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_labels)					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    Dict:
          _snake_case							=   self.num_labels
          _snake_case							=   RemBertForTokenClassification(config=_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   model(_SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,labels=_SCREAMING_SNAKE_CASE					)
          self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.num_labels)					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    int:
          _snake_case							=   self.num_choices
          _snake_case							=   RemBertForMultipleChoice(config=_SCREAMING_SNAKE_CASE					)
          model.to(_SCREAMING_SNAKE_CASE					)
          model.eval()
          _snake_case							=   input_ids.unsqueeze(1					).expand(-1      ,self.num_choices      ,-1					).contiguous()
          _snake_case							=   token_type_ids.unsqueeze(1					).expand(-1      ,self.num_choices      ,-1					).contiguous()
          _snake_case							=   input_mask.unsqueeze(1					).expand(-1      ,self.num_choices      ,-1					).contiguous()
          _snake_case							=   model(
              _SCREAMING_SNAKE_CASE      ,attention_mask=_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,labels=_SCREAMING_SNAKE_CASE      ,)
          self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_choices)					)






     def        _lowercase			(				self					)							->    Any:
          _snake_case							=   self.prepare_config_and_inputs()
          (
              (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,            
          )							=   config_and_inputs
          _snake_case							=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
          return config, inputs_dict







@require_torch
class      _a    (		__lowerCAmelCase      ,      __lowerCAmelCase      ,      unittest.TestCase							):
     SCREAMING_SNAKE_CASE_							:					Dict					       =		(
         (
             RemBertModel,
             RemBertForMaskedLM,
             RemBertForCausalLM,
             RemBertForMultipleChoice,
             RemBertForQuestionAnswering,
             RemBertForSequenceClassification,
             RemBertForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     SCREAMING_SNAKE_CASE_							:					List[Any]					       =		(RemBertForCausalLM,) if is_torch_available() else ()
     SCREAMING_SNAKE_CASE_							:					List[Any]					       =		(
         {
             """feature-extraction""": RemBertModel,
             """fill-mask""": RemBertForMaskedLM,
             """question-answering""": RemBertForQuestionAnswering,
             """text-classification""": RemBertForSequenceClassification,
             """text-generation""": RemBertForCausalLM,
             """token-classification""": RemBertForTokenClassification,
             """zero-shot""": RemBertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )






     def        _lowercase			(				self					)							->    Optional[int]:
          _snake_case							=   RemBertModelTester(self					)
          _snake_case							=   ConfigTester(self      ,config_class=_SCREAMING_SNAKE_CASE      ,hidden_size=37					)






     def        _lowercase			(				self					)							->    Optional[Any]:
          self.config_tester.run_common_tests()






     def        _lowercase			(				self					)							->    Any:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Any:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
               _snake_case							=   type
               self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Tuple:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Union[str, Any]:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_multiple_choice(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Union[str, Any]:
          _snake_case							=   self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_decoder_model_past_large_inputs(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Dict:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Optional[int]:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Tuple:
          _snake_case							=   self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    Dict:
          _snake_case							=   self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_model_as_decoder(*_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self					)							->    List[Any]:
          # This regression test was failing with PyTorch < 1.3
          (
              (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,      (
              _snake_case      
          )							,            
          )							=   self.model_tester.prepare_config_and_inputs_for_decoder()

          _snake_case							=   None

          self.model_tester.create_and_check_model_as_decoder(
              _SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,)






     @slow
     def        _lowercase			(				self					)							->    Tuple:
          for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _snake_case							=   RemBertModel.from_pretrained(_SCREAMING_SNAKE_CASE					)
               self.assertIsNotNone(_SCREAMING_SNAKE_CASE					)







@require_torch
class      _a    (		unittest.TestCase							):






     @slow
     def        _lowercase			(				self					)							->    List[Any]:
          # Test exact values at the last hidden layer
          _snake_case							=   RemBertModel.from_pretrained("google/rembert"					)
          _snake_case							=   torch.tensor([[312, 56_498, 313, 2_125, 313]]					)
          _snake_case							=   torch.tensor([[0, 0, 0, 1, 1]]					)
          with torch.no_grad():
               _snake_case							=   model(_SCREAMING_SNAKE_CASE      ,token_type_ids=_SCREAMING_SNAKE_CASE      ,output_hidden_states=_SCREAMING_SNAKE_CASE					)

          _snake_case							=   1_152

          _snake_case							=   torch.Size((1, 5, hidden_size)					)
          self.assertEqual(output["last_hidden_state"].shape      ,_SCREAMING_SNAKE_CASE					)

          _snake_case							=   torch.tensor(
              [
                  [
                      [0.0_7_5_4, -0.2_0_2_2, 0.1_9_0_4],
                      [-0.3_3_5_4, -0.3_6_9_2, -0.4_7_9_1],
                      [-0.2_3_1_4, -0.6_7_2_9, -0.0_7_4_9],
                      [-0.0_3_9_6, -0.3_1_0_5, -0.4_2_3_4],
                      [-0.1_5_7_1, -0.0_5_2_5, 0.5_3_5_3],
                  ]
              ]					)

          # Running on the original tf implementation gives slightly different results here.
          # Not clear why this variations is present
          # TODO: Find reason for discrepancy
          # expected_original_implementation = [[
          #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
          #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
          #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
          #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
          #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
          # ]]

          self.assertTrue(torch.allclose(output["last_hidden_state"][:, :, :3]      ,_SCREAMING_SNAKE_CASE      ,atol=1e-4					)					)








'''simple docstring'''





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

class      _a       (		enum.Enum							):
     SCREAMING_SNAKE_CASE_							:					Optional[int]					       =		"""timm"""
     SCREAMING_SNAKE_CASE_							:					int					       =		"""transformers"""







def 			__a			(     _UpperCamelCase:					Optional[Iterable[str]]      ,		_UpperCamelCase:					Optional[Iterable[int]]      ,		_UpperCamelCase:					Optional[Iterable[str]]					)		->				Union[str, Any]:

     """simple docstring"""




     if stage_names is None:
          raise ValueError("Stage_names must be set for transformers backbones"					)

     if out_features is not None:
          if not isinstance(_UpperCamelCase      ,		(list,)					):
               raise ValueError(F"""out_features must be a list {type(_UpperCamelCase					)}"""					)
          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(_UpperCamelCase      ,		(list, tuple)					):
               raise ValueError(F"""out_indices must be a list or tuple, got {type(_UpperCamelCase					)}"""					)
          if any(idx >= len(_UpperCamelCase					) 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(_UpperCamelCase					) != len(_UpperCamelCase					):
               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 			__a			(     _UpperCamelCase:					Optional[List[str]]      ,		_UpperCamelCase:					Optional[Union[List[int], Tuple[int]]]      ,		_UpperCamelCase:					List[str]      ,		)		->				Any:

     """simple docstring"""




     if out_indices is None and out_features is None:
          _snake_case							=   [len(_UpperCamelCase					) - 1]
          _snake_case							=   [stage_names[-1]]
     elif out_indices is None and out_features is not None:
          _snake_case							=   [stage_names.index(_UpperCamelCase					) for layer in out_features]
     elif out_features is None and out_indices is not None:
          _snake_case							=   [stage_names[idx] for idx in out_indices]
     return out_features, out_indices





def 			__a			(     _UpperCamelCase:					Optional[List[str]]      ,		_UpperCamelCase:					Optional[Union[List[int], Tuple[int]]]      ,		_UpperCamelCase:					List[str]      ,		)		->				Tuple[List[str], List[int]]:

     """simple docstring"""




     verify_out_features_out_indices(out_features=_UpperCamelCase      ,		out_indices=_UpperCamelCase      ,		stage_names=_UpperCamelCase					)
     _snake_case							,      _snake_case							=   _align_output_features_output_indices(
         out_features=_UpperCamelCase      ,		out_indices=_UpperCamelCase      ,		stage_names=_UpperCamelCase					)
     # Verify that the aligned out_features and out_indices are valid
     verify_out_features_out_indices(out_features=_UpperCamelCase      ,		out_indices=_UpperCamelCase      ,		stage_names=_UpperCamelCase					)
     return output_features, output_indices

class      _a       :
     SCREAMING_SNAKE_CASE_							:					Optional[BackboneType]					       =		None






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    None:
          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']
          _snake_case							=   [stage["module"] for stage in self._backbone.feature_info.info]
          _snake_case							=   [stage["num_chs"] for stage in self._backbone.feature_info.info]
          _snake_case							=   self._backbone.feature_info.out_indices
          _snake_case							=   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					)
          _snake_case							,      _snake_case							=   out_features, out_indices






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    None:
          _snake_case							=   getattr(_SCREAMING_SNAKE_CASE      ,"stage_names"					)
          _snake_case							=   getattr(_SCREAMING_SNAKE_CASE      ,"out_features"      ,_SCREAMING_SNAKE_CASE					)
          _snake_case							=   getattr(_SCREAMING_SNAKE_CASE      ,"out_indices"      ,_SCREAMING_SNAKE_CASE					)

          _snake_case							=   stage_names
          _snake_case							,      _snake_case							=   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
          _snake_case							=   None






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    None:
          _snake_case							=   config

          _snake_case							=   getattr(_SCREAMING_SNAKE_CASE      ,"use_timm_backbone"      ,_SCREAMING_SNAKE_CASE					)
          _snake_case							=   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					)							->    int:
          return self._out_features






     @out_features.setter
     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    str:
          _snake_case							,      _snake_case							=   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					)							->    Any:
          return self._out_indices






     @out_indices.setter
     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    int:
          _snake_case							,      _snake_case							=   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[str]:
          # the current backbones will output the number of channels for each stage
          # even if that stage is not in the out_features list.
          return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names					)}






     @property
     def        _lowercase			(				self					)							->    Any:
          return [self.out_feature_channels[name] for name in self.out_features]






     def        _lowercase			(				self      ,*_SCREAMING_SNAKE_CASE      ,**_SCREAMING_SNAKE_CASE					)							->    str:
          _snake_case							=   dict(inspect.signature(self.forward					).parameters					)
          _snake_case							=   {k: v for k, v in kwargs.items() if k in signature}
          return self(*_SCREAMING_SNAKE_CASE      ,**_SCREAMING_SNAKE_CASE					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE = None      ,_SCREAMING_SNAKE_CASE = None      ,_SCREAMING_SNAKE_CASE = None      ,)							->    Any:
          raise NotImplementedError("This method should be implemented by the derived class."					)






     def        _lowercase			(				self					)							->    Tuple:
          _snake_case							=   super().to_dict()
          _snake_case							=   output.pop("_out_features"					)
          _snake_case							=   output.pop("_out_indices"					)
          return output







class      _a       :






     @property
     def        _lowercase			(				self					)							->    List[str]:
          return self._out_features






     @out_features.setter
     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    str:
          _snake_case							,      _snake_case							=   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					)							->    Any:
          return self._out_indices






     @out_indices.setter
     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    Optional[Any]:
          _snake_case							,      _snake_case							=   get_aligned_output_features_output_indices(
              out_features=_SCREAMING_SNAKE_CASE      ,out_indices=_SCREAMING_SNAKE_CASE      ,stage_names=self.stage_names					)






     def        _lowercase			(				self					)							->    List[str]:
          _snake_case							=   super().to_dict()
          _snake_case							=   output.pop("_out_features"					)
          _snake_case							=   output.pop("_out_indices"					)
          return output


