






import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

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

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


_SCREAMING_SNAKE_CASE   =  1e-4

if is_torch_available():
    import torch

    from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
    from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder






@require_torch
class 					SCREAMING_SNAKE_CASE_  :
       """simple docstring"""


       def __init__(				self	:List[str],      snake_case	:List[str],      snake_case	:List[Any]=16,      snake_case	:str=13,      snake_case	:Tuple=7,      snake_case	:Tuple=14,      snake_case	:List[Any]=10,      snake_case	:Dict=19,      snake_case	:Tuple=5,      snake_case	:int=4,      snake_case	:str=True,      snake_case	:Dict=16,      snake_case	:str=2,      snake_case	:Optional[Any]=4,      snake_case	:Tuple=4,      snake_case	:Tuple="gelu",      snake_case	:List[str]=0.1,      snake_case	:List[Any]=0.1,      snake_case	:Any=[1, 2, 3, 4, 5],      snake_case	:str=25,      snake_case	:str=5,      ):



          """simple docstring"""

          _lowercase              =d_model
          _lowercase              =parent
          _lowercase              =batch_size
          _lowercase              =prediction_length
          _lowercase              =context_length
          _lowercase              =cardinality
          _lowercase              =num_time_features
          _lowercase              =lags_sequence
          _lowercase              =embedding_dimension
          _lowercase              =is_training
          _lowercase              =hidden_size
          _lowercase              =num_hidden_layers
          _lowercase              =num_attention_heads
          _lowercase              =intermediate_size
          _lowercase              =hidden_act
          _lowercase              =hidden_dropout_prob
          _lowercase              =attention_probs_dropout_prob

          _lowercase              =context_length
          _lowercase              =prediction_length + label_length
          _lowercase              =label_length

          _lowercase              =moving_average
          _lowercase              =autocorrelation_factor





       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          return AutoformerConfig(
              d_model=self.d_model,      encoder_layers=self.num_hidden_layers,      decoder_layers=self.num_hidden_layers,      encoder_attention_heads=self.num_attention_heads,      decoder_attention_heads=self.num_attention_heads,      encoder_ffn_dim=self.intermediate_size,      decoder_ffn_dim=self.intermediate_size,      dropout=self.hidden_dropout_prob,      attention_dropout=self.attention_probs_dropout_prob,      prediction_length=self.prediction_length,      context_length=self.context_length,      label_length=self.label_length,      lags_sequence=self.lags_sequence,      num_time_features=self.num_time_features,      num_static_categorical_features=1,      cardinality=[self.cardinality],      embedding_dimension=[self.embedding_dimension],      moving_average=self.moving_average,      )





       def 				UpperCamelCase__						(				self	:List[Any],      snake_case	:Any):



          """simple docstring"""

          _lowercase              =config.context_length + max(config.lags_sequence)

          _lowercase              =ids_tensor([self.batch_size, 1],      config.cardinality[0])
          _lowercase              =floats_tensor([self.batch_size, _past_length, config.num_time_features])
          _lowercase              =floats_tensor([self.batch_size, _past_length])
          _lowercase              =floats_tensor([self.batch_size, _past_length]) > 0.5

          # decoder inputs
          _lowercase              =floats_tensor([self.batch_size, config.prediction_length, config.num_time_features])
          _lowercase              =floats_tensor([self.batch_size, config.prediction_length])

          _lowercase              ={
              'past_values': past_values,
              'static_categorical_features': static_categorical_features,
              'past_time_features': past_time_features,
              'past_observed_mask': past_observed_mask,
              'future_time_features': future_time_features,
              'future_values': future_values,
          }
          return inputs_dict





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          _lowercase              =self.get_config()
          _lowercase              =self.prepare_autoformer_inputs_dict(snake_case)
          return config, inputs_dict





       def 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          _lowercase       ,       _lowercase              =self.prepare_config_and_inputs()
          return config, inputs_dict





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:Dict,      snake_case	:Union[str, Any]):



          """simple docstring"""

          _lowercase              =AutoformerModel(config=snake_case).to(snake_case).eval()
          _lowercase              =model(**snake_case)

          _lowercase              =outputs.encoder_last_hidden_state
          _lowercase              =outputs.last_hidden_state

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =model.get_encoder()
             encoder.save_pretrained(snake_case)
             _lowercase              =AutoformerEncoder.from_pretrained(snake_case).to(snake_case)

          _lowercase       ,       _lowercase       ,       _lowercase       ,       _lowercase       ,       _lowercase              =model.create_network_inputs(**snake_case)
          _lowercase       ,       _lowercase              =model.decomposition_layer(transformer_inputs[:, : config.context_length, ...])

          _lowercase              =torch.cat(
              (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...]),      dim=-1,      )
          _lowercase              =encoder(inputs_embeds=snake_case)[0]
          self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1e-3)

          _lowercase              =(
              torch.mean(transformer_inputs[:, : config.context_length, ...],      dim=1)
              .unsqueeze(1)
              .repeat(1,      config.prediction_length,      1)
          )
          _lowercase              =torch.zeros(
              [transformer_inputs.shape[0], config.prediction_length, transformer_inputs.shape[2]],      device=enc_input.device,      )

          _lowercase              =torch.cat(
              (
                  torch.cat((seasonal_input[:, -config.label_length :, ...], zeros),      dim=1),
                  feature[:, config.context_length - config.label_length :, ...],
              ),      dim=-1,      )
          _lowercase              =torch.cat(
              (
                  torch.cat((trend_input[:, -config.label_length :, ...], mean),      dim=1),
                  feature[:, config.context_length - config.label_length :, ...],
              ),      dim=-1,      )

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =model.get_decoder()
             decoder.save_pretrained(snake_case)
             _lowercase              =AutoformerDecoder.from_pretrained(snake_case).to(snake_case)

          _lowercase              =decoder(
              trend=snake_case,      inputs_embeds=snake_case,      encoder_hidden_states=snake_case,      )[0]

          self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1e-3)






