




"""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 lowerCAmelCase__	:






			'''simple docstring'''



			def __init__(  self  ,      lowercase  ,      lowercase=2  ,      lowercase=8  ,      lowercase=True  ,      lowercase=True  ,      lowercase=True  ,      lowercase=True  ,      lowercase=99  ,      lowercase=16  ,      lowercase=5  ,      lowercase=2  ,      lowercase=36  ,      lowercase="gelu"  ,      lowercase=0.0  ,      lowercase=0.0  ,      lowercase=512  ,      lowercase=16  ,      lowercase=2  ,      lowercase=0.02  ,      lowercase=3  ,      lowercase=4  ,      lowercase=None  ,      ):
							_lowerCamelCase     :			Optional[int]    =    parent
							_lowerCamelCase     :			Tuple    =    batch_size
							_lowerCamelCase     :			int    =    seq_length
							_lowerCamelCase     :			Union[str, Any]    =    is_training
							_lowerCamelCase     :			Any    =    use_input_mask
							_lowerCamelCase     :			Union[str, Any]    =    use_token_type_ids
							_lowerCamelCase     :			str    =    use_labels
							_lowerCamelCase     :			int    =    vocab_size
							_lowerCamelCase     :			Any    =    hidden_size
							_lowerCamelCase     :			int    =    num_hidden_layers
							_lowerCamelCase     :			int    =    num_attention_heads
							_lowerCamelCase     :			Tuple    =    intermediate_size
							_lowerCamelCase     :			Optional[int]    =    hidden_act
							_lowerCamelCase     :			Optional[Any]    =    hidden_dropout_prob
							_lowerCamelCase     :			Union[str, Any]    =    attention_probs_dropout_prob
							_lowerCamelCase     :			Tuple    =    max_position_embeddings
							_lowerCamelCase     :			List[Any]    =    type_vocab_size
							_lowerCamelCase     :			str    =    type_sequence_label_size
							_lowerCamelCase     :			Tuple    =    initializer_range
							_lowerCamelCase     :			str    =    num_labels
							_lowerCamelCase     :			Tuple    =    num_choices
							_lowerCamelCase     :			Optional[Any]    =    scope
			def 					A_    (  self  ):
							_lowerCamelCase     :			Dict    =    ids_tensor([self.batch_size, self.seq_length]  ,      self.vocab_size  )

							_lowerCamelCase     :			Tuple    =    None
							if self.use_input_mask:
											_lowerCamelCase     :			Union[str, Any]    =    random_attention_mask([self.batch_size, self.seq_length]  )

							_lowerCamelCase     :			Union[str, Any]    =    None
							if self.use_token_type_ids:
											_lowerCamelCase     :			Tuple    =    ids_tensor([self.batch_size, self.seq_length]  ,      self.type_vocab_size  )

							_lowerCamelCase     :			List[str]    =    None
							_lowerCamelCase     :			List[str]    =    None
							_lowerCamelCase     :			Optional[int]    =    None
							if self.use_labels:
											_lowerCamelCase     :			Optional[int]    =    ids_tensor([self.batch_size]  ,      self.type_sequence_label_size  )
											_lowerCamelCase     :			str    =    ids_tensor([self.batch_size, self.seq_length]  ,      self.num_labels  )
											_lowerCamelCase     :			str    =    ids_tensor([self.batch_size]  ,      self.num_choices  )

							_lowerCamelCase     :			List[str]    =    self.get_config()

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
			def 					A_    (  self  ):
							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=lowercase  ,      initializer_range=self.initializer_range  ,      )
			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[Any]    =    self.get_config()
							_lowerCamelCase     :			List[str]    =    300
							return config
			def 					A_    (  self  ):
							(
							    (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), 
							)     :			Union[str, Any]    =    self.prepare_config_and_inputs()

							_lowerCamelCase     :			List[Any]    =    True
							_lowerCamelCase     :			Tuple    =    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
							_lowerCamelCase     :			List[str]    =    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 					A_    (  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ):
							_lowerCamelCase     :			List[str]    =    MraModel(config=lowercase  )
							model.to(lowercase  )
							model.eval()
							_lowerCamelCase     :			List[Any]    =    model(lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  )
							_lowerCamelCase     :			Optional[Any]    =    model(lowercase  ,      token_type_ids=lowercase  )
							_lowerCamelCase     :			str    =    model(lowercase  )
							self.parent.assertEqual(result.last_hidden_state.shape  ,      (self.batch_size, self.seq_length, self.hidden_size)  )
			def 					A_    (  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      ):
							_lowerCamelCase     :			Dict    =    True
							_lowerCamelCase     :			Optional[Any]    =    MraModel(lowercase  )
							model.to(lowercase  )
							model.eval()
							_lowerCamelCase     :			str    =    model(
							    lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  ,      encoder_hidden_states=lowercase  ,      encoder_attention_mask=lowercase  ,      )
							_lowerCamelCase     :			Tuple    =    model(
							    lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  ,      encoder_hidden_states=lowercase  ,      )
							_lowerCamelCase     :			int    =    model(lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  )
							self.parent.assertEqual(result.last_hidden_state.shape  ,      (self.batch_size, self.seq_length, self.hidden_size)  )
			def 					A_    (  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ):
							_lowerCamelCase     :			Tuple    =    MraForMaskedLM(config=lowercase  )
							model.to(lowercase  )
							model.eval()
							_lowerCamelCase     :			Optional[Any]    =    model(lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  ,      labels=lowercase  )
							self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.seq_length, self.vocab_size)  )
			def 					A_    (  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ):
							_lowerCamelCase     :			Tuple    =    MraForQuestionAnswering(config=lowercase  )
							model.to(lowercase  )
							model.eval()
							_lowerCamelCase     :			Union[str, Any]    =    model(
							    lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  ,      start_positions=lowercase  ,      end_positions=lowercase  ,      )
							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 					A_    (  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ):
							_lowerCamelCase     :			List[Any]    =    self.num_labels
							_lowerCamelCase     :			Optional[int]    =    MraForSequenceClassification(lowercase  )
							model.to(lowercase  )
							model.eval()
							_lowerCamelCase     :			Any    =    model(lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  ,      labels=lowercase  )
							self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.num_labels)  )
			def 					A_    (  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ):
							_lowerCamelCase     :			int    =    self.num_labels
							_lowerCamelCase     :			Optional[Any]    =    MraForTokenClassification(config=lowercase  )
							model.to(lowercase  )
							model.eval()
							_lowerCamelCase     :			List[Any]    =    model(lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  ,      labels=lowercase  )
							self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.seq_length, self.num_labels)  )
			def 					A_    (  self  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ,      lowercase  ):
							_lowerCamelCase     :			Tuple    =    self.num_choices
							_lowerCamelCase     :			Optional[int]    =    MraForMultipleChoice(config=lowercase  )
							model.to(lowercase  )
							model.eval()
							_lowerCamelCase     :			Dict    =    input_ids.unsqueeze(1  ).expand(-1  ,      self.num_choices  ,      -1  ).contiguous()
							_lowerCamelCase     :			List[Any]    =    token_type_ids.unsqueeze(1  ).expand(-1  ,      self.num_choices  ,      -1  ).contiguous()
							_lowerCamelCase     :			str    =    input_mask.unsqueeze(1  ).expand(-1  ,      self.num_choices  ,      -1  ).contiguous()
							_lowerCamelCase     :			Tuple    =    model(
							    lowercase  ,      attention_mask=lowercase  ,      token_type_ids=lowercase  ,      labels=lowercase  ,      )
							self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.num_choices)  )






			def 					A_    (  self  ):
							_lowerCamelCase     :			Union[str, Any]    =    self.prepare_config_and_inputs()
							(
							    (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), (
							    _lowerCamelCase
							), 
							)     :			Union[str, Any]    =    config_and_inputs
							_lowerCamelCase     :			List[Any]    =    {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
							return config, inputs_dict






@require_torch
class lowerCAmelCase__	(     lowercase,	unittest.TestCase						):






			'''simple docstring'''



			lowerCamelCase__      =			(
			    (
			        MraModel,
			        MraForMaskedLM,
			        MraForMultipleChoice,
			        MraForQuestionAnswering,
			        MraForSequenceClassification,
			        MraForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase__      =			False
			lowerCamelCase__      =			False
			lowerCamelCase__      =			False
			lowerCamelCase__      =			False

			lowerCamelCase__      =			()
			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[Any]    =    MraModelTester(self  )
							_lowerCamelCase     :			int    =    ConfigTester(self  ,      config_class=lowercase  ,      hidden_size=37  )
			def 					A_    (  self  ):
							self.config_tester.run_common_tests()
			def 					A_    (  self  ):
							_lowerCamelCase     :			Union[str, Any]    =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowercase  )
			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[int]    =    self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
											_lowerCamelCase     :			Dict    =    type
											self.model_tester.create_and_check_model(*lowercase  )
			def 					A_    (  self  ):
							_lowerCamelCase     :			Tuple    =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*lowercase  )
			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[Any]    =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*lowercase  )
			def 					A_    (  self  ):
							_lowerCamelCase     :			Any    =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*lowercase  )
			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[Any]    =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*lowercase  )
			def 					A_    (  self  ):
							_lowerCamelCase     :			List[Any]    =    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*lowercase  )
			@slow
			def 					A_    (  self  ):
							for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											_lowerCamelCase     :			Union[str, Any]    =    MraModel.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )






			@unittest.skip(reason='MRA does not output attentions'  )
			def 					A_    (  self  ):
							return






@require_torch
class lowerCAmelCase__	(     unittest.TestCase						):






			'''simple docstring'''



			@slow
			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[Any]    =    MraModel.from_pretrained('uw-madison/mra-base-512-4'  )
							_lowerCamelCase     :			List[Any]    =    torch.arange(256  ).unsqueeze(0  )

							with torch.no_grad():
											_lowerCamelCase     :			Dict    =    model(lowercase  )[0]

							_lowerCamelCase     :			Any    =    torch.Size((1, 256, 768)  )
							self.assertEqual(output.shape  ,      lowercase  )

							_lowerCamelCase     :			Optional[int]    =    torch.tensor(
							    [[[-0.01_40, 0.08_30, -0.03_81], [0.15_46, 0.14_02, 0.02_20], [0.11_62, 0.08_51, 0.01_65]]]  )

							self.assertTrue(torch.allclose(output[:, :3, :3]  ,      lowercase  ,      atol=1E-4  )  )
			@slow
			def 					A_    (  self  ):
							_lowerCamelCase     :			Union[str, Any]    =    MraForMaskedLM.from_pretrained('uw-madison/mra-base-512-4'  )
							_lowerCamelCase     :			Any    =    torch.arange(256  ).unsqueeze(0  )

							with torch.no_grad():
											_lowerCamelCase     :			Tuple    =    model(lowercase  )[0]

							_lowerCamelCase     :			Optional[Any]    =    50265

							_lowerCamelCase     :			str    =    torch.Size((1, 256, vocab_size)  )
							self.assertEqual(output.shape  ,      lowercase  )

							_lowerCamelCase     :			Tuple    =    torch.tensor(
							    [[[9.25_95, -3.60_38, 11.88_19], [9.38_69, -3.26_93, 11.09_56], [11.85_24, -3.49_38, 13.12_10]]]  )

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






			@slow
			def 					A_    (  self  ):
							_lowerCamelCase     :			Dict    =    MraForMaskedLM.from_pretrained('uw-madison/mra-base-4096-8-d3'  )
							_lowerCamelCase     :			Any    =    torch.arange(4096  ).unsqueeze(0  )

							with torch.no_grad():
											_lowerCamelCase     :			str    =    model(lowercase  )[0]

							_lowerCamelCase     :			Tuple    =    50265

							_lowerCamelCase     :			int    =    torch.Size((1, 4096, vocab_size)  )
							self.assertEqual(output.shape  ,      lowercase  )

							_lowerCamelCase     :			int    =    torch.tensor(
							    [[[5.47_89, -2.35_64, 7.50_64], [7.90_67, -1.33_69, 9.96_68], [9.07_12, -1.81_06, 7.03_80]]]  )

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





"""simple docstring"""






import unittest

import numpy as np
import torch

from diffusers import DDIMPipeline, DDIMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class lowerCAmelCase__	(     lowercase,	unittest.TestCase						):






			'''simple docstring'''



			lowerCamelCase__      =			DDIMPipeline
			lowerCamelCase__      =			UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			lowerCamelCase__      =			PipelineTesterMixin.required_optional_params - {
			    """num_images_per_prompt""",
			    """latents""",
			    """callback""",
			    """callback_steps""",
			}
			lowerCamelCase__      =			UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
			lowerCamelCase__      =			False
			def 					A_    (  self  ):
							torch.manual_seed(0  )
							_lowerCamelCase     :			List[Any]    =    UNetaDModel(
							    block_out_channels=(32, 64)  ,      layers_per_block=2  ,      sample_size=32  ,      in_channels=3  ,      out_channels=3  ,      down_block_types=('DownBlock2D', 'AttnDownBlock2D')  ,      up_block_types=('AttnUpBlock2D', 'UpBlock2D')  ,      )
							_lowerCamelCase     :			List[str]    =    DDIMScheduler()
							_lowerCamelCase     :			Optional[int]    =    {'unet': unet, 'scheduler': scheduler}
							return components
			def 					A_    (  self  ,      lowercase  ,      lowercase=0  ):
							if str(lowercase  ).startswith('mps'  ):
											_lowerCamelCase     :			Dict    =    torch.manual_seed(lowercase  )
							else:
											_lowerCamelCase     :			List[str]    =    torch.Generator(device=lowercase  ).manual_seed(lowercase  )
							_lowerCamelCase     :			Tuple    =    {
							    'batch_size': 1,
							    'generator': generator,
							    'num_inference_steps': 2,
							    'output_type': 'numpy',
							}
							return inputs
			def 					A_    (  self  ):
							_lowerCamelCase     :			Any    =    'cpu'

							_lowerCamelCase     :			Tuple    =    self.get_dummy_components()
							_lowerCamelCase     :			Optional[Any]    =    self.pipeline_class(**lowercase  )
							pipe.to(lowercase  )
							pipe.set_progress_bar_config(disable=lowercase  )

							_lowerCamelCase     :			str    =    self.get_dummy_inputs(lowercase  )
							_lowerCamelCase     :			int    =    pipe(**lowercase  ).images
							_lowerCamelCase     :			Any    =    image[0, -3:, -3:, -1]

							self.assertEqual(image.shape  ,      (1, 32, 32, 3)  )
							_lowerCamelCase     :			Tuple    =    np.array(
							    [1.000E00, 5.717E-01, 4.717E-01, 1.000E00, 0.000E00, 1.000E00, 3.000E-04, 0.000E00, 9.000E-04]  )
							_lowerCamelCase     :			str    =    np.abs(image_slice.flatten() - expected_slice  ).max()
							self.assertLessEqual(lowercase  ,      1E-3  )
			def 					A_    (  self  ):
							super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3  )
			def 					A_    (  self  ):
							super().test_save_load_local(expected_max_difference=3E-3  )
			def 					A_    (  self  ):
							super().test_save_load_optional_components(expected_max_difference=3E-3  )






			def 					A_    (  self  ):
							super().test_inference_batch_single_identical(expected_max_diff=3E-3  )






@slow
@require_torch_gpu
class lowerCAmelCase__	(     unittest.TestCase						):






			'''simple docstring'''



			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[Any]    =    'google/ddpm-cifar10-32'

							_lowerCamelCase     :			Optional[Any]    =    UNetaDModel.from_pretrained(lowercase  )
							_lowerCamelCase     :			Dict    =    DDIMScheduler()

							_lowerCamelCase     :			Dict    =    DDIMPipeline(unet=lowercase  ,      scheduler=lowercase  )
							ddim.to(lowercase  )
							ddim.set_progress_bar_config(disable=lowercase  )

							_lowerCamelCase     :			List[str]    =    torch.manual_seed(0  )
							_lowerCamelCase     :			str    =    ddim(generator=lowercase  ,      eta=0.0  ,      output_type='numpy'  ).images

							_lowerCamelCase     :			Dict    =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)
							_lowerCamelCase     :			List[Any]    =    np.array([0.17_23, 0.16_17, 0.16_00, 0.16_26, 0.14_97, 0.15_13, 0.15_05, 0.14_42, 0.14_53]  )

							assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2






			def 					A_    (  self  ):
							_lowerCamelCase     :			Optional[int]    =    'google/ddpm-ema-bedroom-256'

							_lowerCamelCase     :			str    =    UNetaDModel.from_pretrained(lowercase  )
							_lowerCamelCase     :			str    =    DDIMScheduler.from_pretrained(lowercase  )

							_lowerCamelCase     :			Optional[int]    =    DDIMPipeline(unet=lowercase  ,      scheduler=lowercase  )
							ddpm.to(lowercase  )
							ddpm.set_progress_bar_config(disable=lowercase  )

							_lowerCamelCase     :			Tuple    =    torch.manual_seed(0  )
							_lowerCamelCase     :			int    =    ddpm(generator=lowercase  ,      output_type='numpy'  ).images

							_lowerCamelCase     :			Dict    =    image[0, -3:, -3:, -1]

							assert image.shape == (1, 256, 256, 3)
							_lowerCamelCase     :			str    =    np.array([0.00_60, 0.02_01, 0.03_44, 0.00_24, 0.00_18, 0.00_02, 0.00_22, 0.00_00, 0.00_69]  )

							assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2