






"""simple docstring"""



import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest






class 							__SCREAMING_SNAKE_CASE  (			UpperCamelCase						):
							snake_case_       =							(DEISMultistepScheduler,)
							snake_case_       =							(('num_inference_steps', 25),)


							def   _UpperCamelCase    (     self :  Optional[int]     ,      **snake_case :  List[Any]							):
													'''simple docstring'''





													A__  :					List[str]											=					{
													    """num_train_timesteps""": 1000,
													    """beta_start""": 0.0001,
													    """beta_end""": 0.02,
													    """beta_schedule""": """linear""",
													    """solver_order""": 2,
													}

													config.update(**snake_case							)
													return config


							def   _UpperCamelCase    (     self :  Optional[Any]     ,      snake_case :  Tuple=0     ,      **snake_case :  Any							):
													'''simple docstring'''





													A__  :					Optional[Any]											=					dict(self.forward_default_kwargs							)
													A__  :					str											=					kwargs.pop("""num_inference_steps"""     ,      snake_case							)
													A__  :					Dict											=					self.dummy_sample
													A__  :					Any											=					0.1 * sample
													A__  :					Union[str, Any]											=					[residual + 0.2, residual + 0.15, residual + 0.10]

													for scheduler_class in self.scheduler_classes:
																			A__  :					int											=					self.get_scheduler_config(**snake_case							)
																			A__  :					Dict											=					scheduler_class(**snake_case							)
																			scheduler.set_timesteps(snake_case							)
																			# copy over dummy past residuals
																			A__  :					Optional[Any]											=					dummy_past_residuals[: scheduler.config.solver_order]

																			with tempfile.TemporaryDirectory() as tmpdirname:
																									scheduler.save_config(snake_case							)
																									A__  :					List[str]											=					scheduler_class.from_pretrained(snake_case							)
																									new_scheduler.set_timesteps(snake_case							)
																									# copy over dummy past residuals
																									A__  :					Union[str, Any]											=					dummy_past_residuals[: new_scheduler.config.solver_order]

																			A__					,					A__  :					List[str]											=					sample, sample
																			for t in range(snake_case     ,      time_step + scheduler.config.solver_order + 1							):
																									A__  :					Union[str, Any]											=					scheduler.step(snake_case     ,      snake_case     ,      snake_case     ,      **snake_case							).prev_sample
																									A__  :					Union[str, Any]											=					new_scheduler.step(snake_case     ,      snake_case     ,      snake_case     ,      **snake_case							).prev_sample

																									assert torch.sum(torch.abs(output - new_output							)							) < 1e-5, "Scheduler outputs are not identical"


							def   _UpperCamelCase    (     self :  str							):
													'''simple docstring'''





													pass


							def   _UpperCamelCase    (     self :  Dict     ,      snake_case :  Dict=0     ,      **snake_case :  int							):
													'''simple docstring'''





													A__  :					str											=					dict(self.forward_default_kwargs							)
													A__  :					Optional[Any]											=					kwargs.pop("""num_inference_steps"""     ,      snake_case							)
													A__  :					Optional[Any]											=					self.dummy_sample
													A__  :					str											=					0.1 * sample
													A__  :					int											=					[residual + 0.2, residual + 0.15, residual + 0.10]

													for scheduler_class in self.scheduler_classes:
																			A__  :					List[Any]											=					self.get_scheduler_config()
																			A__  :					Optional[int]											=					scheduler_class(**snake_case							)
																			scheduler.set_timesteps(snake_case							)

																			# copy over dummy past residuals (must be after setting timesteps)
																			A__  :					Tuple											=					dummy_past_residuals[: scheduler.config.solver_order]

																			with tempfile.TemporaryDirectory() as tmpdirname:
																									scheduler.save_config(snake_case							)
																									A__  :					Optional[int]											=					scheduler_class.from_pretrained(snake_case							)
																									# copy over dummy past residuals
																									new_scheduler.set_timesteps(snake_case							)

																									# copy over dummy past residual (must be after setting timesteps)
																									A__  :					List[Any]											=					dummy_past_residuals[: new_scheduler.config.solver_order]

																			A__  :					Any											=					scheduler.step(snake_case     ,      snake_case     ,      snake_case     ,      **snake_case							).prev_sample
																			A__  :					List[str]											=					new_scheduler.step(snake_case     ,      snake_case     ,      snake_case     ,      **snake_case							).prev_sample

																			assert torch.sum(torch.abs(output - new_output							)							) < 1e-5, "Scheduler outputs are not identical"


							def   _UpperCamelCase    (     self :  Dict     ,      snake_case :  Optional[int]=None     ,      **snake_case :  Optional[int]							):
													'''simple docstring'''





													if scheduler is None:
																			A__  :					List[str]											=					self.scheduler_classes[0]
																			A__  :					int											=					self.get_scheduler_config(**snake_case							)
																			A__  :					List[Any]											=					scheduler_class(**snake_case							)

													A__  :					Optional[int]											=					self.scheduler_classes[0]
													A__  :					List[Any]											=					self.get_scheduler_config(**snake_case							)
													A__  :					int											=					scheduler_class(**snake_case							)

													A__  :					Optional[Any]											=					10
													A__  :					List[Any]											=					self.dummy_model()
													A__  :					List[Any]											=					self.dummy_sample_deter
													scheduler.set_timesteps(snake_case							)

													for i, t in enumerate(scheduler.timesteps							):
																			A__  :					Optional[Any]											=					model(snake_case     ,      snake_case							)
																			A__  :					Any											=					scheduler.step(snake_case     ,      snake_case     ,      snake_case							).prev_sample

													return sample


							def   _UpperCamelCase    (     self :  int							):
													'''simple docstring'''





													A__  :					Optional[int]											=					dict(self.forward_default_kwargs							)

													A__  :					int											=					kwargs.pop("""num_inference_steps"""     ,      snake_case							)

													for scheduler_class in self.scheduler_classes:
																			A__  :					int											=					self.get_scheduler_config()
																			A__  :					List[Any]											=					scheduler_class(**snake_case							)

																			A__  :					Optional[Any]											=					self.dummy_sample
																			A__  :					List[str]											=					0.1 * sample

																			if num_inference_steps is not None and hasattr(snake_case     ,      """set_timesteps"""							):
																									scheduler.set_timesteps(snake_case							)
																			elif num_inference_steps is not None and not hasattr(snake_case     ,      """set_timesteps"""							):
																									A__  :					Dict											=					num_inference_steps

																			# copy over dummy past residuals (must be done after set_timesteps)
																			A__  :					Tuple											=					[residual + 0.2, residual + 0.15, residual + 0.10]
																			A__  :					Tuple											=					dummy_past_residuals[: scheduler.config.solver_order]

																			A__  :					str											=					scheduler.timesteps[5]
																			A__  :					Optional[int]											=					scheduler.timesteps[6]

																			A__  :					Any											=					scheduler.step(snake_case     ,      snake_case     ,      snake_case     ,      **snake_case							).prev_sample
																			A__  :					Any											=					scheduler.step(snake_case     ,      snake_case     ,      snake_case     ,      **snake_case							).prev_sample

																			self.assertEqual(output_a.shape     ,      sample.shape							)
																			self.assertEqual(output_a.shape     ,      output_a.shape							)


							def   _UpperCamelCase    (     self :  Union[str, Any]							):
													'''simple docstring'''





													A__  :					Optional[int]											=					DEISMultistepScheduler(**self.get_scheduler_config()							)
													A__  :					List[str]											=					self.full_loop(scheduler=snake_case							)
													A__  :					List[str]											=					torch.mean(torch.abs(snake_case							)							)

													assert abs(result_mean.item() - 0.23916							) < 1e-3

													A__  :					str											=					DPMSolverSinglestepScheduler.from_config(scheduler.config							)
													A__  :					List[Any]											=					DPMSolverMultistepScheduler.from_config(scheduler.config							)
													A__  :					Union[str, Any]											=					UniPCMultistepScheduler.from_config(scheduler.config							)
													A__  :					List[Any]											=					DEISMultistepScheduler.from_config(scheduler.config							)

													A__  :					List[str]											=					self.full_loop(scheduler=snake_case							)
													A__  :					List[str]											=					torch.mean(torch.abs(snake_case							)							)

													assert abs(result_mean.item() - 0.23916							) < 1e-3


							def   _UpperCamelCase    (     self :  str							):
													'''simple docstring'''





													for timesteps in [25, 50, 100, 999, 1000]:
																			self.check_over_configs(num_train_timesteps=snake_case							)


							def   _UpperCamelCase    (     self :  Any							):
													'''simple docstring'''





													self.check_over_configs(thresholding=snake_case							)
													for order in [1, 2, 3]:
																			for solver_type in ["logrho"]:
																									for threshold in [0.5, 1.0, 2.0]:
																															for prediction_type in ["epsilon", "sample"]:
																																					self.check_over_configs(
																																					    thresholding=snake_case     ,      prediction_type=snake_case     ,      sample_max_value=snake_case     ,      algorithm_type="""deis"""     ,      solver_order=snake_case     ,      solver_type=snake_case     ,      )


							def   _UpperCamelCase    (     self :  str							):
													'''simple docstring'''





													for prediction_type in ["epsilon", "v_prediction"]:
																			self.check_over_configs(prediction_type=snake_case							)


							def   _UpperCamelCase    (     self :  str							):
													'''simple docstring'''





													for algorithm_type in ["deis"]:
																			for solver_type in ["logrho"]:
																									for order in [1, 2, 3]:
																															for prediction_type in ["epsilon", "sample"]:
																																					self.check_over_configs(
																																					    solver_order=snake_case     ,      solver_type=snake_case     ,      prediction_type=snake_case     ,      algorithm_type=snake_case     ,      )
																																					A__  :					Tuple											=					self.full_loop(
																																					    solver_order=snake_case     ,      solver_type=snake_case     ,      prediction_type=snake_case     ,      algorithm_type=snake_case     ,      )
																																					assert not torch.isnan(snake_case							).any(), "Samples have nan numbers"


							def   _UpperCamelCase    (     self :  List[Any]							):
													'''simple docstring'''





													self.check_over_configs(lower_order_final=snake_case							)
													self.check_over_configs(lower_order_final=snake_case							)


							def   _UpperCamelCase    (     self :  Dict							):
													'''simple docstring'''





													for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1000]:
																			self.check_over_forward(num_inference_steps=snake_case     ,      time_step=0							)


							def   _UpperCamelCase    (     self :  List[Any]							):
													'''simple docstring'''





													A__  :					str											=					self.full_loop()
													A__  :					Union[str, Any]											=					torch.mean(torch.abs(snake_case							)							)

													assert abs(result_mean.item() - 0.23916							) < 1e-3


							def   _UpperCamelCase    (     self :  Optional[Any]							):
													'''simple docstring'''





													A__  :					Dict											=					self.full_loop(prediction_type="""v_prediction"""							)
													A__  :					Union[str, Any]											=					torch.mean(torch.abs(snake_case							)							)

													assert abs(result_mean.item() - 0.091							) < 1e-3


							def   _UpperCamelCase    (     self :  Optional[int]							):
													'''simple docstring'''





													A__  :					Tuple											=					self.scheduler_classes[0]
													A__  :					List[Any]											=					self.get_scheduler_config(thresholding=snake_case     ,      dynamic_thresholding_ratio=0							)
													A__  :					Union[str, Any]											=					scheduler_class(**snake_case							)

													A__  :					List[Any]											=					10
													A__  :					Dict											=					self.dummy_model()
													A__  :					Optional[Any]											=					self.dummy_sample_deter.half()
													scheduler.set_timesteps(snake_case							)

													for i, t in enumerate(scheduler.timesteps							):
																			A__  :					Optional[int]											=					model(snake_case     ,      snake_case							)
																			A__  :					str											=					scheduler.step(snake_case     ,      snake_case     ,      snake_case							).prev_sample

													assert sample.dtype == torch.floataa














