




import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer
class 						_A :






  def __init__(  self		:     List[str]				,						_A		:     str				,						_A		:     List[str]=13				,						_A		:     List[str]=7				,						_A		:     Dict=True				,						_A		:     Any=True				,						_A		:     int=False				,						_A		:     str=True				,						_A		:     str=99				,						_A		:     Dict=32				,						_A		:     str=5				,						_A		:     Any=4				,						_A		:     Dict=37				,						_A		:     int="gelu"				,						_A		:     Optional[Any]=0.1				,						_A		:     List[Any]=0.1				,						_A		:     List[str]=512				,						_A		:     List[str]=16				,						_A		:     Dict=2				,						_A		:     Optional[Any]=0.02				,						_A		:     Any=3				,						_A		:     Optional[int]=4				,						_A		:     Any=None				,						)		->   Union[str, Any]:






       """simple docstring"""







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






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






       """simple docstring"""







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

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

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

       lowercase :       Optional[int]       =			None
       lowercase :       List[Any]       =			None
       lowercase :       Any       =			None
       if self.use_labels:
            lowercase :       Dict       =			ids_tensor([self.batch_size]				,						self.type_sequence_label_size				)
            lowercase :       Tuple       =			ids_tensor([self.batch_size, self.seq_length]				,						self.num_labels				)
            lowercase :       int       =			ids_tensor([self.batch_size]				,						self.num_choices				)

       lowercase :       Dict       =			self.get_config()

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






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






       """simple docstring"""







       return LlamaConfig(
           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=_A				,						initializer_range=self.initializer_range				,						)






  def  __a      (  self		:     List[Any]				,						_A		:     Dict				,						_A		:     List[Any]				,						_A		:     Dict				,						_A		:     Optional[int]				,						_A		:     Optional[Any]				,						_A		:     Union[str, Any]				,						_A		:     Optional[Any]				)		->   Optional[int]:






       """simple docstring"""







       lowercase :       str       =			LlamaModel(config=_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Union[str, Any]       =			model(_A				,						attention_mask=_A				)
       lowercase :       Union[str, Any]       =			model(_A				)
       self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, self.seq_length, self.hidden_size)				)






  def  __a      (  self		:     int				,						_A		:     Optional[int]				,						_A		:     Dict				,						_A		:     Optional[int]				,						_A		:     int				,						_A		:     str				,						_A		:     str				,						_A		:     List[Any]				,						_A		:     int				,						_A		:     List[Any]				,						)		->   List[str]:






       """simple docstring"""







       lowercase :       List[str]       =			True
       lowercase :       str       =			LlamaModel(_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Any       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						encoder_attention_mask=_A				,						)
       lowercase :       int       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						)
       lowercase :       Dict       =			model(_A				,						attention_mask=_A				)
       self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, self.seq_length, self.hidden_size)				)






  def  __a      (  self		:     List[str]				,						_A		:     List[Any]				,						_A		:     Tuple				,						_A		:     Optional[Any]				,						_A		:     Optional[int]				,						_A		:     Dict				,						_A		:     Optional[Any]				,						_A		:     List[str]				,						_A		:     Optional[int]				,						_A		:     int				,						)		->   Tuple:






       """simple docstring"""







       lowercase :       Dict       =			LlamaForCausalLM(config=_A				)
       model.to(_A				)
       model.eval()
       lowercase :       str       =			model(_A				,						attention_mask=_A				,						labels=_A				)
       self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)				)






  def  __a      (  self		:     Union[str, Any]				,						_A		:     List[str]				,						_A		:     List[Any]				,						_A		:     Dict				,						_A		:     Optional[int]				,						_A		:     Union[str, Any]				,						_A		:     Optional[Any]				,						_A		:     Any				,						_A		:     Any				,						_A		:     Dict				,						)		->   Dict:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			True
       lowercase :       List[Any]       =			True
       lowercase :       Union[str, Any]       =			LlamaForCausalLM(config=_A				)
       model.to(_A				)
       model.eval()

       # first forward pass
       lowercase :       List[str]       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						encoder_attention_mask=_A				,						use_cache=_A				,						)
       lowercase :       Dict       =			outputs.past_key_values

       # create hypothetical multiple next token and extent to next_input_ids
       lowercase :       Union[str, Any]       =			ids_tensor((self.batch_size, 3)				,						config.vocab_size				)
       lowercase :       List[str]       =			ids_tensor((self.batch_size, 3)				,						vocab_size=2				)

       # append to next input_ids and
       lowercase :       Union[str, Any]       =			torch.cat([input_ids, next_tokens]				,						dim=-1				)
       lowercase :       Any       =			torch.cat([input_mask, next_mask]				,						dim=-1				)

       lowercase :       Union[str, Any]       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						encoder_attention_mask=_A				,						output_hidden_states=_A				,						)['''hidden_states'''][0]
       lowercase :       List[Any]       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						encoder_attention_mask=_A				,						past_key_values=_A				,						output_hidden_states=_A				,						)['''hidden_states'''][0]

       # select random slice
       lowercase :       Optional[Any]       =			ids_tensor((1,)				,						output_from_past.shape[-1]				).item()
       lowercase :       Optional[Any]       =			output_from_no_past[:, -3:, random_slice_idx].detach()
       lowercase :       Dict       =			output_from_past[:, :, random_slice_idx].detach()

       self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]				)

       # test that outputs are equal for slice
       self.parent.assertTrue(torch.allclose(_A				,						_A				,						atol=1E-3				)				)






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






       """simple docstring"""







       lowercase :       Union[str, Any]       =			self.prepare_config_and_inputs()
       (
           (
           lowercase
       )      ,	(
           lowercase
       )      ,	(
           lowercase
       )      ,	(
           lowercase
       )      ,	(
           lowercase
       )      ,	(
           lowercase
       )      ,	(
           lowercase
       )      ,	
       ) :       Any       =			config_and_inputs
       lowercase :       Union[str, Any]       =			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict


