





'''simple docstring'''




import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
							import torch

							from transformers import (
							    MODEL_FOR_PRETRAINING_MAPPING,
							    MegatronBertForCausalLM,
							    MegatronBertForMaskedLM,
							    MegatronBertForMultipleChoice,
							    MegatronBertForNextSentencePrediction,
							    MegatronBertForPreTraining,
							    MegatronBertForQuestionAnswering,
							    MegatronBertForSequenceClassification,
							    MegatronBertForTokenClassification,
							    MegatronBertModel,
							)




class      UpperCamelCase__					:




					"""simple docstring"""


					def __init__(						self				:			Optional[int]      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			str=13      ,						lowerCamelCase_				:			Optional[Any]=7      ,						lowerCamelCase_				:			Any=True      ,						lowerCamelCase_				:			Tuple=True      ,						lowerCamelCase_				:			List[Any]=True      ,						lowerCamelCase_				:			Any=True      ,						lowerCamelCase_				:			Union[str, Any]=99      ,						lowerCamelCase_				:			Union[str, Any]=64      ,						lowerCamelCase_				:			Tuple=32      ,						lowerCamelCase_				:			Union[str, Any]=5      ,						lowerCamelCase_				:			int=4      ,						lowerCamelCase_				:			List[str]=37      ,						lowerCamelCase_				:			Union[str, Any]="gelu"      ,						lowerCamelCase_				:			Dict=0.1      ,						lowerCamelCase_				:			List[str]=0.1      ,						lowerCamelCase_				:			Union[str, Any]=5_12      ,						lowerCamelCase_				:			Union[str, Any]=16      ,						lowerCamelCase_				:			Dict=2      ,						lowerCamelCase_				:			Tuple=0.02      ,						lowerCamelCase_				:			Dict=3      ,						lowerCamelCase_				:			Optional[Any]=4      ,						lowerCamelCase_				:			List[Any]=None      ,						):

							'''simple docstring'''



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


					def lowerCamelCase_			(						self				:			Any						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Tuple		=			ids_tensor([self.batch_size, self.seq_length]      ,						self.vocab_size						)

							SCREAMING_SNAKE_CASE :				Optional[Any]		=			None
							if self.use_input_mask:
									SCREAMING_SNAKE_CASE :				List[str]		=			random_attention_mask([self.batch_size, self.seq_length]						)

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

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

							SCREAMING_SNAKE_CASE :				List[Any]		=			self.get_config()

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


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



							return MegatronBertConfig(
							    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      ,						embedding_size=self.embedding_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=lowerCamelCase_      ,						initializer_range=self.initializer_range      ,						)


					def lowerCamelCase_			(						self				:			Tuple      ,						lowerCamelCase_				:			str      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			Union[str, Any]      ,						lowerCamelCase_				:			Tuple      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[Any]		=			MegatronBertModel(config=lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				List[Any]		=			model(lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				List[str]		=			model(lowerCamelCase_      ,						token_type_ids=lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				Tuple		=			model(lowerCamelCase_						)

							self.parent.assertEqual(result.last_hidden_state.shape      ,						(self.batch_size, self.seq_length, self.hidden_size)						)
							self.parent.assertEqual(result.pooler_output.shape      ,						(self.batch_size, self.hidden_size)						)


					def lowerCamelCase_			(						self				:			List[Any]      ,						lowerCamelCase_				:			Optional[Any]      ,						lowerCamelCase_				:			Dict      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			Union[str, Any]      ,						lowerCamelCase_				:			Dict      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			Optional[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				str		=			MegatronBertForMaskedLM(config=lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				List[str]		=			model(lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_      ,						labels=lowerCamelCase_						)
							self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.seq_length, self.vocab_size)						)


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

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[str]		=			MegatronBertForCausalLM(config=lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				List[Any]		=			model(lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_      ,						labels=lowerCamelCase_						)
							self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.seq_length, self.vocab_size)						)


					def lowerCamelCase_			(						self				:			Dict      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			int      ,						lowerCamelCase_				:			Optional[int]      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			int      ,						lowerCamelCase_				:			Any						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Any		=			MegatronBertForNextSentencePrediction(config=lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				Dict		=			model(
							    lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_      ,						labels=lowerCamelCase_      ,						)
							self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, 2)						)


					def lowerCamelCase_			(						self				:			Optional[int]      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			Tuple      ,						lowerCamelCase_				:			Tuple      ,						lowerCamelCase_				:			Dict      ,						lowerCamelCase_				:			Dict      ,						lowerCamelCase_				:			Union[str, Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				int		=			MegatronBertForPreTraining(config=lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				str		=			model(
							    lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_      ,						labels=lowerCamelCase_      ,						next_sentence_label=lowerCamelCase_      ,						)
							self.parent.assertEqual(result.prediction_logits.shape      ,						(self.batch_size, self.seq_length, self.vocab_size)						)
							self.parent.assertEqual(result.seq_relationship_logits.shape      ,						(self.batch_size, 2)						)


					def lowerCamelCase_			(						self				:			str      ,						lowerCamelCase_				:			Union[str, Any]      ,						lowerCamelCase_				:			int      ,						lowerCamelCase_				:			Optional[int]      ,						lowerCamelCase_				:			List[Any]      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			Any						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[int]		=			MegatronBertForQuestionAnswering(config=lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				Optional[int]		=			model(
							    lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_      ,						start_positions=lowerCamelCase_      ,						end_positions=lowerCamelCase_      ,						)
							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 lowerCamelCase_			(						self				:			Optional[Any]      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			Tuple      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			Union[str, Any]      ,						lowerCamelCase_				:			Optional[int]      ,						lowerCamelCase_				:			List[Any]      ,						lowerCamelCase_				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			self.num_labels
							SCREAMING_SNAKE_CASE :				str		=			MegatronBertForSequenceClassification(lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				Any		=			model(lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_      ,						labels=lowerCamelCase_						)
							self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.num_labels)						)


					def lowerCamelCase_			(						self				:			int      ,						lowerCamelCase_				:			Optional[Any]      ,						lowerCamelCase_				:			str      ,						lowerCamelCase_				:			Optional[Any]      ,						lowerCamelCase_				:			Any      ,						lowerCamelCase_				:			Union[str, Any]      ,						lowerCamelCase_				:			Optional[int]      ,						lowerCamelCase_				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				str		=			self.num_labels
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			MegatronBertForTokenClassification(config=lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			model(lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_      ,						labels=lowerCamelCase_						)
							self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.seq_length, self.num_labels)						)


					def lowerCamelCase_			(						self				:			Union[str, Any]      ,						lowerCamelCase_				:			List[Any]      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			Dict      ,						lowerCamelCase_				:			str      ,						lowerCamelCase_				:			int      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			Dict						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			self.num_choices
							SCREAMING_SNAKE_CASE :				List[Any]		=			MegatronBertForMultipleChoice(config=lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				List[str]		=			input_ids.unsqueeze(1						).expand(-1      ,						self.num_choices      ,						-1						).contiguous()
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			token_type_ids.unsqueeze(1						).expand(-1      ,						self.num_choices      ,						-1						).contiguous()
							SCREAMING_SNAKE_CASE :				int		=			input_mask.unsqueeze(1						).expand(-1      ,						self.num_choices      ,						-1						).contiguous()
							SCREAMING_SNAKE_CASE :				List[Any]		=			model(
							    lowerCamelCase_      ,						attention_mask=lowerCamelCase_      ,						token_type_ids=lowerCamelCase_      ,						labels=lowerCamelCase_      ,						)
							self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.num_choices)						)



					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



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







@require_torch
class      UpperCamelCase__					(					lowercase_						,					lowercase_						,					unittest.TestCase			):




					"""simple docstring"""
					SCREAMING_SNAKE_CASE__				=	(
					    (
					        MegatronBertModel,
					        MegatronBertForMaskedLM,
					        MegatronBertForCausalLM,
					        MegatronBertForMultipleChoice,
					        MegatronBertForNextSentencePrediction,
					        MegatronBertForPreTraining,
					        MegatronBertForQuestionAnswering,
					        MegatronBertForSequenceClassification,
					        MegatronBertForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					SCREAMING_SNAKE_CASE__				=	(
					    {
					        '''feature-extraction''': MegatronBertModel,
					        '''fill-mask''': MegatronBertForMaskedLM,
					        '''question-answering''': MegatronBertForQuestionAnswering,
					        '''text-classification''': MegatronBertForSequenceClassification,
					        '''text-generation''': MegatronBertForCausalLM,
					        '''token-classification''': MegatronBertForTokenClassification,
					        '''zero-shot''': MegatronBertForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					SCREAMING_SNAKE_CASE__				=	True
					# test_resize_embeddings = False
					SCREAMING_SNAKE_CASE__				=	False


					def lowerCamelCase_			(						self				:			str      ,						lowerCamelCase_				:			Union[str, Any]      ,						lowerCamelCase_				:			List[Any]      ,						lowerCamelCase_				:			Optional[int]=False						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			super()._prepare_for_class(lowerCamelCase_      ,						lowerCamelCase_      ,						return_labels=lowerCamelCase_						)

							if return_labels:
									if model_class in get_values(lowerCamelCase_						):
											SCREAMING_SNAKE_CASE :				Dict		=			torch.zeros(
											    (self.model_tester.batch_size, self.model_tester.seq_length)      ,						dtype=torch.long      ,						device=lowerCamelCase_						)
											SCREAMING_SNAKE_CASE :				Union[str, Any]		=			torch.zeros(
											    self.model_tester.batch_size      ,						dtype=torch.long      ,						device=lowerCamelCase_						)
							return inputs_dict


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Any		=			MegatronBertModelTester(self						)
							SCREAMING_SNAKE_CASE :				Dict		=			ConfigTester(self      ,						config_class=lowerCamelCase_      ,						hidden_size=37						)


					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							self.config_tester.run_common_tests()


					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			Any						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			Tuple						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase_						)


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

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Dict		=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			Tuple						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Dict		=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			Any						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[Any]		=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase_						)



					def lowerCamelCase_			(						self				:			Dict						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase_						)







def 		__A		(  lowerCamelCase_	):


		"""simple docstring"""



		return torch.tensor(
		    lowerCamelCase_			,			dtype=torch.long			,			device=lowerCamelCase_			,			)


__UpperCAmelCase            =       1e-4




@require_torch
@require_sentencepiece
@require_tokenizers
class      UpperCamelCase__					(					unittest.TestCase			):




					"""simple docstring"""


					@slow
					@unittest.skip("""Model is not available."""						)
					def lowerCamelCase_			(						self				:			Any						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[str]		=			"""nvidia/megatron-bert-uncased-345m"""
							if "MYDIR" in os.environ:
									SCREAMING_SNAKE_CASE :				str		=			os.path.join(os.environ["""MYDIR"""]      ,						lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				List[Any]		=			MegatronBertModel.from_pretrained(lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.half()
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			_long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]						)
							with torch.no_grad():
									SCREAMING_SNAKE_CASE :				List[str]		=			model(lowerCamelCase_						)[0]
							SCREAMING_SNAKE_CASE :				Any		=			torch.Size((1, 9, 10_24)						)
							self.assertEqual(output.shape      ,						lowerCamelCase_						)

							SCREAMING_SNAKE_CASE :				int		=			[-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
							for ii in range(3						):
									for jj in range(3						):
											SCREAMING_SNAKE_CASE :				Dict		=			output[0, ii, jj]
											SCREAMING_SNAKE_CASE :				Union[str, Any]		=			expected[3 * ii + jj]
											SCREAMING_SNAKE_CASE :				Optional[int]		=			"""ii={} jj={} a={} b={}""".format(lowerCamelCase_      ,						lowerCamelCase_      ,						lowerCamelCase_      ,						lowerCamelCase_						)
											self.assertTrue(math.isclose(lowerCamelCase_      ,						lowerCamelCase_      ,						rel_tol=lowerCamelCase_      ,						abs_tol=lowerCamelCase_						)      ,						msg=lowerCamelCase_						)










'''simple docstring'''




import torch

from diffusers import UnCLIPScheduler

from .test_schedulers import SchedulerCommonTest




class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""
					SCREAMING_SNAKE_CASE__				=	(UnCLIPScheduler,)


					def lowerCamelCase_			(						self				:			List[str]      ,						**lowerCamelCase_				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[Any]		=			{
							    """num_train_timesteps""": 10_00,
							    """variance_type""": """fixed_small_log""",
							    """clip_sample""": True,
							    """clip_sample_range""": 1.0,
							    """prediction_type""": """epsilon""",
							}

							config.update(**lowerCamelCase_						)
							return config


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							for timesteps in [1, 5, 1_00, 10_00]:
									self.check_over_configs(num_train_timesteps=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



							for variance in ["fixed_small_log", "learned_range"]:
									self.check_over_configs(variance_type=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							for clip_sample in [True, False]:
									self.check_over_configs(clip_sample=lowerCamelCase_						)


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

							'''simple docstring'''



							for clip_sample_range in [1, 5, 10, 20]:
									self.check_over_configs(clip_sample_range=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							for prediction_type in ["epsilon", "sample"]:
									self.check_over_configs(prediction_type=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							for time_step in [0, 5_00, 9_99]:
									for prev_timestep in [None, 5, 1_00, 2_50, 5_00, 7_50]:
											if prev_timestep is not None and prev_timestep >= time_step:
													continue

											self.check_over_forward(time_step=lowerCamelCase_      ,						prev_timestep=lowerCamelCase_						)


					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				int		=			self.scheduler_classes[0]
							SCREAMING_SNAKE_CASE :				Dict		=			self.get_scheduler_config(variance_type="""fixed_small_log"""						)
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			scheduler_class(**lowerCamelCase_						)

							assert torch.sum(torch.abs(scheduler._get_variance(0						) - 1.0000e-10						)						) < 1e-5
							assert torch.sum(torch.abs(scheduler._get_variance(4_87						) - 0.0_549_625						)						) < 1e-5
							assert torch.sum(torch.abs(scheduler._get_variance(9_99						) - 0.9_994_987						)						) < 1e-5


					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Tuple		=			self.scheduler_classes[0]
							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.get_scheduler_config(variance_type="""learned_range"""						)
							SCREAMING_SNAKE_CASE :				List[str]		=			scheduler_class(**lowerCamelCase_						)

							SCREAMING_SNAKE_CASE :				List[str]		=			0.5

							assert scheduler._get_variance(1      ,						predicted_variance=lowerCamelCase_						) - -10.1_712_790 < 1e-5
							assert scheduler._get_variance(4_87      ,						predicted_variance=lowerCamelCase_						) - -5.7_998_052 < 1e-5
							assert scheduler._get_variance(9_99      ,						predicted_variance=lowerCamelCase_						) - -0.0_010_011 < 1e-5


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Tuple		=			self.scheduler_classes[0]
							SCREAMING_SNAKE_CASE :				int		=			self.get_scheduler_config()
							SCREAMING_SNAKE_CASE :				List[Any]		=			scheduler_class(**lowerCamelCase_						)

							SCREAMING_SNAKE_CASE :				Optional[Any]		=			scheduler.timesteps

							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.dummy_model()
							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.dummy_sample_deter
							SCREAMING_SNAKE_CASE :				str		=			torch.manual_seed(0						)

							for i, t in enumerate(lowerCamelCase_						):
									# 1. predict noise residual
									SCREAMING_SNAKE_CASE :				List[str]		=			model(lowerCamelCase_      ,						lowerCamelCase_						)

									# 2. predict previous mean of sample x_t-1
									SCREAMING_SNAKE_CASE :				Optional[int]		=			scheduler.step(lowerCamelCase_      ,						lowerCamelCase_      ,						lowerCamelCase_      ,						generator=lowerCamelCase_						).prev_sample

									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			pred_prev_sample

							SCREAMING_SNAKE_CASE :				Optional[int]		=			torch.sum(torch.abs(lowerCamelCase_						)						)
							SCREAMING_SNAKE_CASE :				int		=			torch.mean(torch.abs(lowerCamelCase_						)						)

							assert abs(result_sum.item() - 252.2_682_495						) < 1e-2
							assert abs(result_mean.item() - 0.3_284_743						) < 1e-3


					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			self.scheduler_classes[0]
							SCREAMING_SNAKE_CASE :				Dict		=			self.get_scheduler_config()
							SCREAMING_SNAKE_CASE :				List[str]		=			scheduler_class(**lowerCamelCase_						)

							scheduler.set_timesteps(25						)

							SCREAMING_SNAKE_CASE :				List[str]		=			scheduler.timesteps

							SCREAMING_SNAKE_CASE :				Optional[int]		=			self.dummy_model()
							SCREAMING_SNAKE_CASE :				Tuple		=			self.dummy_sample_deter
							SCREAMING_SNAKE_CASE :				Any		=			torch.manual_seed(0						)

							for i, t in enumerate(lowerCamelCase_						):
									# 1. predict noise residual
									SCREAMING_SNAKE_CASE :				int		=			model(lowerCamelCase_      ,						lowerCamelCase_						)

									if i + 1 == timesteps.shape[0]:
											SCREAMING_SNAKE_CASE :				Union[str, Any]		=			None
									else:
											SCREAMING_SNAKE_CASE :				int		=			timesteps[i + 1]

									# 2. predict previous mean of sample x_t-1
									SCREAMING_SNAKE_CASE :				int		=			scheduler.step(
									    lowerCamelCase_      ,						lowerCamelCase_      ,						lowerCamelCase_      ,						prev_timestep=lowerCamelCase_      ,						generator=lowerCamelCase_						).prev_sample

									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			pred_prev_sample

							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			torch.sum(torch.abs(lowerCamelCase_						)						)
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			torch.mean(torch.abs(lowerCamelCase_						)						)

							assert abs(result_sum.item() - 258.2_044_983						) < 1e-2
							assert abs(result_mean.item() - 0.3_362_038						) < 1e-3


					def lowerCamelCase_			(						self				:			Dict						):

							'''simple docstring'''



							pass



					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



							pass



