

import unittest

from transformers import FunnelConfig, FunnelTokenizer, 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    MODEL_FOR_PRETRAINING_MAPPING,
							    FunnelBaseModel,
							    FunnelForMaskedLM,
							    FunnelForMultipleChoice,
							    FunnelForPreTraining,
							    FunnelForQuestionAnswering,
							    FunnelForSequenceClassification,
							    FunnelForTokenClassification,
							    FunnelModel,
							)



class  __snake_case					:

		def __init__(   self  ,						_A  ,						_A=13  ,						_A=7  ,						_A=True  ,						_A=True  ,						_A=True  ,						_A=True  ,						_A=99  ,						_A=[1, 1, 2]  ,						_A=1  ,						_A=32  ,						_A=4  ,						_A=8  ,						_A=37  ,						_A="gelu_new"  ,						_A=0.1  ,						_A=0.1  ,						_A=0.0  ,						_A=512  ,						_A=3  ,						_A=0.0_2  ,						_A=3  ,						_A=4  ,						_A=None  ,						_A=False  ,						):
						SCREAMING_SNAKE_CASE_													=  parent
						SCREAMING_SNAKE_CASE_													=  batch_size
						SCREAMING_SNAKE_CASE_													=  seq_length
						SCREAMING_SNAKE_CASE_													=  is_training
						SCREAMING_SNAKE_CASE_													=  use_input_mask
						SCREAMING_SNAKE_CASE_													=  use_token_type_ids
						SCREAMING_SNAKE_CASE_													=  use_labels
						SCREAMING_SNAKE_CASE_													=  vocab_size
						SCREAMING_SNAKE_CASE_													=  block_sizes
						SCREAMING_SNAKE_CASE_													=  num_decoder_layers
						SCREAMING_SNAKE_CASE_													=  d_model
						SCREAMING_SNAKE_CASE_													=  n_head
						SCREAMING_SNAKE_CASE_													=  d_head
						SCREAMING_SNAKE_CASE_													=  d_inner
						SCREAMING_SNAKE_CASE_													=  hidden_act
						SCREAMING_SNAKE_CASE_													=  hidden_dropout
						SCREAMING_SNAKE_CASE_													=  attention_dropout
						SCREAMING_SNAKE_CASE_													=  activation_dropout
						SCREAMING_SNAKE_CASE_													=  max_position_embeddings
						SCREAMING_SNAKE_CASE_													=  type_vocab_size
						SCREAMING_SNAKE_CASE_													=  2
						SCREAMING_SNAKE_CASE_													=  num_labels
						SCREAMING_SNAKE_CASE_													=  num_choices
						SCREAMING_SNAKE_CASE_													=  scope
						SCREAMING_SNAKE_CASE_													=  initializer_std

						# Used in the tests to check the size of the first attention layer
						SCREAMING_SNAKE_CASE_													=  n_head
						# Used in the tests to check the size of the first hidden state
						SCREAMING_SNAKE_CASE_													=  self.d_model
						# Used in the tests to check the number of output hidden states/attentions
						SCREAMING_SNAKE_CASE_													=  sum(self.block_sizes) + (0 if base else self.num_decoder_layers)
						# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
						# the last hidden state of the first block (which is the first hidden state of the decoder).
						if not base:
										SCREAMING_SNAKE_CASE_													=  self.num_hidden_layers + 2

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  ids_tensor([self.batch_size, self.seq_length]  ,						self.vocab_size)

						SCREAMING_SNAKE_CASE_													=  None
						if self.use_input_mask:
										SCREAMING_SNAKE_CASE_													=  ids_tensor([self.batch_size, self.seq_length]  ,						vocab_size=2)

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

						SCREAMING_SNAKE_CASE_													=  None
						SCREAMING_SNAKE_CASE_													=  None
						SCREAMING_SNAKE_CASE_													=  None
						if self.use_labels:
										SCREAMING_SNAKE_CASE_													=  ids_tensor([self.batch_size]  ,						self.type_sequence_label_size)
										SCREAMING_SNAKE_CASE_													=  ids_tensor([self.batch_size, self.seq_length]  ,						self.num_labels)
										SCREAMING_SNAKE_CASE_													=  ids_tensor([self.batch_size]  ,						self.num_choices)
										SCREAMING_SNAKE_CASE_													=  ids_tensor([self.batch_size, self.seq_length]  ,						1)

						SCREAMING_SNAKE_CASE_													=  self.get_config()

						return (
						    config,
						    input_ids,
						    token_type_ids,
						    input_mask,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						    fake_token_labels,
						)

		def   lowerCAmelCase__       (   self):
						return FunnelConfig(
						    vocab_size=self.vocab_size  ,						block_sizes=self.block_sizes  ,						num_decoder_layers=self.num_decoder_layers  ,						d_model=self.d_model  ,						n_head=self.n_head  ,						d_head=self.d_head  ,						d_inner=self.d_inner  ,						hidden_act=self.hidden_act  ,						hidden_dropout=self.hidden_dropout  ,						attention_dropout=self.attention_dropout  ,						activation_dropout=self.activation_dropout  ,						max_position_embeddings=self.max_position_embeddings  ,						type_vocab_size=self.type_vocab_size  ,						initializer_std=self.initializer_std  ,						)

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						):
						SCREAMING_SNAKE_CASE_													=  FunnelModel(config=_A)
						model.to(_A)
						model.eval()
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						attention_mask=_A  ,						token_type_ids=_A)
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						token_type_ids=_A)
						SCREAMING_SNAKE_CASE_													=  model(_A)
						self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.d_model))

						SCREAMING_SNAKE_CASE_													=  False
						SCREAMING_SNAKE_CASE_													=  model(_A)
						self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.d_model))

						SCREAMING_SNAKE_CASE_													=  False
						SCREAMING_SNAKE_CASE_													=  model(_A)
						self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.d_model))

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						):
						SCREAMING_SNAKE_CASE_													=  FunnelBaseModel(config=_A)
						model.to(_A)
						model.eval()
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						attention_mask=_A  ,						token_type_ids=_A)
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						token_type_ids=_A)
						SCREAMING_SNAKE_CASE_													=  model(_A)
						self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, 2, self.d_model))

						SCREAMING_SNAKE_CASE_													=  False
						SCREAMING_SNAKE_CASE_													=  model(_A)
						self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, 3, self.d_model))

						SCREAMING_SNAKE_CASE_													=  False
						SCREAMING_SNAKE_CASE_													=  model(_A)
						self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, 2, self.d_model))

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						):
						SCREAMING_SNAKE_CASE_													=  self.num_labels
						SCREAMING_SNAKE_CASE_													=  FunnelForPreTraining(config=_A)
						model.to(_A)
						model.eval()
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						attention_mask=_A  ,						token_type_ids=_A  ,						labels=_A)
						self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length))

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						):
						SCREAMING_SNAKE_CASE_													=  FunnelForMaskedLM(config=_A)
						model.to(_A)
						model.eval()
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						attention_mask=_A  ,						token_type_ids=_A  ,						labels=_A)
						self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.vocab_size))

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						):
						SCREAMING_SNAKE_CASE_													=  self.num_labels
						SCREAMING_SNAKE_CASE_													=  FunnelForSequenceClassification(_A)
						model.to(_A)
						model.eval()
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						attention_mask=_A  ,						token_type_ids=_A  ,						labels=_A)
						self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_labels))

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						):
						SCREAMING_SNAKE_CASE_													=  self.num_choices
						SCREAMING_SNAKE_CASE_													=  FunnelForMultipleChoice(config=_A)
						model.to(_A)
						model.eval()
						SCREAMING_SNAKE_CASE_													=  input_ids.unsqueeze(1).expand(-1  ,						self.num_choices  ,						-1).contiguous()
						SCREAMING_SNAKE_CASE_													=  token_type_ids.unsqueeze(1).expand(-1  ,						self.num_choices  ,						-1).contiguous()
						SCREAMING_SNAKE_CASE_													=  input_mask.unsqueeze(1).expand(-1  ,						self.num_choices  ,						-1).contiguous()
						SCREAMING_SNAKE_CASE_													=  model(
						    _A  ,						attention_mask=_A  ,						token_type_ids=_A  ,						labels=_A  ,						)
						self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_choices))

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						):
						SCREAMING_SNAKE_CASE_													=  self.num_labels
						SCREAMING_SNAKE_CASE_													=  FunnelForTokenClassification(config=_A)
						model.to(_A)
						model.eval()
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						attention_mask=_A  ,						token_type_ids=_A  ,						labels=_A)
						self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.num_labels))

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						_A  ,						):
						SCREAMING_SNAKE_CASE_													=  FunnelForQuestionAnswering(config=_A)
						model.to(_A)
						model.eval()
						SCREAMING_SNAKE_CASE_													=  model(
						    _A  ,						attention_mask=_A  ,						token_type_ids=_A  ,						start_positions=_A  ,						end_positions=_A  ,						)
						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):
						SCREAMING_SNAKE_CASE_													=  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_						
						)							,     (
						    SCREAMING_SNAKE_CASE_						
						)							,     						
						)													=  config_and_inputs
						SCREAMING_SNAKE_CASE_													=  {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
						return config, inputs_dict






@require_torch
class  __snake_case					(      A__       ,	A__       ,	unittest.TestCase    ):
		__lowerCAmelCase       :		Any  	=					False
		__lowerCAmelCase       :		Optional[Any]  	=					False
		__lowerCAmelCase       :		Any  	=					(
		    (
		        FunnelModel,
		        FunnelForMaskedLM,
		        FunnelForPreTraining,
		        FunnelForQuestionAnswering,
		        FunnelForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		__lowerCAmelCase       :		List[Any]  	=					(
		    {
		        'feature-extraction': (FunnelBaseModel, FunnelModel),
		        'fill-mask': FunnelForMaskedLM,
		        'question-answering': FunnelForQuestionAnswering,
		        'text-classification': FunnelForSequenceClassification,
		        'token-classification': FunnelForTokenClassification,
		        'zero-shot': FunnelForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)

		def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A=False):
						SCREAMING_SNAKE_CASE_													=  super()._prepare_for_class(_A  ,						_A  ,						return_labels=_A)

						if return_labels:
										if model_class in get_values(_A):
														SCREAMING_SNAKE_CASE_													=  torch.zeros(
														    (self.model_tester.batch_size, self.model_tester.seq_length)  ,						dtype=torch.long  ,						device=_A)
						return inputs_dict

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  FunnelModelTester(self)
						SCREAMING_SNAKE_CASE_													=  ConfigTester(self  ,						config_class=_A)

		def   lowerCAmelCase__       (   self):
						self.config_tester.run_common_tests()

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*_A)

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*_A)

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*_A)

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*_A)

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*_A)


		def   lowerCAmelCase__       (   self  ,						_A):
						if hasattr(_A  ,						'weight') and module.weight is not None:
										module.weight.data.fill_(3)
						if hasattr(_A  ,						'bias') and module.bias is not None:
										module.bias.data.fill_(3)

						for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
										if hasattr(_A  ,						_A) and getattr(_A  ,						_A) is not None:
														SCREAMING_SNAKE_CASE_													=  getattr(_A  ,						_A)
														weight.data.fill_(3)