@require_torch
class 						_A (  _lowerCamelCase      ,       _lowerCamelCase      ,       _lowerCamelCase      ,       unittest.TestCase				):
  _UpperCamelCase       :      str			=      (LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
  _UpperCamelCase       :      Any			=      (LlamaForCausalLM,) if is_torch_available() else ()
  _UpperCamelCase       :      Dict			=      (
      {
          '''feature-extraction''': LlamaModel,
          '''text-classification''': LlamaForSequenceClassification,
          '''text-generation''': LlamaForCausalLM,
          '''zero-shot''': LlamaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  _UpperCamelCase       :      str			=      False
  _UpperCamelCase       :      Optional[int]			=      False






  def  __a      (  self		:     List[str]				)		->   int:






       """simple docstring"""







       lowercase :       Optional[Any]       =			LlamaModelTester(self				)
       lowercase :       Optional[int]       =			ConfigTester(self				,						config_class=_A				,						hidden_size=37				)






  def  __a      (  self		:     List[str]				)		->   Dict:






       """simple docstring"""







       self.config_tester.run_common_tests()






  def  __a      (  self		:     Tuple				)		->   List[str]:






       """simple docstring"""







       lowercase :       Dict       =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_A				)






  def  __a      (  self		:     Dict				)		->   str:






       """simple docstring"""







       lowercase :       str       =			self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
            lowercase :       List[Any]       =			type
            self.model_tester.create_and_check_model(*_A				)






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






       """simple docstring"""







       lowercase      ,	lowercase :       Union[str, Any]       =			self.model_tester.prepare_config_and_inputs_for_common()
       lowercase :       Union[str, Any]       =			3
       lowercase :       Union[str, Any]       =			input_dict['''input_ids''']
       lowercase :       Any       =			input_ids.ne(1				).to(_A				)
       lowercase :       List[str]       =			ids_tensor([self.model_tester.batch_size]				,						self.model_tester.type_sequence_label_size				)
       lowercase :       Optional[int]       =			LlamaForSequenceClassification(_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Dict       =			model(_A				,						attention_mask=_A				,						labels=_A				)
       self.assertEqual(result.logits.shape				,						(self.model_tester.batch_size, self.model_tester.num_labels)				)






  def  __a      (  self		:     Any				)		->   int:






       """simple docstring"""







       lowercase      ,	lowercase :       Union[str, Any]       =			self.model_tester.prepare_config_and_inputs_for_common()
       lowercase :       Tuple       =			3
       lowercase :       List[str]       =			'''single_label_classification'''
       lowercase :       Union[str, Any]       =			input_dict['''input_ids''']
       lowercase :       int       =			input_ids.ne(1				).to(_A				)
       lowercase :       Union[str, Any]       =			ids_tensor([self.model_tester.batch_size]				,						self.model_tester.type_sequence_label_size				)
       lowercase :       Union[str, Any]       =			LlamaForSequenceClassification(_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Tuple       =			model(_A				,						attention_mask=_A				,						labels=_A				)
       self.assertEqual(result.logits.shape				,						(self.model_tester.batch_size, self.model_tester.num_labels)				)






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






       """simple docstring"""







       lowercase      ,	lowercase :       List[Any]       =			self.model_tester.prepare_config_and_inputs_for_common()
       lowercase :       Optional[Any]       =			3
       lowercase :       Dict       =			'''multi_label_classification'''
       lowercase :       int       =			input_dict['''input_ids''']
       lowercase :       str       =			input_ids.ne(1				).to(_A				)
       lowercase :       Any       =			ids_tensor(
           [self.model_tester.batch_size, config.num_labels]				,						self.model_tester.type_sequence_label_size				).to(torch.float				)
       lowercase :       Optional[Any]       =			LlamaForSequenceClassification(_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Tuple       =			model(_A				,						attention_mask=_A				,						labels=_A				)
       self.assertEqual(result.logits.shape				,						(self.model_tester.batch_size, self.model_tester.num_labels)				)






  @unittest.skip('''LLaMA buffers include complex numbers, which breaks this test'''				)
  def  __a      (  self		:     Union[str, Any]				)		->   str:






       """simple docstring"""







       pass






  @parameterized.expand([('''linear''',), ('''dynamic''',)]				)
  def  __a      (  self		:     int				,						_A		:     List[str]				)		->   Dict:






       """simple docstring"""







       lowercase      ,	lowercase :       Tuple       =			self.model_tester.prepare_config_and_inputs_for_common()
       lowercase :       int       =			ids_tensor([1, 10]				,						config.vocab_size				)
       lowercase :       Dict       =			ids_tensor([1, int(config.max_position_embeddings * 1.5				)]				,						config.vocab_size				)

       set_seed(42				)  # Fixed seed at init time so the two models get the same random weights
       lowercase :       Dict       =			LlamaModel(_A				)
       original_model.to(_A				)
       original_model.eval()
       lowercase :       Tuple       =			original_model(_A				).last_hidden_state
       lowercase :       Optional[Any]       =			original_model(_A				).last_hidden_state

       set_seed(42				)  # Fixed seed at init time so the two models get the same random weights
       lowercase :       List[Any]       =			{'''type''': scaling_type, '''factor''': 10.0}
       lowercase :       Union[str, Any]       =			LlamaModel(_A				)
       scaled_model.to(_A				)
       scaled_model.eval()
       lowercase :       Union[str, Any]       =			scaled_model(_A				).last_hidden_state
       lowercase :       int       =			scaled_model(_A				).last_hidden_state

       # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
       # maximum sequence length, so the outputs for the short input should match.
       if scaling_type == "dynamic":
            self.assertTrue(torch.allclose(_A				,						_A				,						atol=1E-5				)				)
       else:
            self.assertFalse(torch.allclose(_A				,						_A				,						atol=1E-5				)				)

       # The output should be different for long inputs
       self.assertFalse(torch.allclose(_A				,						_A				,						atol=1E-5				)				)


@require_torch
class 						_A (  unittest.TestCase				):






  @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''				)
  @slow
  def  __a      (  self		:     int				)		->   List[str]:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
       lowercase :       List[str]       =			LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-7b-hf'''				,						device_map='''auto'''				)
       lowercase :       Union[str, Any]       =			model(torch.tensor([input_ids]				)				)
       # Expected mean on dim = -1
       lowercase :       Optional[int]       =			torch.tensor([[-6.6_550, -4.1_227, -4.9_859, -3.2_406, 0.8_262, -3.0_033, 1.2_964, -3.3_699]]				)
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       lowercase :       List[str]       =			torch.tensor([-12.8_281, -7.4_453, -0.4_639, -8.0_625, -7.2_500, -8.0_000, -6.4_883, -7.7_695, -7.8_438, -7.0_312, -6.2_188, -7.1_328, -1.8_496, 1.9_961, -8.6_250, -6.7_227, -12.8_281, -6.9_492, -7.0_742, -7.7_852, -7.5_820, -7.9_062, -6.9_375, -7.9_805, -8.3_438, -8.1_562, -8.0_469, -7.6_250, -7.7_422, -7.3_398,]				)
       # fmt: on
       torch.testing.assert_close(out[0, 0, :30]				,						_A				,						atol=1E-5				,						rtol=1E-5				)






  @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''				)
  @slow
  def  __a      (  self		:     List[str]				)		->   Any:






       """simple docstring"""







       lowercase :       int       =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
       lowercase :       List[str]       =			LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-hf'''				,						device_map='''auto'''				)
       lowercase :       str       =			model(torch.tensor(_A				)				)
       # Expected mean on dim = -1
       lowercase :       Any       =			torch.tensor([[-2.0_622, -1.2_794, -1.1_638, -0.9_788, -1.4_603, -1.0_238, -1.7_893, -1.4_411]]				)
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       lowercase :       Optional[int]       =			torch.tensor([-8.1_406, -8.0_547, 2.7_461, -1.2_344, -0.1_448, -1.8_262, -1.0_020, -1.8_154, -1.6_895, -1.8_516, -2.3_574, -0.9_277, 3.7_598, 6.5_742, -1.2_998, -0.1_177, -8.1_406, -2.9_688, -2.9_199, -3.1_699, -3.5_254, -2.3_555, -2.7_988, -3.4_141, -2.8_262, -4.5_195, -3.3_379, -3.3_164, -2.7_832, -3.0_273]				)
       # fmt: on
       torch.testing.assert_close(out[0, 0, :30]				,						_A				,						atol=1E-5				,						rtol=1E-5				)






  @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''				)
  @slow
  def  __a      (  self		:     int				)		->   Dict:






       """simple docstring"""







       lowercase :       str       =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
       lowercase :       int       =			LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''				,						device_map='''auto'''				)
       lowercase :       str       =			model(torch.tensor(_A				)				)
       # Expected mean on dim = -1
       lowercase :       Optional[Any]       =			torch.tensor([[-0.8_562, -1.8_520, -0.7_551, -0.4_162, -1.5_161, -1.2_038, -2.4_823, -2.3_254]]				)
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       lowercase :       Any       =			torch.tensor([-2.2_227, 4.8_828, 0.9_023, -0.4_578, -0.7_871, -0.1_033, -0.6_221, -0.5_786, -0.7_803, -1.0_674, -1.2_920, -0.1_570, 0.8_008, 2.0_723, -0.9_497, 0.2_771, -2.2_227, -0.7_612, -1.4_346, -1.2_061, -1.6_426, -0.3_000, -0.7_139, -1.1_934, -1.8_691, -1.6_973, -1.5_947, -1.2_705, -0.3_523, -0.5_513]				)
       # fmt: on
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)






  @unittest.skip(
      '''Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test'''				)
  @slow
  def  __a      (  self		:     int				)		->   List[Any]:






       """simple docstring"""







       lowercase :       int       =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
       lowercase :       List[Any]       =			LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-70b-hf'''				,						device_map='''auto'''				)
       lowercase :       Any       =			model(torch.tensor(_A				)				)

       lowercase :       Union[str, Any]       =			torch.tensor(
           [[-4.2_327, -3.3_360, -4.6_665, -4.7_631, -1.8_180, -3.4_170, -1.4_211, -3.1_810]]				,						dtype=torch.floataa				)
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)
       # fmt: off
       lowercase :       List[str]       =			torch.tensor([-9.4_922, -3.9_551, 1.7_998, -5.6_758, -5.1_055, -5.8_984, -4.8_320, -6.8_086, -6.5_391, -5.6_172, -5.5_820, -5.5_352, 1.7_881, 3.6_289, -6.5_117, -3.4_785, -9.5_000, -6.0_352, -6.8_125, -6.0_195, -6.6_836, -5.4_727, -6.2_812, -6.0_391, -7.3_398, -7.4_297, -7.4_844, -6.5_820, -5.8_789, -5.5_312]				)
       # fmt: on
       torch.testing.assert_close(out[0, 0, :30]				,						_A				,						atol=1E-5				,						rtol=1E-5				)






  @unittest.skip('''Model is curently gated'''				)
  @slow
  def  __a      (  self		:     int				)		->   Optional[int]:






       """simple docstring"""







       lowercase :       Any       =			'''Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the "princi'''
       lowercase :       str       =			'''Simply put, the theory of relativity states that '''
       lowercase :       Optional[Any]       =			LlamaTokenizer.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''				)
       lowercase :       Union[str, Any]       =			tokenizer.encode(_A				,						return_tensors='''pt'''				)
       lowercase :       Tuple       =			LlamaForCausalLM.from_pretrained(
           '''meta-llama/Llama-2-13b-chat-hf'''				,						device_map='''sequential'''				,						use_safetensors=_A				)

       # greedy generation outputs
       lowercase :       List[str]       =			model.generate(_A				,						max_new_tokens=64				,						top_p=_A				,						temperature=1				,						do_sample=_A				)
       lowercase :       List[Any]       =			tokenizer.decode(generated_ids[0]				,						skip_special_tokens=_A				)
       self.assertEqual(_A				,						_A				)