@require_torch
class 					SCREAMING_SNAKE_CASE_  (       _a  ,						_a  ,						unittest.TestCase ):
       """simple docstring"""


       __lowerCAmelCase    :		int					=(AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
       __lowerCAmelCase    :		Tuple					=(AutoformerForPrediction,) if is_torch_available() else ()
       __lowerCAmelCase    :		Tuple					={'''feature-extraction''': AutoformerModel} if is_torch_available() else {}
       __lowerCAmelCase    :		Optional[int]					=False
       __lowerCAmelCase    :		List[str]					=False
       __lowerCAmelCase    :		Optional[Any]					=False
       __lowerCAmelCase    :		Optional[int]					=False
       __lowerCAmelCase    :		str					=False
       __lowerCAmelCase    :		str					=False





       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          _lowercase              =AutoformerModelTester(self)
          _lowercase              =ConfigTester(self,      config_class=snake_case,      has_text_modality=snake_case)





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          self.config_tester.run_common_tests()





       def 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          _lowercase       ,       _lowercase              =self.model_tester.prepare_config_and_inputs()
          for model_class in self.all_model_classes:
             _lowercase              =model_class(snake_case)

             with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(snake_case)
                _lowercase       ,       _lowercase              =model_class.from_pretrained(snake_case,      output_loading_info=snake_case)
             self.assertEqual(info['missing_keys'],      [])





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.check_encoder_decoder_model_standalone(*snake_case)





       @unittest.skip(reason='Model has no tokens embeddings')
       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          pass





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          _lowercase              =inspect.signature(getattr(snake_case,      'forward'))
          # The main input is the name of the argument after `self`
          _lowercase              =list(model_signature.parameters.keys())[1]
          self.assertEqual(AutoformerModel.main_input_name,      snake_case)





       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          _lowercase       ,       _lowercase              =self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             _lowercase              =model_class(snake_case)
             _lowercase              =inspect.signature(model.forward)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             _lowercase              =[*signature.parameters.keys()]

             _lowercase              =[
                 'past_values',
                 'past_time_features',
                 'past_observed_mask',
                 'static_categorical_features',
                 'static_real_features',
                 'future_values',
                 'future_time_features',
             ]

             if model.__class__.__name__ in ["AutoformerForPrediction"]:
                expected_arg_names.append('future_observed_mask')

             expected_arg_names.extend(
                 [
                     'decoder_attention_mask',
                     'head_mask',
                     'decoder_head_mask',
                     'cross_attn_head_mask',
                     'encoder_outputs',
                     'past_key_values',
                     'output_hidden_states',
                     'output_attentions',
                     'use_cache',
                     'return_dict',
                 ])

             self.assertListEqual(arg_names[: len(snake_case)],      snake_case)





       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          _lowercase       ,       _lowercase              =self.model_tester.prepare_config_and_inputs_for_common()
          _lowercase              =True

          _lowercase              =getattr(self.model_tester,      'seq_length',      snake_case)
          _lowercase              =getattr(self.model_tester,      'decoder_seq_length',      snake_case)
          _lowercase              =getattr(self.model_tester,      'encoder_seq_length',      snake_case)
          _lowercase              =getattr(self.model_tester,      'd_model',      snake_case)
          _lowercase              =getattr(self.model_tester,      'num_attention_heads',      snake_case)
          _lowercase              =d_model // num_attention_heads

          for model_class in self.all_model_classes:
             _lowercase              =True
             _lowercase              =False
             _lowercase              =True
             _lowercase              =model_class(snake_case)
             model.to(snake_case)
             model.eval()
             with torch.no_grad():
                _lowercase              =model(**self._prepare_for_class(snake_case,      snake_case))
             _lowercase              =outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
             self.assertEqual(len(snake_case),      self.model_tester.num_hidden_layers)

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             _lowercase              =True
             _lowercase              =model_class(snake_case)
             model.to(snake_case)
             model.eval()
             with torch.no_grad():
                _lowercase              =model(**self._prepare_for_class(snake_case,      snake_case))
             _lowercase              =outputs.encoder_attentions
             self.assertEqual(len(snake_case),      self.model_tester.num_hidden_layers)

             self.assertListEqual(
                 list(attentions[0].shape[-3:]),      [self.model_tester.num_attention_heads, encoder_seq_length, dim],      )
             _lowercase              =len(snake_case)

             _lowercase              =7

             if "last_hidden_state" in outputs:
                correct_outlen += 1

             if "trend" in outputs:
                correct_outlen += 1

             if "past_key_values" in outputs:
                correct_outlen += 1  # past_key_values have been returned

             if "loss" in outputs:
                correct_outlen += 1

             if "params" in outputs:
                correct_outlen += 1

             self.assertEqual(snake_case,      snake_case)

             # decoder attentions
             _lowercase              =outputs.decoder_attentions
             self.assertIsInstance(snake_case,      (list, tuple))
             self.assertEqual(len(snake_case),      self.model_tester.num_hidden_layers)
             self.assertListEqual(
                 list(decoder_attentions[0].shape[-3:]),      [self.model_tester.num_attention_heads, decoder_seq_length, dim],      )

             # cross attentions
             _lowercase              =outputs.cross_attentions
             self.assertIsInstance(snake_case,      (list, tuple))
             self.assertEqual(len(snake_case),      self.model_tester.num_hidden_layers)
             self.assertListEqual(
                 list(cross_attentions[0].shape[-3:]),      [self.model_tester.num_attention_heads, decoder_seq_length, dim],      )

          # Check attention is always last and order is fine
          _lowercase              =True
          _lowercase              =True
          _lowercase              =model_class(snake_case)
          model.to(snake_case)
          model.eval()
          with torch.no_grad():
             _lowercase              =model(**self._prepare_for_class(snake_case,      snake_case))

          self.assertEqual(out_len + 2,      len(snake_case))

          _lowercase              =outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

          self.assertEqual(len(snake_case),      self.model_tester.num_hidden_layers)
          self.assertListEqual(
              list(self_attentions[0].shape[-3:]),      [self.model_tester.num_attention_heads, encoder_seq_length, dim],      )





       @is_flaky()
       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          super().test_retain_grad_hidden_states_attentions()





def        _snake_case   (_snake_case				:       Union[str, Any]="train-batch.pt")					->	str:
   _lowercase              =hf_hub_download(repo_id='hf-internal-testing/tourism-monthly-batch'      ,    filename=_snake_case      ,    repo_type='dataset')
   _lowercase              =torch.load(_snake_case      ,    map_location=_snake_case)
   return batch






@require_torch
@slow
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          _lowercase              =AutoformerModel.from_pretrained('huggingface/autoformer-tourism-monthly').to(snake_case)
          _lowercase              =prepare_batch()

          with torch.no_grad():
             _lowercase              =model(
                 past_values=batch['past_values'],      past_time_features=batch['past_time_features'],      past_observed_mask=batch['past_observed_mask'],      static_categorical_features=batch['static_categorical_features'],      future_values=batch['future_values'],      future_time_features=batch['future_time_features'],      )[0]

          _lowercase              =torch.Size(
              (64, model.config.prediction_length + model.config.label_length, model.config.feature_size))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[0.3_5_9_3, -1.3_3_9_8, 0.6_3_3_0], [0.2_2_7_9, 1.5_3_9_6, -0.1_7_9_2], [0.0_4_5_0, 1.3_2_2_5, -0.2_3_3_5]],      device=snake_case)
          self.assertTrue(torch.allclose(output[0, :3, :3],      snake_case,      atol=snake_case))





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly').to(snake_case)
          _lowercase              =prepare_batch('val-batch.pt')
          with torch.no_grad():
             _lowercase              =model(
                 past_values=batch['past_values'],      past_time_features=batch['past_time_features'],      past_observed_mask=batch['past_observed_mask'],      static_categorical_features=batch['static_categorical_features'],      ).encoder_last_hidden_state
          _lowercase              =torch.Size((64, model.config.context_length, model.config.d_model))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[-0.0_7_3_4, -0.9_0_3_6, 0.8_3_5_8], [4.7_1_8_6, 2.4_1_1_3, 1.9_5_8_1], [1.7_9_5_3, 2.3_5_5_8, 1.2_9_7_0]],      device=snake_case)
          self.assertTrue(torch.allclose(output[0, :3, :3],      snake_case,      atol=snake_case))





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly').to(snake_case)
          _lowercase              =prepare_batch('val-batch.pt')
          with torch.no_grad():
             _lowercase              =model.generate(
                 static_categorical_features=batch['static_categorical_features'],      past_time_features=batch['past_time_features'],      past_values=batch['past_values'],      future_time_features=batch['future_time_features'],      past_observed_mask=batch['past_observed_mask'],      )
          _lowercase              =torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length))
          self.assertEqual(outputs.sequences.shape,      snake_case)

          _lowercase              =torch.tensor([3_1_3_0.6_7_6_3, 4_0_5_6.5_2_9_3, 7_0_5_3.0_7_8_6],      device=snake_case)
          _lowercase              =outputs.sequences.mean(dim=1)
          self.assertTrue(torch.allclose(mean_prediction[0, -3:],      snake_case,      rtol=1e-1))










