





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





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase					,						__lowerCamelCase=13					,						__lowerCamelCase=7					,						__lowerCamelCase=True					,						__lowerCamelCase=True					,						__lowerCamelCase=True					,						__lowerCamelCase=True					,						__lowerCamelCase=99					,						__lowerCamelCase=64					,						__lowerCamelCase=32					,						__lowerCamelCase=5					,						__lowerCamelCase=4					,						__lowerCamelCase=37					,						__lowerCamelCase="gelu"					,						__lowerCamelCase=0.1					,						__lowerCamelCase=0.1					,						__lowerCamelCase=512					,						__lowerCamelCase=16					,						__lowerCamelCase=2					,						__lowerCamelCase=0.0_2					,						__lowerCamelCase=3					,						__lowerCamelCase=4					,						__lowerCamelCase=None					,						):



     '''simple docstring'''




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

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						int						 =     ids_tensor([self.batch_size, self.seq_length]					,						self.vocab_size  )

     __A    :						Union[str, Any]						 =     None
     if self.use_input_mask:
         __A    :						Optional[int]						 =     random_attention_mask([self.batch_size, self.seq_length]  )

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

     __A    :						Dict						 =     None
     __A    :						int						 =     None
     __A    :						str						 =     None
     if self.use_labels:
         __A    :						Union[str, Any]						 =     ids_tensor([self.batch_size]					,						self.type_sequence_label_size  )
         __A    :						Optional[int]						 =     ids_tensor([self.batch_size, self.seq_length]					,						self.num_labels  )
         __A    :						str						 =     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  ):



     '''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=lowercase_					,						initializer_range=self.initializer_range					,						)

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Optional[Any]						 =     MegatronBertModel(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     __A    :						List[str]						 =     model(lowercase_					,						attention_mask=lowercase_					,						token_type_ids=lowercase_  )
     __A    :						int						 =     model(lowercase_					,						token_type_ids=lowercase_  )
     __A    :						Optional[Any]						 =     model(lowercase_  )

     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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						List[str]						 =     MegatronBertForMaskedLM(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     __A    :						Tuple						 =     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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Tuple						 =     MegatronBertForCausalLM(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     __A    :						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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Tuple						 =     MegatronBertForNextSentencePrediction(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     __A    :						Dict						 =     model(
         lowercase_					,						attention_mask=lowercase_					,						token_type_ids=lowercase_					,						labels=lowercase_					,						)
     self.parent.assertEqual(result.logits.shape					,						(self.batch_size, 2)  )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Union[str, Any]						 =     MegatronBertForPreTraining(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     __A    :						List[str]						 =     model(
         lowercase_					,						attention_mask=lowercase_					,						token_type_ids=lowercase_					,						labels=lowercase_					,						next_sentence_label=lowercase_					,						)
     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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Any						 =     MegatronBertForQuestionAnswering(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     __A    :						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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Tuple						 =     self.num_labels
     __A    :						Union[str, Any]						 =     MegatronBertForSequenceClassification(lowercase_  )
     model.to(lowercase_  )
     model.eval()
     __A    :						Tuple						 =     model(lowercase_					,						attention_mask=lowercase_					,						token_type_ids=lowercase_					,						labels=lowercase_  )
     self.parent.assertEqual(result.logits.shape					,						(self.batch_size, self.num_labels)  )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						int						 =     self.num_labels
     __A    :						List[str]						 =     MegatronBertForTokenClassification(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     __A    :						Union[str, 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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




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







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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






@require_torch
class 		__snake_case					(		_UpperCAmelCase							,							_UpperCAmelCase							,							unittest.TestCase     ):





 """simple docstring"""


 _lowerCamelCase        =     (
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 _lowerCamelCase        =     (
     {
         """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 {}
 )
 _lowerCamelCase        =     True
 # test_resize_embeddings = False
 _lowerCamelCase        =     False

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase=False  ):



     '''simple docstring'''




     __A    :						str						 =     super()._prepare_for_class(lowercase_					,						lowercase_					,						return_labels=lowercase_  )

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

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[Any]						 =     MegatronBertModelTester(self  )
     __A    :						List[Any]						 =     ConfigTester(self					,						config_class=lowercase_					,						hidden_size=37  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     self.config_tester.run_common_tests()

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						str						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_model(*lowercase_  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[Any]						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase_  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						int						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase_  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						int						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase_  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Any						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowercase_  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Any						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase_  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Dict						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase_  )







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Any						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase_  )







def 		__lowercase     (   snake_case_  :  int						)   ->Any:



    '''simple docstring'''







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


a_ =			1e-4






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





 """simple docstring"""


 @slow
 @unittest.skip('''Model is not available.'''  )
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						str						 =     '''nvidia/megatron-bert-uncased-345m'''
     if "MYDIR" in os.environ:
         __A    :						str						 =     os.path.join(os.environ['''MYDIR''']					,						lowercase_  )
     __A    :						Tuple						 =     MegatronBertModel.from_pretrained(lowercase_  )
     model.to(lowercase_  )
     model.half()
     __A    :						List[str]						 =     _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]  )
     with torch.no_grad():
         __A    :						str						 =     model(lowercase_  )[0]
     __A    :						List[Any]						 =     torch.Size((1, 9, 1024)  )
     self.assertEqual(output.shape					,						lowercase_  )

     __A    :						List[str]						 =     [-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
     for ii in range(3  ):
         for jj in range(3  ):
             __A    :						Tuple						 =     output[0, ii, jj]
             __A    :						Optional[int]						 =     expected[3 * ii + jj]
             __A    :						int						 =     '''ii={} jj={} a={} b={}'''.format(lowercase_					,						lowercase_					,						lowercase_					,						lowercase_  )
             self.assertTrue(math.isclose(lowercase_					,						lowercase_					,						rel_tol=lowercase_					,						abs_tol=lowercase_  )					,						msg=lowercase_  )







import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest


class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    lowerCamelCase			=							(DPMSolverSinglestepScheduler,)
    lowerCamelCase			=							(('num_inference_steps', 25),)


    def 						snake_case__     (      self :       Tuple,**lowercase_ :       Dict		)->    Optional[int]:





        '''simple docstring'''







        A__          =   {
            'num_train_timesteps': 1_0_0_0,
            'beta_start': 0.0_001,
            'beta_end': 0.02,
            'beta_schedule': 'linear',
            'solver_order': 2,
            'prediction_type': 'epsilon',
            'thresholding': False,
            'sample_max_value': 1.0,
            'algorithm_type': 'dpmsolver++',
            'solver_type': 'midpoint',
            'lambda_min_clipped': -float('inf'		),
            'variance_type': None,
        }

        config.update(**lowercase_		)
        return config


    def 						snake_case__     (      self :       str,lowercase_ :       Optional[Any]=0,**lowercase_ :       Any		)->    List[Any]:





        '''simple docstring'''







        A__          =   dict(self.forward_default_kwargs		)
        A__          =   kwargs.pop('num_inference_steps',lowercase_		)
        A__          =   self.dummy_sample
        A__          =   0.1 * sample
        A__          =   [residual + 0.2, residual + 0.15, residual + 0.10]

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

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

            A__      ,						A__          =   sample, sample
            for t in range(lowercase_,time_step + scheduler.config.solver_order + 1		):
                A__          =   scheduler.step(lowercase_,lowercase_,lowercase_,**lowercase_		).prev_sample
                A__          =   new_scheduler.step(lowercase_,lowercase_,lowercase_,**lowercase_		).prev_sample

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


    def 						snake_case__     (      self :       List[str]		)->    List[Any]:





        '''simple docstring'''







        pass


    def 						snake_case__     (      self :       Tuple,lowercase_ :       Union[str, Any]=0,**lowercase_ :       Union[str, Any]		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   dict(self.forward_default_kwargs		)
        A__          =   kwargs.pop('num_inference_steps',lowercase_		)
        A__          =   self.dummy_sample
        A__          =   0.1 * sample
        A__          =   [residual + 0.2, residual + 0.15, residual + 0.10]

        for scheduler_class in self.scheduler_classes:
            A__          =   self.get_scheduler_config()
            A__          =   scheduler_class(**lowercase_		)
            scheduler.set_timesteps(lowercase_		)

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

            with tempfile.TemporaryDirectory() as tmpdirname:
                scheduler.save_config(lowercase_		)
                A__          =   scheduler_class.from_pretrained(lowercase_		)
                # copy over dummy past residuals
                new_scheduler.set_timesteps(lowercase_		)

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

            A__          =   scheduler.step(lowercase_,lowercase_,lowercase_,**lowercase_		).prev_sample
            A__          =   new_scheduler.step(lowercase_,lowercase_,lowercase_,**lowercase_		).prev_sample

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


    def 						snake_case__     (      self :       Optional[Any],lowercase_ :       Optional[int]=None,**lowercase_ :       int		)->    int:





        '''simple docstring'''







        if scheduler is None:
            A__          =   self.scheduler_classes[0]
            A__          =   self.get_scheduler_config(**lowercase_		)
            A__          =   scheduler_class(**lowercase_		)

        A__          =   self.scheduler_classes[0]
        A__          =   self.get_scheduler_config(**lowercase_		)
        A__          =   scheduler_class(**lowercase_		)

        A__          =   1_0
        A__          =   self.dummy_model()
        A__          =   self.dummy_sample_deter
        scheduler.set_timesteps(lowercase_		)

        for i, t in enumerate(scheduler.timesteps		):
            A__          =   model(lowercase_,lowercase_		)
            A__          =   scheduler.step(lowercase_,lowercase_,lowercase_		).prev_sample

        return sample


    def 						snake_case__     (      self :       Any		)->    str:





        '''simple docstring'''







        A__          =   DPMSolverSinglestepScheduler(**self.get_scheduler_config()		)
        A__          =   5_0
        A__          =   self.dummy_model()
        A__          =   self.dummy_sample_deter
        scheduler.set_timesteps(lowercase_		)

        # make sure that the first t is uneven
        for i, t in enumerate(scheduler.timesteps[3:]		):
            A__          =   model(lowercase_,lowercase_		)
            A__          =   scheduler.step(lowercase_,lowercase_,lowercase_		).prev_sample

        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_574		) < 1E-3


    def 						snake_case__     (      self :       Optional[Any]		)->    List[Any]:





        '''simple docstring'''







        for timesteps in [2_5, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
            self.check_over_configs(num_train_timesteps=lowercase_		)


    def 						snake_case__     (      self :       int		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   DPMSolverSinglestepScheduler(**self.get_scheduler_config()		)
        A__          =   self.full_loop(scheduler=lowercase_		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_791		) < 1E-3

        A__          =   DEISMultistepScheduler.from_config(scheduler.config		)
        A__          =   DPMSolverMultistepScheduler.from_config(scheduler.config		)
        A__          =   UniPCMultistepScheduler.from_config(scheduler.config		)
        A__          =   DPMSolverSinglestepScheduler.from_config(scheduler.config		)

        A__          =   self.full_loop(scheduler=lowercase_		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_791		) < 1E-3


    def 						snake_case__     (      self :       Tuple		)->    Any:





        '''simple docstring'''







        self.check_over_configs(thresholding=lowercase_		)
        for order in [1, 2, 3]:
            for solver_type in ["midpoint", "heun"]:
                for threshold in [0.5, 1.0, 2.0]:
                    for prediction_type in ["epsilon", "sample"]:
                        self.check_over_configs(
                            thresholding=lowercase_,prediction_type=lowercase_,sample_max_value=lowercase_,algorithm_type='dpmsolver++',solver_order=lowercase_,solver_type=lowercase_,)


    def 						snake_case__     (      self :       List[Any]		)->    int:





        '''simple docstring'''







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


    def 						snake_case__     (      self :       Dict		)->    List[Any]:





        '''simple docstring'''







        for algorithm_type in ["dpmsolver", "dpmsolver++"]:
            for solver_type in ["midpoint", "heun"]:
                for order in [1, 2, 3]:
                    for prediction_type in ["epsilon", "sample"]:
                        self.check_over_configs(
                            solver_order=lowercase_,solver_type=lowercase_,prediction_type=lowercase_,algorithm_type=lowercase_,)
                        A__          =   self.full_loop(
                            solver_order=lowercase_,solver_type=lowercase_,prediction_type=lowercase_,algorithm_type=lowercase_,)
                        assert not torch.isnan(lowercase_		).any(), "Samples have nan numbers"


    def 						snake_case__     (      self :       Optional[int]		)->    Tuple:





        '''simple docstring'''







        self.check_over_configs(lower_order_final=lowercase_		)
        self.check_over_configs(lower_order_final=lowercase_		)


    def 						snake_case__     (      self :       Tuple		)->    Optional[int]:





        '''simple docstring'''







        self.check_over_configs(lambda_min_clipped=-float('inf'		)		)
        self.check_over_configs(lambda_min_clipped=-5.1		)


    def 						snake_case__     (      self :       Optional[Any]		)->    Tuple:





        '''simple docstring'''







        self.check_over_configs(variance_type=lowercase_		)
        self.check_over_configs(variance_type='learned_range'		)


    def 						snake_case__     (      self :       str		)->    Any:





        '''simple docstring'''







        for num_inference_steps in [1, 2, 3, 5, 1_0, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
            self.check_over_forward(num_inference_steps=lowercase_,time_step=0		)


    def 						snake_case__     (      self :       Tuple		)->    Tuple:





        '''simple docstring'''







        A__          =   self.full_loop()
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_791		) < 1E-3


    def 						snake_case__     (      self :       Any		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   self.full_loop(use_karras_sigmas=lowercase_		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_248		) < 1E-3


    def 						snake_case__     (      self :       Union[str, Any]		)->    Tuple:





        '''simple docstring'''







        A__          =   self.full_loop(prediction_type='v_prediction'		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.1_453		) < 1E-3


    def 						snake_case__     (      self :       Tuple		)->    int:





        '''simple docstring'''







        A__          =   self.full_loop(prediction_type='v_prediction',use_karras_sigmas=lowercase_		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.0_649		) < 1E-3


    def 						snake_case__     (      self :       List[Any]		)->    int:





        '''simple docstring'''







        A__          =   self.scheduler_classes[0]
        A__          =   self.get_scheduler_config(thresholding=lowercase_,dynamic_thresholding_ratio=0		)
        A__          =   scheduler_class(**lowercase_		)

        A__          =   1_0
        A__          =   self.dummy_model()
        A__          =   self.dummy_sample_deter.half()
        scheduler.set_timesteps(lowercase_		)

        for i, t in enumerate(scheduler.timesteps		):
            A__          =   model(lowercase_,lowercase_		)
            A__          =   scheduler.step(lowercase_,lowercase_,lowercase_		).prev_sample

        assert sample.dtype == torch.floataa