import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest
class 						_A (  _lowerCamelCase				):
  _UpperCamelCase       :      Tuple			=      (PNDMScheduler,)
  _UpperCamelCase       :      Optional[int]			=      (('''num_inference_steps''', 5_0),)






  def  __a      (  self		:     Union[str, Any]				,						**_A		:     Optional[Any]				)		->   Optional[int]:






       """simple docstring"""







       lowercase :       List[str]       =			{
           '''num_train_timesteps''': 1_000,
           '''beta_start''': 0.0_001,
           '''beta_end''': 0.02,
           '''beta_schedule''': '''linear''',
       }

       config.update(**_A				)
       return config






  def  __a      (  self		:     Optional[Any]				,						_A		:     Any=0				,						**_A		:     int				)		->   Dict:






       """simple docstring"""







       lowercase :       List[Any]       =			dict(self.forward_default_kwargs				)
       lowercase :       Union[str, Any]       =			kwargs.pop('''num_inference_steps'''				,						_A				)
       lowercase :       Optional[Any]       =			self.dummy_sample
       lowercase :       str       =			0.1 * sample
       lowercase :       Union[str, Any]       =			[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

       for scheduler_class in self.scheduler_classes:
            lowercase :       Dict       =			self.get_scheduler_config(**_A				)
            lowercase :       int       =			scheduler_class(**_A				)
            scheduler.set_timesteps(_A				)
            # copy over dummy past residuals
            lowercase :       List[str]       =			dummy_past_residuals[:]

            with tempfile.TemporaryDirectory() as tmpdirname:
                 scheduler.save_config(_A				)
                 lowercase :       Union[str, Any]       =			scheduler_class.from_pretrained(_A				)
                 new_scheduler.set_timesteps(_A				)
                 # copy over dummy past residuals
                 lowercase :       int       =			dummy_past_residuals[:]

            lowercase :       List[str]       =			scheduler.step_prk(_A				,						_A				,						_A				,						**_A				).prev_sample
            lowercase :       List[Any]       =			new_scheduler.step_prk(_A				,						_A				,						_A				,						**_A				).prev_sample

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

            lowercase :       Dict       =			scheduler.step_plms(_A				,						_A				,						_A				,						**_A				).prev_sample
            lowercase :       List[Any]       =			new_scheduler.step_plms(_A				,						_A				,						_A				,						**_A				).prev_sample

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






  def  __a      (  self		:     int				)		->   List[str]:






       """simple docstring"""







       pass






  def  __a      (  self		:     Optional[Any]				,						_A		:     Optional[int]=0				,						**_A		:     str				)		->   Optional[Any]:






       """simple docstring"""







       lowercase :       Optional[int]       =			dict(self.forward_default_kwargs				)
       lowercase :       Optional[int]       =			kwargs.pop('''num_inference_steps'''				,						_A				)
       lowercase :       Union[str, Any]       =			self.dummy_sample
       lowercase :       Any       =			0.1 * sample
       lowercase :       int       =			[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

       for scheduler_class in self.scheduler_classes:
            lowercase :       List[str]       =			self.get_scheduler_config()
            lowercase :       List[str]       =			scheduler_class(**_A				)
            scheduler.set_timesteps(_A				)

            # copy over dummy past residuals (must be after setting timesteps)
            lowercase :       Union[str, Any]       =			dummy_past_residuals[:]

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

                 # copy over dummy past residual (must be after setting timesteps)
                 lowercase :       Optional[Any]       =			dummy_past_residuals[:]

            lowercase :       int       =			scheduler.step_prk(_A				,						_A				,						_A				,						**_A				).prev_sample
            lowercase :       Tuple       =			new_scheduler.step_prk(_A				,						_A				,						_A				,						**_A				).prev_sample

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

            lowercase :       Any       =			scheduler.step_plms(_A				,						_A				,						_A				,						**_A				).prev_sample
            lowercase :       Union[str, Any]       =			new_scheduler.step_plms(_A				,						_A				,						_A				,						**_A				).prev_sample

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






  def  __a      (  self		:     Union[str, Any]				,						**_A		:     Optional[int]				)		->   Any:






       """simple docstring"""







       lowercase :       str       =			self.scheduler_classes[0]
       lowercase :       Optional[int]       =			self.get_scheduler_config(**_A				)
       lowercase :       int       =			scheduler_class(**_A				)

       lowercase :       Any       =			10
       lowercase :       Optional[int]       =			self.dummy_model()
       lowercase :       Tuple       =			self.dummy_sample_deter
       scheduler.set_timesteps(_A				)

       for i, t in enumerate(scheduler.prk_timesteps				):
            lowercase :       List[str]       =			model(_A				,						_A				)
            lowercase :       Any       =			scheduler.step_prk(_A				,						_A				,						_A				).prev_sample

       for i, t in enumerate(scheduler.plms_timesteps				):
            lowercase :       Dict       =			model(_A				,						_A				)
            lowercase :       List[Any]       =			scheduler.step_plms(_A				,						_A				,						_A				).prev_sample

       return sample






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






       """simple docstring"""







       lowercase :       int       =			dict(self.forward_default_kwargs				)

       lowercase :       Union[str, Any]       =			kwargs.pop('''num_inference_steps'''				,						_A				)

       for scheduler_class in self.scheduler_classes:
            lowercase :       Optional[Any]       =			self.get_scheduler_config()
            lowercase :       List[Any]       =			scheduler_class(**_A				)

            lowercase :       Dict       =			self.dummy_sample
            lowercase :       str       =			0.1 * sample

            if num_inference_steps is not None and hasattr(_A				,						'''set_timesteps'''				):
                 scheduler.set_timesteps(_A				)
            elif num_inference_steps is not None and not hasattr(_A				,						'''set_timesteps'''				):
                 lowercase :       List[str]       =			num_inference_steps

            # copy over dummy past residuals (must be done after set_timesteps)
            lowercase :       Any       =			[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
            lowercase :       Optional[Any]       =			dummy_past_residuals[:]

            lowercase :       Optional[int]       =			scheduler.step_prk(_A				,						0				,						_A				,						**_A				).prev_sample
            lowercase :       str       =			scheduler.step_prk(_A				,						1				,						_A				,						**_A				).prev_sample

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

            lowercase :       Dict       =			scheduler.step_plms(_A				,						0				,						_A				,						**_A				).prev_sample
            lowercase :       Tuple       =			scheduler.step_plms(_A				,						1				,						_A				,						**_A				).prev_sample

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






  def  __a      (  self		:     Any				)		->   Optional[Any]:






       """simple docstring"""







       for timesteps in [100, 1_000]:
            self.check_over_configs(num_train_timesteps=_A				)






  def  __a      (  self		:     Tuple				)		->   Any:






       """simple docstring"""







       for steps_offset in [0, 1]:
            self.check_over_configs(steps_offset=_A				)

       lowercase :       str       =			self.scheduler_classes[0]
       lowercase :       int       =			self.get_scheduler_config(steps_offset=1				)
       lowercase :       Union[str, Any]       =			scheduler_class(**_A				)
       scheduler.set_timesteps(10				)
       assert torch.equal(
           scheduler.timesteps				,						torch.LongTensor(
               [901, 851, 851, 801, 801, 751, 751, 701, 701, 651, 651, 601, 601, 501, 401, 301, 201, 101, 1]				)				,						)






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






       """simple docstring"""







       for beta_start, beta_end in zip([0.0_001, 0.001]				,						[0.002, 0.02]				):
            self.check_over_configs(beta_start=_A				,						beta_end=_A				)






  def  __a      (  self		:     List[str]				)		->   Tuple:






       """simple docstring"""







       for schedule in ["linear", "squaredcos_cap_v2"]:
            self.check_over_configs(beta_schedule=_A				)






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






       """simple docstring"""







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






  def  __a      (  self		:     Tuple				)		->   str:






       """simple docstring"""







       for t in [1, 5, 10]:
            self.check_over_forward(time_step=_A				)






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






       """simple docstring"""







       for t, num_inference_steps in zip([1, 5, 10]				,						[10, 50, 100]				):
            self.check_over_forward(num_inference_steps=_A				)






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






       """simple docstring"""







       lowercase :       str       =			27

       for scheduler_class in self.scheduler_classes:
            lowercase :       List[str]       =			self.dummy_sample
            lowercase :       Union[str, Any]       =			0.1 * sample

            lowercase :       int       =			self.get_scheduler_config()
            lowercase :       Union[str, Any]       =			scheduler_class(**_A				)

            scheduler.set_timesteps(_A				)

            # before power of 3 fix, would error on first step, so we only need to do two
            for i, t in enumerate(scheduler.prk_timesteps[:2]				):
                 lowercase :       Union[str, Any]       =			scheduler.step_prk(_A				,						_A				,						_A				).prev_sample






  def  __a      (  self		:     List[Any]				)		->   List[Any]:






       """simple docstring"""







       with self.assertRaises(_A				):
            lowercase :       Dict       =			self.scheduler_classes[0]
            lowercase :       Union[str, Any]       =			self.get_scheduler_config()
            lowercase :       Optional[int]       =			scheduler_class(**_A				)

            scheduler.step_plms(self.dummy_sample				,						1				,						self.dummy_sample				).prev_sample






  def  __a      (  self		:     str				)		->   Dict:






       """simple docstring"""







       lowercase :       List[str]       =			self.full_loop()
       lowercase :       List[Any]       =			torch.sum(torch.abs(_A				)				)
       lowercase :       Any       =			torch.mean(torch.abs(_A				)				)

       assert abs(result_sum.item() - 198.1_318				) < 1E-2
       assert abs(result_mean.item() - 0.2_580				) < 1E-3






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






       """simple docstring"""







       lowercase :       List[str]       =			self.full_loop(prediction_type='''v_prediction'''				)
       lowercase :       Optional[int]       =			torch.sum(torch.abs(_A				)				)
       lowercase :       Tuple       =			torch.mean(torch.abs(_A				)				)

       assert abs(result_sum.item() - 67.3_986				) < 1E-2
       assert abs(result_mean.item() - 0.0_878				) < 1E-3






  def  __a      (  self		:     str				)		->   Optional[Any]:






       """simple docstring"""







       lowercase :       Any       =			self.full_loop(set_alpha_to_one=_A				,						beta_start=0.01				)
       lowercase :       List[str]       =			torch.sum(torch.abs(_A				)				)
       lowercase :       Optional[int]       =			torch.mean(torch.abs(_A				)				)

       assert abs(result_sum.item() - 230.0_399				) < 1E-2
       assert abs(result_mean.item() - 0.2_995				) < 1E-3






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






       """simple docstring"""







       lowercase :       Union[str, Any]       =			self.full_loop(set_alpha_to_one=_A				,						beta_start=0.01				)
       lowercase :       int       =			torch.sum(torch.abs(_A				)				)
       lowercase :       int       =			torch.mean(torch.abs(_A				)				)

       assert abs(result_sum.item() - 186.9_482				) < 1E-2
       assert abs(result_mean.item() - 0.2_434				) < 1E-3