




import unittest
from unittest.util import safe_repr

from transformers import AutoTokenizer, RwkvConfig, is_torch_available
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, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
							    RwkvForCausalLM,
							    RwkvModel,
							)
							from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_0
else:
							UpperCamelCase__		=					False


class     __SCREAMING_SNAKE_CASE							:

					def __init__(    self			,	__lowercase			,	__lowercase=14			,	__lowercase=7			,	__lowercase=True			,	__lowercase=False			,	__lowercase=True			,	__lowercase=True			,	__lowercase=True			,	__lowercase=99			,	__lowercase=32			,	__lowercase=5			,	__lowercase=37			,	__lowercase="gelu"			,	__lowercase=0.1			,	__lowercase=0.1			,	__lowercase=512			,	__lowercase=16			,	__lowercase=2			,	__lowercase=3			,	__lowercase=4			,	__lowercase=None			,	):







						'''simple docstring'''




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

					def      lowercase_							(    self ):







						'''simple docstring'''




						return RwkvConfig.from_pretrained('sgugger/rwkv-4-pile-7b' )

					def      lowercase_							(    self			,	__lowercase=False			,	__lowercase=False			,	__lowercase=False ):







						'''simple docstring'''




						lowerCamelCase_							:		str							       =							ids_tensor([self.batch_size, self.seq_length]			,	self.vocab_size )

						lowerCamelCase_							:		Tuple							       =							None
						if self.use_input_mask:
							lowerCamelCase_							:		int							       =							random_attention_mask([self.batch_size, self.seq_length] )

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

						lowerCamelCase_							:		Tuple							       =							None
						if self.use_mc_token_ids:
							lowerCamelCase_							:		int							       =							ids_tensor([self.batch_size, self.num_choices]			,	self.seq_length )

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

						lowerCamelCase_							:		List[str]							       =							self.get_config(
						    gradient_checkpointing=__lowercase			,	scale_attn_by_inverse_layer_idx=__lowercase			,	reorder_and_upcast_attn=__lowercase			,	)

						return (
						    config,
						    input_ids,
						    input_mask,
						    None,
						    token_type_ids,
						    mc_token_ids,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						)

					def      lowercase_							(    self			,	__lowercase=False			,	__lowercase=False			,	__lowercase=False ):







						'''simple docstring'''




						return RwkvConfig(
						    vocab_size=self.vocab_size			,	hidden_size=self.hidden_size			,	num_hidden_layers=self.num_hidden_layers			,	intermediate_size=self.intermediate_size			,	activation_function=self.hidden_act			,	resid_pdrop=self.hidden_dropout_prob			,	attn_pdrop=self.attention_probs_dropout_prob			,	n_positions=self.max_position_embeddings			,	type_vocab_size=self.type_vocab_size			,	use_cache=__lowercase			,	bos_token_id=self.bos_token_id			,	eos_token_id=self.eos_token_id			,	pad_token_id=self.pad_token_id			,	gradient_checkpointing=__lowercase			,	scale_attn_by_inverse_layer_idx=__lowercase			,	reorder_and_upcast_attn=__lowercase			,	)

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Dict							       =							self.get_config()
						lowerCamelCase_							:		Optional[int]							       =							300
						return config

					def      lowercase_							(    self ):







						'''simple docstring'''




						(
						    lowerCamelCase_      
						)							:		List[str]							       =							self.prepare_config_and_inputs()

						lowerCamelCase_							:		str							       =							floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
						lowerCamelCase_							:		Optional[int]							       =							ids_tensor([self.batch_size, self.seq_length]			,	vocab_size=2 )

						return (
						    config,
						    input_ids,
						    input_mask,
						    head_mask,
						    token_type_ids,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	*__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		str							       =							True
						lowerCamelCase_							:		Any							       =							RwkvModel(config=__lowercase )
						model.to(__lowercase )
						model.eval()

						lowerCamelCase_							:		str							       =							model(__lowercase )

						self.parent.assertEqual(result.last_hidden_state.shape			,	(self.batch_size, self.seq_length, self.hidden_size) )
						self.parent.assertEqual(len(result.hidden_states )			,	config.num_hidden_layers + 1 )

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	*__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		Dict							       =							RwkvForCausalLM(__lowercase )
						model.to(__lowercase )
						model.eval()

						lowerCamelCase_							:		int							       =							model(__lowercase			,	labels=__lowercase )
						self.parent.assertEqual(result.loss.shape			,	() )
						self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length, self.vocab_size) )

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	*__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		List[Any]							       =							RwkvModel(config=__lowercase )
						model.to(__lowercase )
						model.eval()

						lowerCamelCase_							:		Optional[Any]							       =							model(__lowercase )
						lowerCamelCase_							:		Any							       =							outputs.last_hidden_state

						lowerCamelCase_							:		List[Any]							       =							model(input_ids[:, :2] )
						lowerCamelCase_							:		List[str]							       =							outputs.last_hidden_state

						# Using the state computed on the first inputs, we will get the same output
						lowerCamelCase_							:		int							       =							model(input_ids[:, 2:]			,	state=outputs.state )
						lowerCamelCase_							:		List[str]							       =							outputs.last_hidden_state

						self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two]			,	dim=1 )			,	__lowercase			,	atol=1E-5 ) )

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	*__lowercase			,	__lowercase=False ):







						'''simple docstring'''




						lowerCamelCase_							:		Tuple							       =							RwkvForCausalLM(__lowercase )
						model.to(__lowercase )
						if gradient_checkpointing:
							model.gradient_checkpointing_enable()

						lowerCamelCase_							:		Union[str, Any]							       =							model(__lowercase			,	labels=__lowercase )
						self.parent.assertEqual(result.loss.shape			,	() )
						self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length, self.vocab_size) )
						result.loss.backward()





					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Tuple							       =							self.prepare_config_and_inputs()

						(
						    lowerCamelCase_      
						)							:		Any							       =							config_and_inputs

						lowerCamelCase_							:		Tuple							       =							{'input_ids': input_ids}

						return config, inputs_dict