"""simple docstring"""



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  :


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





													A__  :					Union[str, Any]											=					parent
													A__  :					Optional[Any]											=					batch_size
													A__  :					Dict											=					seq_length
													A__  :					str											=					is_training
													A__  :					Tuple											=					use_input_mask
													A__  :					Dict											=					use_token_type_ids
													A__  :					Dict											=					use_labels
													A__  :					int											=					vocab_size
													A__  :					List[str]											=					hidden_size
													A__  :					Union[str, Any]											=					num_hidden_layers
													A__  :					int											=					num_attention_heads
													A__  :					List[str]											=					intermediate_size
													A__  :					int											=					hidden_act
													A__  :					str											=					hidden_dropout_prob
													A__  :					Tuple											=					attention_probs_dropout_prob
													A__  :					Any											=					max_position_embeddings
													A__  :					Optional[int]											=					type_vocab_size
													A__  :					int											=					type_sequence_label_size
													A__  :					Optional[Any]											=					initializer_range
													A__  :					int											=					num_labels
													A__  :					Optional[int]											=					num_choices
													A__  :					Optional[int]											=					scope


							def   _UpperCamelCase    (     self :  Optional[int]							):
													'''simple docstring'''





													A__  :					Dict											=					ids_tensor([self.batch_size, self.seq_length]     ,      self.vocab_size							)

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

													A__  :					Optional[int]											=					None
													if self.use_token_type_ids:
																			A__  :					List[Any]											=					ids_tensor([self.batch_size, self.seq_length]     ,      self.type_vocab_size							)

													A__  :					Dict											=					None
													A__  :					List[str]											=					None
													A__  :					Union[str, Any]											=					None
													if self.use_labels:
																			A__  :					str											=					ids_tensor([self.batch_size]     ,      self.type_sequence_label_size							)
																			A__  :					str											=					ids_tensor([self.batch_size, self.seq_length]     ,      self.num_labels							)
																			A__  :					Any											=					ids_tensor([self.batch_size]     ,      self.num_choices							)

													A__  :					Optional[int]											=					self.get_config()

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


							def   _UpperCamelCase    (     self :  List[str]							):
													'''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 :  Tuple							):
													'''simple docstring'''





													A__  :					Any											=					self.get_config()
													A__  :					List[str]											=					300
													return config


							def   _UpperCamelCase    (     self :  Union[str, Any]							):
													'''simple docstring'''





													(
													    (
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					
													)  :					Tuple											=					self.prepare_config_and_inputs()

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





													A__  :					List[str]											=					MraModel(config=snake_case							)
													model.to(snake_case							)
													model.eval()
													A__  :					Dict											=					model(snake_case     ,      attention_mask=snake_case     ,      token_type_ids=snake_case							)
													A__  :					List[str]											=					model(snake_case     ,      token_type_ids=snake_case							)
													A__  :					Union[str, Any]											=					model(snake_case							)
													self.parent.assertEqual(result.last_hidden_state.shape     ,      (self.batch_size, self.seq_length, self.hidden_size)							)


							def   _UpperCamelCase    (     self :  Optional[Any]     ,      snake_case :  List[Any]     ,      snake_case :  Any     ,      snake_case :  Optional[Any]     ,      snake_case :  Union[str, Any]     ,      snake_case :  Tuple     ,      snake_case :  Dict     ,      snake_case :  str     ,      snake_case :  Dict     ,      snake_case :  str     ,      ):
													'''simple docstring'''





													A__  :					Dict											=					True
													A__  :					Optional[Any]											=					MraModel(snake_case							)
													model.to(snake_case							)
													model.eval()
													A__  :					Union[str, Any]											=					model(
													    snake_case     ,      attention_mask=snake_case     ,      token_type_ids=snake_case     ,      encoder_hidden_states=snake_case     ,      encoder_attention_mask=snake_case     ,      )
													A__  :					str											=					model(
													    snake_case     ,      attention_mask=snake_case     ,      token_type_ids=snake_case     ,      encoder_hidden_states=snake_case     ,      )
													A__  :					Optional[int]											=					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 :  int     ,      snake_case :  Optional[Any]     ,      snake_case :  Optional[int]     ,      snake_case :  int     ,      snake_case :  str     ,      snake_case :  Union[str, Any]     ,      snake_case :  Dict     ,      snake_case :  List[str]							):
													'''simple docstring'''





													A__  :					Union[str, Any]											=					MraForMaskedLM(config=snake_case							)
													model.to(snake_case							)
													model.eval()
													A__  :					List[Any]											=					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 :  Optional[Any]     ,      snake_case :  Dict     ,      snake_case :  Dict     ,      snake_case :  Dict     ,      snake_case :  List[str]     ,      snake_case :  List[str]     ,      snake_case :  Tuple     ,      snake_case :  Union[str, Any]							):
													'''simple docstring'''





													A__  :					Dict											=					MraForQuestionAnswering(config=snake_case							)
													model.to(snake_case							)
													model.eval()
													A__  :					str											=					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 :  Tuple     ,      snake_case :  List[Any]     ,      snake_case :  Optional[Any]     ,      snake_case :  Optional[Any]     ,      snake_case :  List[str]     ,      snake_case :  Optional[int]     ,      snake_case :  List[str]     ,      snake_case :  Union[str, Any]							):
													'''simple docstring'''





													A__  :					str											=					self.num_labels
													A__  :					Optional[Any]											=					MraForSequenceClassification(snake_case							)
													model.to(snake_case							)
													model.eval()
													A__  :					str											=					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 :  Union[str, Any]     ,      snake_case :  Dict     ,      snake_case :  str     ,      snake_case :  List[Any]     ,      snake_case :  Any     ,      snake_case :  Dict     ,      snake_case :  Tuple     ,      snake_case :  Optional[Any]							):
													'''simple docstring'''





													A__  :					str											=					self.num_labels
													A__  :					Union[str, Any]											=					MraForTokenClassification(config=snake_case							)
													model.to(snake_case							)
													model.eval()
													A__  :					str											=					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 :  Tuple     ,      snake_case :  Optional[Any]     ,      snake_case :  Optional[int]     ,      snake_case :  int     ,      snake_case :  Optional[Any]     ,      snake_case :  List[str]     ,      snake_case :  Dict     ,      snake_case :  Optional[Any]							):
													'''simple docstring'''





													A__  :					List[str]											=					self.num_choices
													A__  :					str											=					MraForMultipleChoice(config=snake_case							)
													model.to(snake_case							)
													model.eval()
													A__  :					int											=					input_ids.unsqueeze(1							).expand(-1     ,      self.num_choices     ,      -1							).contiguous()
													A__  :					Dict											=					token_type_ids.unsqueeze(1							).expand(-1     ,      self.num_choices     ,      -1							).contiguous()
													A__  :					Tuple											=					input_mask.unsqueeze(1							).expand(-1     ,      self.num_choices     ,      -1							).contiguous()
													A__  :					str											=					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 :  List[Any]							):
													'''simple docstring'''





													A__  :					List[str]											=					self.prepare_config_and_inputs()
													(
													    (
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					(
													    A__
													)					,					
													)  :					Dict											=					config_and_inputs
													A__  :					Optional[int]											=					{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
													return config, inputs_dict






@require_torch
class 							__SCREAMING_SNAKE_CASE  (			UpperCamelCase    ,						unittest.TestCase						):
							snake_case_       =							(
							    (
							        MraModel,
							        MraForMaskedLM,
							        MraForMultipleChoice,
							        MraForQuestionAnswering,
							        MraForSequenceClassification,
							        MraForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							snake_case_       =							False
							snake_case_       =							False
							snake_case_       =							False
							snake_case_       =							False

							snake_case_       =							()


							def   _UpperCamelCase    (     self :  int							):
													'''simple docstring'''





													A__  :					Optional[Any]											=					MraModelTester(self							)
													A__  :					List[str]											=					ConfigTester(self     ,      config_class=snake_case     ,      hidden_size=37							)


							def   _UpperCamelCase    (     self :  Tuple							):
													'''simple docstring'''





													self.config_tester.run_common_tests()


							def   _UpperCamelCase    (     self :  Tuple							):
													'''simple docstring'''





													A__  :					int											=					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*snake_case							)


							def   _UpperCamelCase    (     self :  Tuple							):
													'''simple docstring'''





													A__  :					Optional[int]											=					self.model_tester.prepare_config_and_inputs()
													for type in ["absolute", "relative_key", "relative_key_query"]:
																			A__  :					List[str]											=					type
																			self.model_tester.create_and_check_model(*snake_case							)


							def   _UpperCamelCase    (     self :  Union[str, Any]							):
													'''simple docstring'''





													A__  :					int											=					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_lm(*snake_case							)


							def   _UpperCamelCase    (     self :  Union[str, Any]							):
													'''simple docstring'''





													A__  :					Any											=					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_multiple_choice(*snake_case							)


							def   _UpperCamelCase    (     self :  Optional[int]							):
													'''simple docstring'''





													A__  :					Optional[int]											=					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_question_answering(*snake_case							)


							def   _UpperCamelCase    (     self :  int							):
													'''simple docstring'''





													A__  :					Union[str, Any]											=					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_sequence_classification(*snake_case							)


							def   _UpperCamelCase    (     self :  Union[str, Any]							):
													'''simple docstring'''





													A__  :					Any											=					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_token_classification(*snake_case							)


							@slow
							def   _UpperCamelCase    (     self :  Any							):
													'''simple docstring'''





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


							@unittest.skip(reason="""MRA does not output attentions"""							)
							def   _UpperCamelCase    (     self :  Tuple							):
													'''simple docstring'''





													return






@require_torch
class 							__SCREAMING_SNAKE_CASE  (			unittest.TestCase						):


							@slow
							def   _UpperCamelCase    (     self :  Union[str, Any]							):
													'''simple docstring'''





													A__  :					str											=					MraModel.from_pretrained("""uw-madison/mra-base-512-4"""							)
													A__  :					Any											=					torch.arange(256							).unsqueeze(0							)

													with torch.no_grad():
																			A__  :					List[Any]											=					model(snake_case							)[0]

													A__  :					List[Any]											=					torch.Size((1, 256, 768)							)
													self.assertEqual(output.shape     ,      snake_case							)

													A__  :					int											=					torch.tensor(
													    [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]]							)

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


							@slow
							def   _UpperCamelCase    (     self :  List[Any]							):
													'''simple docstring'''





													A__  :					Union[str, Any]											=					MraForMaskedLM.from_pretrained("""uw-madison/mra-base-512-4"""							)
													A__  :					Tuple											=					torch.arange(256							).unsqueeze(0							)

													with torch.no_grad():
																			A__  :					List[Any]											=					model(snake_case							)[0]

													A__  :					Dict											=					5_0265

													A__  :					List[str]											=					torch.Size((1, 256, vocab_size)							)
													self.assertEqual(output.shape     ,      snake_case							)

													A__  :					List[Any]											=					torch.tensor(
													    [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]]							)

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


							@slow
							def   _UpperCamelCase    (     self :  Dict							):
													'''simple docstring'''





													A__  :					Any											=					MraForMaskedLM.from_pretrained("""uw-madison/mra-base-4096-8-d3"""							)
													A__  :					List[Any]											=					torch.arange(4096							).unsqueeze(0							)

													with torch.no_grad():
																			A__  :					List[Any]											=					model(snake_case							)[0]

													A__  :					Union[str, Any]											=					5_0265

													A__  :					Optional[Any]											=					torch.Size((1, 4096, vocab_size)							)
													self.assertEqual(output.shape     ,      snake_case							)

													A__  :					Optional[int]											=					torch.tensor(
													    [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]]							)

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