@require_torch
class  __snake_case					(      A__       ,	unittest.TestCase    ):
		__lowerCAmelCase       :		Union[str, Any]  	=					False
		__lowerCAmelCase       :		Any  	=					False
		__lowerCAmelCase       :		Any  	=					(
		    (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
		)

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  FunnelModelTester(self  ,						base=_A)
						SCREAMING_SNAKE_CASE_													=  ConfigTester(self  ,						config_class=_A)

		def   lowerCAmelCase__       (   self):
						self.config_tester.run_common_tests()

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_base_model(*_A)

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*_A)

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*_A)

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs_for_common()
						SCREAMING_SNAKE_CASE_													=  True

						for model_class in self.all_model_classes:
										if model_class.__name__ == "FunnelBaseModel":
														continue
										SCREAMING_SNAKE_CASE_													=  model_class(_A)
										model.to(_A)
										model.train()
										SCREAMING_SNAKE_CASE_													=  self._prepare_for_class(_A  ,						_A  ,						return_labels=_A)
										SCREAMING_SNAKE_CASE_													=  model(**_A).loss
										loss.backward()


		def   lowerCAmelCase__       (   self  ,						_A):
						if hasattr(_A  ,						'weight') and module.weight is not None:
										module.weight.data.fill_(3)
						if hasattr(_A  ,						'bias') and module.bias is not None:
										module.bias.data.fill_(3)

						for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
										if hasattr(_A  ,						_A) and getattr(_A  ,						_A) is not None:
														SCREAMING_SNAKE_CASE_													=  getattr(_A  ,						_A)
														weight.data.fill_(3)






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

		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  13
						SCREAMING_SNAKE_CASE_													=  7
						SCREAMING_SNAKE_CASE_													=  torch.arange(0  ,						batch_size * sequence_length).long().reshape(_A  ,						_A)
						SCREAMING_SNAKE_CASE_													=  [0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
						SCREAMING_SNAKE_CASE_													=  torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths])

						SCREAMING_SNAKE_CASE_													=  FunnelModel.from_pretrained('sgugger/funnel-random-tiny')
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						token_type_ids=_A)[0].abs()

						SCREAMING_SNAKE_CASE_													=  torch.tensor(2_3_4_4.8_3_5_2)
						SCREAMING_SNAKE_CASE_													=  torch.tensor(0.8_0_5_2)
						self.assertTrue(torch.allclose(output.sum()  ,						_A  ,						atol=1E-4))
						self.assertTrue(torch.allclose(output.mean()  ,						_A  ,						atol=1E-4))

						SCREAMING_SNAKE_CASE_													=  torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]])
						SCREAMING_SNAKE_CASE_													=  model(_A  ,						attention_mask=_A  ,						token_type_ids=_A)[0].abs()

						SCREAMING_SNAKE_CASE_													=  torch.tensor(2_3_4_3.8_4_2_5)
						SCREAMING_SNAKE_CASE_													=  torch.tensor(0.8_0_4_9)
						self.assertTrue(torch.allclose(output.sum()  ,						_A  ,						atol=1E-4))
						self.assertTrue(torch.allclose(output.mean()  ,						_A  ,						atol=1E-4))


		@slow
		def   lowerCAmelCase__       (   self):
						SCREAMING_SNAKE_CASE_													=  FunnelTokenizer.from_pretrained('huggingface/funnel-small')
						SCREAMING_SNAKE_CASE_													=  FunnelModel.from_pretrained('huggingface/funnel-small')
						SCREAMING_SNAKE_CASE_													=  tokenizer('Hello! I am the Funnel Transformer model.'  ,						return_tensors='pt')
						SCREAMING_SNAKE_CASE_													=  model(**_A)[0]

						SCREAMING_SNAKE_CASE_													=  torch.tensor(2_3_5.7_2_4_6)
						SCREAMING_SNAKE_CASE_													=  torch.tensor(0.0_2_5_6)
						self.assertTrue(torch.allclose(output.sum()  ,						_A  ,						atol=1E-4))
						self.assertTrue(torch.allclose(output.mean()  ,						_A  ,						atol=1E-4))