@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0       ,					reason="""See https://github.com/huggingface/transformers/pull/24204"""  )
@require_torch
class     __SCREAMING_SNAKE_CASE							(		lowerCamelCase       ,					lowerCamelCase       ,					lowerCamelCase       ,					unittest.TestCase  ):
					_lowerCamelCase		  =			(RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
					_lowerCamelCase		  =			(
					    {"""feature-extraction""": RwkvModel, """text-generation""": RwkvForCausalLM} if is_torch_available() else {}
					)
					# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
					_lowerCamelCase		  =			False
					_lowerCamelCase		  =			False
					_lowerCamelCase		  =			False
					_lowerCamelCase		  =			False
					_lowerCamelCase		  =			False  # Rwkv does not support head masking

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[int]							       =							RwkvModelTester(self )
						lowerCamelCase_							:		List[str]							       =							ConfigTester(
						    self			,	config_class=__lowercase			,	n_embd=37			,	common_properties=['hidden_size', 'num_hidden_layers'] )

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase=None ):







						'''simple docstring'''




						if isinstance(__lowercase			,	torch.Tensor ):
							lowerCamelCase_							:		List[Any]							       =							member.max().item(), member.min().item()
						elif isinstance(__lowercase			,	__lowercase ) or isinstance(__lowercase			,	__lowercase ):
							lowerCamelCase_							:		Union[str, Any]							       =							max(__lowercase ), min(__lowercase )

						if not isinstance(__lowercase			,	__lowercase ):
							raise TypeError('container should be a list or tuple' )
						elif len(__lowercase ) != 2:
							raise ValueError('container should have 2 elements' )

						lowerCamelCase_							:		str							       =							container

						lowerCamelCase_							:		Tuple							       =							(min_value >= expected_min) and (max_value <= expected_max)

						if not is_inside_interval:
							lowerCamelCase_							:		Any							       =							'%s not found in %s' % (safe_repr(__lowercase ), safe_repr(__lowercase ))
							self.fail(self._formatMessage(__lowercase			,	__lowercase ) )

					def      lowercase_							(    self ):







						'''simple docstring'''




						self.config_tester.run_common_tests()

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Dict							       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_rwkv_model(*__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




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

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_state_equivalency(*__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		str							       =							self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
							lowerCamelCase_							:		Optional[int]							       =							model_class(config=__lowercase )
							for name, param in model.named_parameters():
								if "time_decay" in name:
									if param.requires_grad:
										self.assertTrue(param.data.max().item() == 3.0 )
										self.assertTrue(param.data.min().item() == -5.0 )
								elif "time_first" in name:
									if param.requires_grad:
										# check if it's a ones like
										self.assertTrue(torch.allclose(param.data			,	torch.ones_like(param.data )			,	atol=1E-5			,	rtol=1E-5 ) )
								elif any(x in name for x in ['time_mix_key', 'time_mix_receptance'] ):
									if param.requires_grad:
										self.assertInterval(
										    param.data			,	[0.0, 1.0]			,	msg=f"Parameter {name} of model {model_class} seems not properly initialized"			,	)
								elif "time_mix_value" in name:
									if param.requires_grad:
										self.assertInterval(
										    param.data			,	[0.0, 1.3]			,	msg=f"Parameter {name} of model {model_class} seems not properly initialized"			,	)

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[int]							       =							self.model_tester.prepare_config_and_inputs_for_common()
						lowerCamelCase_							:		Any							       =							True

						lowerCamelCase_							:		Any							       =							getattr(self.model_tester			,	'seq_length'			,	__lowercase )

						for model_class in self.all_model_classes:
							lowerCamelCase_							:		Any							       =							True
							lowerCamelCase_							:		Union[str, Any]							       =							False
							lowerCamelCase_							:		Tuple							       =							True
							lowerCamelCase_							:		Tuple							       =							model_class(__lowercase )
							model.to(__lowercase )
							model.eval()

							lowerCamelCase_							:		List[Any]							       =							self._prepare_for_class(__lowercase			,	__lowercase )
							lowerCamelCase_							:		List[str]							       =							inputs['input_ids'].shape[0]
							with torch.no_grad():
								lowerCamelCase_							:		Optional[int]							       =							model(**__lowercase )
							lowerCamelCase_							:		Dict							       =							outputs.attentions
							self.assertEqual(len(__lowercase )			,	self.model_tester.num_hidden_layers )

							# check that output_attentions also work using config
							del inputs_dict["output_attentions"]
							lowerCamelCase_							:		Any							       =							True
							lowerCamelCase_							:		Tuple							       =							model_class(__lowercase )
							model.to(__lowercase )
							model.eval()

							lowerCamelCase_							:		Dict							       =							self._prepare_for_class(__lowercase			,	__lowercase )
							lowerCamelCase_							:		Any							       =							inputs['input_ids'].shape[0]
							with torch.no_grad():
								lowerCamelCase_							:		Tuple							       =							model(**__lowercase )
							lowerCamelCase_							:		Union[str, Any]							       =							outputs.attentions
							self.assertEqual(len(__lowercase )			,	self.model_tester.num_hidden_layers )

							self.assertListEqual(
							    list(attentions[0].shape[-3:] )			,	[batch_size, seq_len, config.hidden_size]			,	)
							lowerCamelCase_							:		List[Any]							       =							len(__lowercase )

							# Check attention is always last and order is fine
							lowerCamelCase_							:		int							       =							True
							lowerCamelCase_							:		Any							       =							True
							lowerCamelCase_							:		Union[str, Any]							       =							model_class(__lowercase )
							model.to(__lowercase )
							model.eval()

							lowerCamelCase_							:		str							       =							self._prepare_for_class(__lowercase			,	__lowercase )
							lowerCamelCase_							:		Optional[Any]							       =							inputs['input_ids'].shape[0]
							with torch.no_grad():
								lowerCamelCase_							:		Union[str, Any]							       =							model(**__lowercase )

							lowerCamelCase_							:		List[Any]							       =							1
							self.assertEqual(out_len + added_hidden_states			,	len(__lowercase ) )

							lowerCamelCase_							:		Dict							       =							outputs.attentions

							self.assertEqual(len(__lowercase )			,	self.model_tester.num_hidden_layers )
							self.assertListEqual(
							    list(self_attentions[0].shape[-3:] )			,	[batch_size, seq_len, config.hidden_size]			,	)





					@slow
					def      lowercase_							(    self ):







						'''simple docstring'''




						for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							lowerCamelCase_							:		Optional[int]							       =							RwkvModel.from_pretrained(__lowercase )
							self.assertIsNotNone(__lowercase )


@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0       ,					reason="""See https://github.com/huggingface/transformers/pull/24204"""  )
@slow
class     __SCREAMING_SNAKE_CASE							(		unittest.TestCase  ):

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		List[str]							       =							'RWKV/rwkv-4-169m-pile'
						lowerCamelCase_							:		str							       =							AutoTokenizer.from_pretrained(self.model_id )

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[int]							       =							'Hello my name is Jasmine and I am a newbie to the'
						lowerCamelCase_							:		Tuple							       =							RwkvForCausalLM.from_pretrained(self.model_id ).to(__lowercase )

						lowerCamelCase_							:		int							       =							self.tokenizer('Hello my name is'			,	return_tensors='pt' ).input_ids.to(__lowercase )
						lowerCamelCase_							:		Tuple							       =							model.generate(__lowercase			,	max_new_tokens=10 )
						lowerCamelCase_							:		Union[str, Any]							       =							self.tokenizer.decode(output[0].tolist() )

						self.assertEqual(__lowercase			,	__lowercase )





					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		str							       =							'Hello my name is Jasmine and I am a newbie to the'

						lowerCamelCase_							:		str							       =							self.tokenizer('Hello my name is'			,	return_tensors='pt' ).input_ids.to(__lowercase )
						lowerCamelCase_							:		Tuple							       =							RwkvForCausalLM.from_pretrained(self.model_id			,	torch_dtype=torch.bfloataa ).to(__lowercase )

						lowerCamelCase_							:		List[str]							       =							model.generate(__lowercase			,	max_new_tokens=10 )
						lowerCamelCase_							:		Dict							       =							self.tokenizer.decode(output[0].tolist() )

						self.assertEqual(__lowercase			,	__lowercase )






import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class     __SCREAMING_SNAKE_CASE							(		lowerCamelCase  ):
					_lowerCamelCase		  =			(HeunDiscreteScheduler,)
					_lowerCamelCase		  =			10

					def      lowercase_							(    self			,	**__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		List[str]							       =							{
						    'num_train_timesteps': 1_100,
						    'beta_start': 0.00_01,
						    'beta_end': 0.02,
						    'beta_schedule': 'linear',
						}

						config.update(**__lowercase )
						return config

					def      lowercase_							(    self ):







						'''simple docstring'''




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

					def      lowercase_							(    self ):







						'''simple docstring'''




						for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]			,	[0.00_02, 0.0_02, 0.02] ):
							self.check_over_configs(beta_start=__lowercase			,	beta_end=__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




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

					def      lowercase_							(    self ):







						'''simple docstring'''




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

					def      lowercase_							(    self ):







						'''simple docstring'''




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

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Tuple							       =							self.scheduler_classes[0]
						lowerCamelCase_							:		List[str]							       =							self.get_scheduler_config()
						lowerCamelCase_							:		Union[str, Any]							       =							scheduler_class(**__lowercase )

						scheduler.set_timesteps(self.num_inference_steps )

						lowerCamelCase_							:		Optional[Any]							       =							self.dummy_model()
						lowerCamelCase_							:		Optional[int]							       =							self.dummy_sample_deter * scheduler.init_noise_sigma
						lowerCamelCase_							:		Dict							       =							sample.to(__lowercase )

						for i, t in enumerate(scheduler.timesteps ):
							lowerCamelCase_							:		Any							       =							scheduler.scale_model_input(__lowercase			,	__lowercase )

							lowerCamelCase_							:		int							       =							model(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Optional[int]							       =							scheduler.step(__lowercase			,	__lowercase			,	__lowercase )
							lowerCamelCase_							:		Tuple							       =							output.prev_sample

						lowerCamelCase_							:		str							       =							torch.sum(torch.abs(__lowercase ) )
						lowerCamelCase_							:		Tuple							       =							torch.mean(torch.abs(__lowercase ) )

						if torch_device in ["cpu", "mps"]:
							assert abs(result_sum.item() - 0.12_33 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3
						else:
							# CUDA
							assert abs(result_sum.item() - 0.12_33 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Any							       =							self.scheduler_classes[0]
						lowerCamelCase_							:		str							       =							self.get_scheduler_config(prediction_type='v_prediction' )
						lowerCamelCase_							:		Optional[Any]							       =							scheduler_class(**__lowercase )

						scheduler.set_timesteps(self.num_inference_steps )

						lowerCamelCase_							:		Any							       =							self.dummy_model()
						lowerCamelCase_							:		str							       =							self.dummy_sample_deter * scheduler.init_noise_sigma
						lowerCamelCase_							:		List[Any]							       =							sample.to(__lowercase )

						for i, t in enumerate(scheduler.timesteps ):
							lowerCamelCase_							:		int							       =							scheduler.scale_model_input(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Tuple							       =							model(__lowercase			,	__lowercase )

							lowerCamelCase_							:		int							       =							scheduler.step(__lowercase			,	__lowercase			,	__lowercase )
							lowerCamelCase_							:		Dict							       =							output.prev_sample

						lowerCamelCase_							:		Any							       =							torch.sum(torch.abs(__lowercase ) )
						lowerCamelCase_							:		List[Any]							       =							torch.mean(torch.abs(__lowercase ) )

						if torch_device in ["cpu", "mps"]:
							assert abs(result_sum.item() - 4.6934E-07 ) < 1E-2
							assert abs(result_mean.item() - 6.1112E-10 ) < 1E-3
						else:
							# CUDA
							assert abs(result_sum.item() - 4.693428650170972E-07 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							       =							self.scheduler_classes[0]
						lowerCamelCase_							:		int							       =							self.get_scheduler_config()
						lowerCamelCase_							:		Dict							       =							scheduler_class(**__lowercase )

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

						lowerCamelCase_							:		str							       =							self.dummy_model()
						lowerCamelCase_							:		Optional[int]							       =							self.dummy_sample_deter.to(__lowercase ) * scheduler.init_noise_sigma

						for t in scheduler.timesteps:
							lowerCamelCase_							:		Tuple							       =							scheduler.scale_model_input(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Tuple							       =							model(__lowercase			,	__lowercase )

							lowerCamelCase_							:		List[Any]							       =							scheduler.step(__lowercase			,	__lowercase			,	__lowercase )
							lowerCamelCase_							:		Union[str, Any]							       =							output.prev_sample

						lowerCamelCase_							:		Dict							       =							torch.sum(torch.abs(__lowercase ) )
						lowerCamelCase_							:		Any							       =							torch.mean(torch.abs(__lowercase ) )

						if str(__lowercase ).startswith('cpu' ):
							# The following sum varies between 148 and 156 on mps. Why?
							assert abs(result_sum.item() - 0.12_33 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3
						elif str(__lowercase ).startswith('mps' ):
							# Larger tolerance on mps
							assert abs(result_mean.item() - 0.00_02 ) < 1E-2
						else:
							# CUDA
							assert abs(result_sum.item() - 0.12_33 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3





					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							       =							self.scheduler_classes[0]
						lowerCamelCase_							:		Optional[Any]							       =							self.get_scheduler_config()
						lowerCamelCase_							:		Dict							       =							scheduler_class(**__lowercase			,	use_karras_sigmas=__lowercase )

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

						lowerCamelCase_							:		int							       =							self.dummy_model()
						lowerCamelCase_							:		Optional[Any]							       =							self.dummy_sample_deter.to(__lowercase ) * scheduler.init_noise_sigma
						lowerCamelCase_							:		Tuple							       =							sample.to(__lowercase )

						for t in scheduler.timesteps:
							lowerCamelCase_							:		str							       =							scheduler.scale_model_input(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Tuple							       =							model(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Dict							       =							scheduler.step(__lowercase			,	__lowercase			,	__lowercase )
							lowerCamelCase_							:		Dict							       =							output.prev_sample

						lowerCamelCase_							:		Tuple							       =							torch.sum(torch.abs(__lowercase ) )
						lowerCamelCase_							:		Optional[Any]							       =							torch.mean(torch.abs(__lowercase ) )

						assert abs(result_sum.item() - 0.0_00_15 ) < 1E-2
						assert abs(result_mean.item() - 1.9869554535034695E-07 ) < 1E-2