import unittest

from transformers import MraConfig, 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


if is_torch_available():
    import torch

    from transformers import (
        MraForMaskedLM,
        MraForMultipleChoice,
        MraForQuestionAnswering,
        MraForSequenceClassification,
        MraForTokenClassification,
        MraModel,
    )
    from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST






class 					SCREAMING_SNAKE_CASE_  :
       """simple docstring"""


       def __init__(				self	:str,      snake_case	:Any,      snake_case	:str=2,      snake_case	:int=8,      snake_case	:List[Any]=True,      snake_case	:str=True,      snake_case	:Optional[int]=True,      snake_case	:str=True,      snake_case	:Optional[Any]=99,      snake_case	:Tuple=16,      snake_case	:Optional[Any]=5,      snake_case	:Optional[Any]=2,      snake_case	:Optional[Any]=36,      snake_case	:List[str]="gelu",      snake_case	:str=0.0,      snake_case	:List[Any]=0.0,      snake_case	:List[str]=512,      snake_case	:Dict=16,      snake_case	:int=2,      snake_case	:Tuple=0.0_2,      snake_case	:List[str]=3,      snake_case	:Any=4,      snake_case	:Union[str, Any]=None,      ):



          """simple docstring"""

          _lowercase              =parent
          _lowercase              =batch_size
          _lowercase              =seq_length
          _lowercase              =is_training
          _lowercase              =use_input_mask
          _lowercase              =use_token_type_ids
          _lowercase              =use_labels
          _lowercase              =vocab_size
          _lowercase              =hidden_size
          _lowercase              =num_hidden_layers
          _lowercase              =num_attention_heads
          _lowercase              =intermediate_size
          _lowercase              =hidden_act
          _lowercase              =hidden_dropout_prob
          _lowercase              =attention_probs_dropout_prob
          _lowercase              =max_position_embeddings
          _lowercase              =type_vocab_size
          _lowercase              =type_sequence_label_size
          _lowercase              =initializer_range
          _lowercase              =num_labels
          _lowercase              =num_choices
          _lowercase              =scope





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          _lowercase              =ids_tensor([self.batch_size, self.seq_length],      self.vocab_size)

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

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

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

          _lowercase              =self.get_config()

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





       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          return MraConfig(
              vocab_size=self.vocab_size,      hidden_size=self.hidden_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=snake_case,      initializer_range=self.initializer_range,      )





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =self.get_config()
          _lowercase              =300
          return config





       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          (
              (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       
          )              =self.prepare_config_and_inputs()

          _lowercase              =True
          _lowercase              =floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
          _lowercase              =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 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:Optional[int],      snake_case	:List[str],      snake_case	:int,      snake_case	:Tuple,      snake_case	:Optional[int],      snake_case	:Dict,      snake_case	:str):



          """simple docstring"""

          _lowercase              =MraModel(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case)
          _lowercase              =model(snake_case,      token_type_ids=snake_case)
          _lowercase              =model(snake_case)
          self.parent.assertEqual(result.last_hidden_state.shape,      (self.batch_size, self.seq_length, self.hidden_size))





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:Any,      snake_case	:str,      snake_case	:Optional[int],      snake_case	:int,      snake_case	:int,      snake_case	:int,      snake_case	:Optional[Any],      snake_case	:Tuple,      snake_case	:Tuple,      ):



          """simple docstring"""

          _lowercase              =True
          _lowercase              =MraModel(snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      encoder_hidden_states=snake_case,      encoder_attention_mask=snake_case,      )
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      encoder_hidden_states=snake_case,      )
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case)
          self.parent.assertEqual(result.last_hidden_state.shape,      (self.batch_size, self.seq_length, self.hidden_size))





       def 				UpperCamelCase__						(				self	:str,      snake_case	:Optional[int],      snake_case	:Optional[int],      snake_case	:int,      snake_case	:Optional[int],      snake_case	:List[str],      snake_case	:Optional[int],      snake_case	:Optional[Any]):



          """simple docstring"""

          _lowercase              =MraForMaskedLM(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      labels=snake_case)
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.vocab_size))





       def 				UpperCamelCase__						(				self	:Tuple,      snake_case	:Union[str, Any],      snake_case	:Dict,      snake_case	:Any,      snake_case	:Union[str, Any],      snake_case	:List[Any],      snake_case	:Tuple,      snake_case	:Any):



          """simple docstring"""

          _lowercase              =MraForQuestionAnswering(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      start_positions=snake_case,      end_positions=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 				UpperCamelCase__						(				self	:int,      snake_case	:List[str],      snake_case	:int,      snake_case	:Optional[Any],      snake_case	:Union[str, Any],      snake_case	:str,      snake_case	:List[Any],      snake_case	:Optional[Any]):



          """simple docstring"""

          _lowercase              =self.num_labels
          _lowercase              =MraForSequenceClassification(snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      labels=snake_case)
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.num_labels))





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:Optional[int],      snake_case	:List[Any],      snake_case	:Optional[Any],      snake_case	:int,      snake_case	:int,      snake_case	:Dict,      snake_case	:List[Any]):



          """simple docstring"""

          _lowercase              =self.num_labels
          _lowercase              =MraForTokenClassification(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =model(snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      labels=snake_case)
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.num_labels))





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:List[str],      snake_case	:Optional[int],      snake_case	:Dict,      snake_case	:str,      snake_case	:Optional[Any],      snake_case	:Tuple,      snake_case	:Dict):



          """simple docstring"""

          _lowercase              =self.num_choices
          _lowercase              =MraForMultipleChoice(config=snake_case)
          model.to(snake_case)
          model.eval()
          _lowercase              =input_ids.unsqueeze(1).expand(-1,      self.num_choices,      -1).contiguous()
          _lowercase              =token_type_ids.unsqueeze(1).expand(-1,      self.num_choices,      -1).contiguous()
          _lowercase              =input_mask.unsqueeze(1).expand(-1,      self.num_choices,      -1).contiguous()
          _lowercase              =model(
              snake_case,      attention_mask=snake_case,      token_type_ids=snake_case,      labels=snake_case,      )
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.num_choices))





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =self.prepare_config_and_inputs()
          (
              (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       (
              _lowercase
          )       ,       
          )              =config_and_inputs
          _lowercase              ={'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
          return config, inputs_dict






@require_torch
class 					SCREAMING_SNAKE_CASE_  (       _a  ,						unittest.TestCase ):
       """simple docstring"""


       __lowerCAmelCase    :		Optional[int]					=(
           (
               MraModel,
               MraForMaskedLM,
               MraForMultipleChoice,
               MraForQuestionAnswering,
               MraForSequenceClassification,
               MraForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       __lowerCAmelCase    :		int					=False
       __lowerCAmelCase    :		Union[str, Any]					=False
       __lowerCAmelCase    :		str					=False
       __lowerCAmelCase    :		str					=False

       __lowerCAmelCase    :		Dict					=()





       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          _lowercase              =MraModelTester(self)
          _lowercase              =ConfigTester(self,      config_class=snake_case,      hidden_size=37)





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          self.config_tester.run_common_tests()





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case)





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
             _lowercase              =type
             self.model_tester.create_and_check_model(*snake_case)





       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*snake_case)





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_multiple_choice(*snake_case)





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*snake_case)





       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*snake_case)





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*snake_case)





       @slow
       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _lowercase              =MraModel.from_pretrained(snake_case)
             self.assertIsNotNone(snake_case)





       @unittest.skip(reason='MRA does not output attentions')
       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          return