import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest



class  __snake_case					(      lowerCAmelCase__    ):
  __lowerCAmelCase       :		Union[str, Any]  	=					(HeunDiscreteScheduler,)
  __lowerCAmelCase       :		Optional[int]  	=					10

  def   lowerCAmelCase__       (   self  ,						**_A):
      SCREAMING_SNAKE_CASE_													=  {
          'num_train_timesteps': 1100,
          'beta_start': 0.0_0_0_1,
          'beta_end': 0.0_2,
          'beta_schedule': 'linear',
      }

      config.update(**_A)
      return config

  def   lowerCAmelCase__       (   self):
      for timesteps in [10, 50, 100, 1000]:
          self.check_over_configs(num_train_timesteps=_A)

  def   lowerCAmelCase__       (   self):
      for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]  ,						[0.0_0_0_2, 0.0_0_2, 0.0_2]):
          self.check_over_configs(beta_start=_A  ,						beta_end=_A)

  def   lowerCAmelCase__       (   self):
      for schedule in ["linear", "scaled_linear", "exp"]:
          self.check_over_configs(beta_schedule=_A)

  def   lowerCAmelCase__       (   self):
      for clip_sample_range in [1.0, 2.0, 3.0]:
          self.check_over_configs(clip_sample_range=_A  ,						clip_sample=_A)

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

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  self.scheduler_classes[0]
      SCREAMING_SNAKE_CASE_													=  self.get_scheduler_config()
      SCREAMING_SNAKE_CASE_													=  scheduler_class(**_A)

      scheduler.set_timesteps(self.num_inference_steps)

      SCREAMING_SNAKE_CASE_													=  self.dummy_model()
      SCREAMING_SNAKE_CASE_													=  self.dummy_sample_deter * scheduler.init_noise_sigma
      SCREAMING_SNAKE_CASE_													=  sample.to(_A)

      for i, t in enumerate(scheduler.timesteps):
          SCREAMING_SNAKE_CASE_													=  scheduler.scale_model_input(_A  ,						_A)

          SCREAMING_SNAKE_CASE_													=  model(_A  ,						_A)

          SCREAMING_SNAKE_CASE_													=  scheduler.step(_A  ,						_A  ,						_A)
          SCREAMING_SNAKE_CASE_													=  output.prev_sample

      SCREAMING_SNAKE_CASE_													=  torch.sum(torch.abs(_A))
      SCREAMING_SNAKE_CASE_													=  torch.mean(torch.abs(_A))

      if torch_device in ["cpu", "mps"]:
          assert abs(result_sum.item() - 0.1_2_3_3) < 1E-2
          assert abs(result_mean.item() - 0.0_0_0_2) < 1E-3
      else:
          # CUDA
          assert abs(result_sum.item() - 0.1_2_3_3) < 1E-2
          assert abs(result_mean.item() - 0.0_0_0_2) < 1E-3

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  self.scheduler_classes[0]
      SCREAMING_SNAKE_CASE_													=  self.get_scheduler_config(prediction_type='v_prediction')
      SCREAMING_SNAKE_CASE_													=  scheduler_class(**_A)

      scheduler.set_timesteps(self.num_inference_steps)

      SCREAMING_SNAKE_CASE_													=  self.dummy_model()
      SCREAMING_SNAKE_CASE_													=  self.dummy_sample_deter * scheduler.init_noise_sigma
      SCREAMING_SNAKE_CASE_													=  sample.to(_A)

      for i, t in enumerate(scheduler.timesteps):
          SCREAMING_SNAKE_CASE_													=  scheduler.scale_model_input(_A  ,						_A)

          SCREAMING_SNAKE_CASE_													=  model(_A  ,						_A)

          SCREAMING_SNAKE_CASE_													=  scheduler.step(_A  ,						_A  ,						_A)
          SCREAMING_SNAKE_CASE_													=  output.prev_sample

      SCREAMING_SNAKE_CASE_													=  torch.sum(torch.abs(_A))
      SCREAMING_SNAKE_CASE_													=  torch.mean(torch.abs(_A))

      if torch_device in ["cpu", "mps"]:
          assert abs(result_sum.item() - 4.6_934E-07) < 1E-2
          assert abs(result_mean.item() - 6.1_112E-10) < 1E-3
      else:
          # CUDA
          assert abs(result_sum.item() - 4.693_428_650_170_972E-07) < 1E-2
          assert abs(result_mean.item() - 0.0_0_0_2) < 1E-3

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  self.scheduler_classes[0]
      SCREAMING_SNAKE_CASE_													=  self.get_scheduler_config()
      SCREAMING_SNAKE_CASE_													=  scheduler_class(**_A)

      scheduler.set_timesteps(self.num_inference_steps  ,						device=_A)

      SCREAMING_SNAKE_CASE_													=  self.dummy_model()
      SCREAMING_SNAKE_CASE_													=  self.dummy_sample_deter.to(_A) * scheduler.init_noise_sigma

      for t in scheduler.timesteps:
          SCREAMING_SNAKE_CASE_													=  scheduler.scale_model_input(_A  ,						_A)

          SCREAMING_SNAKE_CASE_													=  model(_A  ,						_A)

          SCREAMING_SNAKE_CASE_													=  scheduler.step(_A  ,						_A  ,						_A)
          SCREAMING_SNAKE_CASE_													=  output.prev_sample

      SCREAMING_SNAKE_CASE_													=  torch.sum(torch.abs(_A))
      SCREAMING_SNAKE_CASE_													=  torch.mean(torch.abs(_A))

      if str(_A).startswith('cpu'):
          # The following sum varies between 148 and 156 on mps. Why?
          assert abs(result_sum.item() - 0.1_2_3_3) < 1E-2
          assert abs(result_mean.item() - 0.0_0_0_2) < 1E-3
      elif str(_A).startswith('mps'):
          # Larger tolerance on mps
          assert abs(result_mean.item() - 0.0_0_0_2) < 1E-2
      else:
          # CUDA
          assert abs(result_sum.item() - 0.1_2_3_3) < 1E-2
          assert abs(result_mean.item() - 0.0_0_0_2) < 1E-3


  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  self.scheduler_classes[0]
      SCREAMING_SNAKE_CASE_													=  self.get_scheduler_config()
      SCREAMING_SNAKE_CASE_													=  scheduler_class(**_A  ,						use_karras_sigmas=_A)

      scheduler.set_timesteps(self.num_inference_steps  ,						device=_A)

      SCREAMING_SNAKE_CASE_													=  self.dummy_model()
      SCREAMING_SNAKE_CASE_													=  self.dummy_sample_deter.to(_A) * scheduler.init_noise_sigma
      SCREAMING_SNAKE_CASE_													=  sample.to(_A)

      for t in scheduler.timesteps:
          SCREAMING_SNAKE_CASE_													=  scheduler.scale_model_input(_A  ,						_A)

          SCREAMING_SNAKE_CASE_													=  model(_A  ,						_A)

          SCREAMING_SNAKE_CASE_													=  scheduler.step(_A  ,						_A  ,						_A)
          SCREAMING_SNAKE_CASE_													=  output.prev_sample

      SCREAMING_SNAKE_CASE_													=  torch.sum(torch.abs(_A))
      SCREAMING_SNAKE_CASE_													=  torch.mean(torch.abs(_A))

      assert abs(result_sum.item() - 0.0_0_0_1_5) < 1E-2
      assert abs(result_mean.item() - 1.9_869_554_535_034_695E-07) < 1E-2