@require_torch
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       @slow
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          _lowercase              =MraModel.from_pretrained('uw-madison/mra-base-512-4')
          _lowercase              =torch.arange(256).unsqueeze(0)

          with torch.no_grad():
             _lowercase              =model(snake_case)[0]

          _lowercase              =torch.Size((1, 256, 768))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[-0.0_1_4_0, 0.0_8_3_0, -0.0_3_8_1], [0.1_5_4_6, 0.1_4_0_2, 0.0_2_2_0], [0.1_1_6_2, 0.0_8_5_1, 0.0_1_6_5]]])

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





       @slow
       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          _lowercase              =MraForMaskedLM.from_pretrained('uw-madison/mra-base-512-4')
          _lowercase              =torch.arange(256).unsqueeze(0)

          with torch.no_grad():
             _lowercase              =model(snake_case)[0]

          _lowercase              =5_0265

          _lowercase              =torch.Size((1, 256, vocab_size))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[9.2_5_9_5, -3.6_0_3_8, 1_1.8_8_1_9], [9.3_8_6_9, -3.2_6_9_3, 1_1.0_9_5_6], [1_1.8_5_2_4, -3.4_9_3_8, 1_3.1_2_1_0]]])

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





       @slow
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          _lowercase              =MraForMaskedLM.from_pretrained('uw-madison/mra-base-4096-8-d3')
          _lowercase              =torch.arange(4096).unsqueeze(0)

          with torch.no_grad():
             _lowercase              =model(snake_case)[0]

          _lowercase              =5_0265

          _lowercase              =torch.Size((1, 4096, vocab_size))
          self.assertEqual(output.shape,      snake_case)

          _lowercase              =torch.tensor(
              [[[5.4_7_8_9, -2.3_5_6_4, 7.5_0_6_4], [7.9_0_6_7, -1.3_3_6_9, 9.9_6_6_8], [9.0_7_1_2, -1.8_1_0_6, 7.0_3_8_0]]])

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


