




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						:    Tuple      =  False




class  __lowerCAmelCase							:






       def __init__( self:       Tuple							,       _lowerCAmelCase:       Optional[int]							,       _lowerCAmelCase:       Dict=14							,       _lowerCAmelCase:       Any=7							,       _lowerCAmelCase:       str=True							,       _lowerCAmelCase:       int=False							,       _lowerCAmelCase:       str=True							,       _lowerCAmelCase:       Tuple=True							,       _lowerCAmelCase:       str=True							,       _lowerCAmelCase:       Optional[Any]=99							,       _lowerCAmelCase:       List[str]=32							,       _lowerCAmelCase:       Dict=5							,       _lowerCAmelCase:       Union[str, Any]=37							,       _lowerCAmelCase:       Dict="gelu"							,       _lowerCAmelCase:       Union[str, Any]=0.1							,       _lowerCAmelCase:       Any=0.1							,       _lowerCAmelCase:       List[str]=5_12							,       _lowerCAmelCase:       str=16							,       _lowerCAmelCase:       Optional[Any]=2							,       _lowerCAmelCase:       Tuple=3							,       _lowerCAmelCase:       str=4							,       _lowerCAmelCase:       int=None							,       ):
         lowercase     :Union[str, Any]   						=    parent
         lowercase     :Dict   						=    batch_size
         lowercase     :Optional[Any]   						=    seq_length
         lowercase     :Optional[int]   						=    is_training
         lowercase     :List[Any]   						=    use_token_type_ids
         lowercase     :Any   						=    use_input_mask
         lowercase     :str   						=    use_labels
         lowercase     :List[str]   						=    use_mc_token_ids
         lowercase     :List[Any]   						=    vocab_size
         lowercase     :Optional[Any]   						=    hidden_size
         lowercase     :Dict   						=    num_hidden_layers
         lowercase     :str   						=    intermediate_size
         lowercase     :Any   						=    hidden_act
         lowercase     :Tuple   						=    hidden_dropout_prob
         lowercase     :Tuple   						=    attention_probs_dropout_prob
         lowercase     :int   						=    max_position_embeddings
         lowercase     :List[Any]   						=    type_vocab_size
         lowercase     :str   						=    type_sequence_label_size
         lowercase     :Dict   						=    num_labels
         lowercase     :Optional[Any]   						=    num_choices
         lowercase     :int   						=    scope
         lowercase     :Dict   						=    vocab_size - 1
         lowercase     :List[str]   						=    vocab_size - 1
         lowercase     :List[Any]   						=    vocab_size - 1






       def        SCREAMING_SNAKE_CASE     ( self:       Tuple					):
         return RwkvConfig.from_pretrained("sgugger/rwkv-4-pile-7b"					)






       def        SCREAMING_SNAKE_CASE     ( self:       Any							,       _lowerCAmelCase:       str=False							,       _lowerCAmelCase:       int=False							,       _lowerCAmelCase:       Optional[Any]=False					):
         lowercase     :List[Any]   						=    ids_tensor([self.batch_size, self.seq_length]							,       self.vocab_size					)

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

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

         lowercase     :Optional[Any]   						=    None
         if self.use_mc_token_ids:
           lowercase     :List[Any]   						=    ids_tensor([self.batch_size, self.num_choices]							,       self.seq_length					)

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

         lowercase     :int   						=    self.get_config(
             gradient_checkpointing=_lowerCAmelCase							,       scale_attn_by_inverse_layer_idx=_lowerCAmelCase							,       reorder_and_upcast_attn=_lowerCAmelCase							,       )

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






       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]							,       _lowerCAmelCase:       Optional[int]=False							,       _lowerCAmelCase:       Tuple=False							,       _lowerCAmelCase:       int=False					):
         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=_lowerCAmelCase							,       bos_token_id=self.bos_token_id							,       eos_token_id=self.eos_token_id							,       pad_token_id=self.pad_token_id							,       gradient_checkpointing=_lowerCAmelCase							,       scale_attn_by_inverse_layer_idx=_lowerCAmelCase							,       reorder_and_upcast_attn=_lowerCAmelCase							,       )






       def        SCREAMING_SNAKE_CASE     ( self:       str					):
         lowercase     :Tuple   						=    self.get_config()
         lowercase     :Optional[Any]   						=    3_00
         return config






       def        SCREAMING_SNAKE_CASE     ( self:       int					):
         (
             (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     
         )     :Dict   						=    self.prepare_config_and_inputs()

         lowercase     :Optional[int]   						=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
         lowercase     :Union[str, Any]   						=    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        SCREAMING_SNAKE_CASE     ( self:       str							,       _lowerCAmelCase:       List[Any]							,       _lowerCAmelCase:       Optional[int]							,       _lowerCAmelCase:       Any							,       _lowerCAmelCase:       List[str]							,       _lowerCAmelCase:       str							,       *_lowerCAmelCase:       Tuple					):
         lowercase     :Tuple   						=    True
         lowercase     :List[str]   						=    RwkvModel(config=_lowerCAmelCase					)
         model.to(_lowerCAmelCase					)
         model.eval()

         lowercase     :List[Any]   						=    model(_lowerCAmelCase					)

         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        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]							,       _lowerCAmelCase:       List[Any]							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       Union[str, Any]							,       _lowerCAmelCase:       Optional[int]							,       *_lowerCAmelCase:       str					):
         lowercase     :Optional[int]   						=    RwkvForCausalLM(_lowerCAmelCase					)
         model.to(_lowerCAmelCase					)
         model.eval()

         lowercase     :Any   						=    model(_lowerCAmelCase							,       labels=_lowerCAmelCase					)
         self.parent.assertEqual(result.loss.shape							,       ()					)
         self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.seq_length, self.vocab_size)					)






       def        SCREAMING_SNAKE_CASE     ( self:       Tuple							,       _lowerCAmelCase:       Any							,       _lowerCAmelCase:       Any							,       _lowerCAmelCase:       Dict							,       _lowerCAmelCase:       List[str]							,       _lowerCAmelCase:       Optional[Any]							,       *_lowerCAmelCase:       List[Any]					):
         lowercase     :Any   						=    RwkvModel(config=_lowerCAmelCase					)
         model.to(_lowerCAmelCase					)
         model.eval()

         lowercase     :List[Any]   						=    model(_lowerCAmelCase					)
         lowercase     :Union[str, Any]   						=    outputs.last_hidden_state

         lowercase     :Optional[int]   						=    model(input_ids[:, :2]					)
         lowercase     :Dict   						=    outputs.last_hidden_state

         # Using the state computed on the first inputs, we will get the same output
         lowercase     :Union[str, Any]   						=    model(input_ids[:, 2:]							,       state=outputs.state					)
         lowercase     :Tuple   						=    outputs.last_hidden_state

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






       def        SCREAMING_SNAKE_CASE     ( self:       List[str]							,       _lowerCAmelCase:       Union[str, Any]							,       _lowerCAmelCase:       List[Any]							,       _lowerCAmelCase:       str							,       _lowerCAmelCase:       Tuple							,       _lowerCAmelCase:       Any							,       *_lowerCAmelCase:       Union[str, Any]							,       _lowerCAmelCase:       Union[str, Any]=False					):
         lowercase     :Union[str, Any]   						=    RwkvForCausalLM(_lowerCAmelCase					)
         model.to(_lowerCAmelCase					)
         if gradient_checkpointing:
           model.gradient_checkpointing_enable()

         lowercase     :Tuple   						=    model(_lowerCAmelCase							,       labels=_lowerCAmelCase					)
         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        SCREAMING_SNAKE_CASE     ( self:       Any					):
         lowercase     :List[Any]   						=    self.prepare_config_and_inputs()

         (
             (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     (
             lowercase
         )						,     
         )     :str   						=    config_and_inputs

         lowercase     :List[Any]   						=    {"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  __lowerCAmelCase							(				lowerCAmelCase   ,			lowerCAmelCase   ,			lowerCAmelCase   ,			unittest.TestCase):
       _a			       =    (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
       _a			       =    (
           {'''feature-extraction''': RwkvModel, '''text-generation''': RwkvForCausalLM} if is_torch_available() else {}
       )
       # all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
       _a			       =    False
       _a			       =    False
       _a			       =    False
       _a			       =    False
       _a			       =    False  # Rwkv does not support head masking






       def        SCREAMING_SNAKE_CASE     ( self:       List[Any]					):
         lowercase     :Tuple   						=    RwkvModelTester(self					)
         lowercase     :Union[str, Any]   						=    ConfigTester(
             self							,       config_class=_lowerCAmelCase							,       n_embd=37							,       common_properties=["hidden_size", "num_hidden_layers"]					)






       def        SCREAMING_SNAKE_CASE     ( self:       List[Any]							,       _lowerCAmelCase:       Tuple							,       _lowerCAmelCase:       int							,       _lowerCAmelCase:       Any=None					):
         if isinstance(_lowerCAmelCase							,       torch.Tensor					):
           lowercase						,     lowercase     :Dict   						=    member.max().item(), member.min().item()
         elif isinstance(_lowerCAmelCase							,       _lowerCAmelCase					) or isinstance(_lowerCAmelCase							,       _lowerCAmelCase					):
           lowercase						,     lowercase     :List[Any]   						=    max(_lowerCAmelCase					), min(_lowerCAmelCase					)

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

         lowercase						,     lowercase     :List[str]   						=    container

         lowercase     :Optional[Any]   						=    (min_value >= expected_min) and (max_value <= expected_max)

         if not is_inside_interval:
           lowercase     :int   						=    "%s not found in %s" % (safe_repr(_lowerCAmelCase					), safe_repr(_lowerCAmelCase					))
           self.fail(self._formatMessage(_lowerCAmelCase							,       _lowerCAmelCase					)					)






       def        SCREAMING_SNAKE_CASE     ( self:       str					):
         self.config_tester.run_common_tests()






       def        SCREAMING_SNAKE_CASE     ( self:       Dict					):
         lowercase     :int   						=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_rwkv_model(*_lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]					):
         lowercase     :Optional[int]   						=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_causl_lm(*_lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       str					):
         lowercase     :str   						=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_state_equivalency(*_lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       str					):
         lowercase						,     lowercase     :Any   						=    self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           lowercase     :Optional[int]   						=    model_class(config=_lowerCAmelCase					)
           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        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]					):
         lowercase						,     lowercase     :int   						=    self.model_tester.prepare_config_and_inputs_for_common()
         lowercase     :Optional[int]   						=    True

         lowercase     :Optional[int]   						=    getattr(self.model_tester							,       "seq_length"							,       _lowerCAmelCase					)

         for model_class in self.all_model_classes:
           lowercase     :int   						=    True
           lowercase     :Tuple   						=    False
           lowercase     :int   						=    True
           lowercase     :List[Any]   						=    model_class(_lowerCAmelCase					)
           model.to(_lowerCAmelCase					)
           model.eval()

           lowercase     :Any   						=    self._prepare_for_class(_lowerCAmelCase							,       _lowerCAmelCase					)
           lowercase     :Tuple   						=    inputs["input_ids"].shape[0]
           with torch.no_grad():
             lowercase     :List[str]   						=    model(**_lowerCAmelCase					)
           lowercase     :Optional[Any]   						=    outputs.attentions
           self.assertEqual(len(_lowerCAmelCase					)							,       self.model_tester.num_hidden_layers					)

           # check that output_attentions also work using config
           del inputs_dict["output_attentions"]
           lowercase     :List[Any]   						=    True
           lowercase     :str   						=    model_class(_lowerCAmelCase					)
           model.to(_lowerCAmelCase					)
           model.eval()

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

           self.assertListEqual(
               list(attentions[0].shape[-3:]					)							,       [batch_size, seq_len, config.hidden_size]							,       )
           lowercase     :str   						=    len(_lowerCAmelCase					)

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

           lowercase     :Optional[int]   						=    self._prepare_for_class(_lowerCAmelCase							,       _lowerCAmelCase					)
           lowercase     :Union[str, Any]   						=    inputs["input_ids"].shape[0]
           with torch.no_grad():
             lowercase     :int   						=    model(**_lowerCAmelCase					)

           lowercase     :Optional[int]   						=    1
           self.assertEqual(out_len + added_hidden_states							,       len(_lowerCAmelCase					)					)

           lowercase     :str   						=    outputs.attentions

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






       @slow
       def        SCREAMING_SNAKE_CASE     ( self:       List[Any]					):
         for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase     :int   						=    RwkvModel.from_pretrained(_lowerCAmelCase					)
           self.assertIsNotNone(_lowerCAmelCase					)







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






       def        SCREAMING_SNAKE_CASE     ( self:       List[str]					):
         lowercase     :List[str]   						=    "RWKV/rwkv-4-169m-pile"
         lowercase     :Union[str, Any]   						=    AutoTokenizer.from_pretrained(self.model_id					)






       def        SCREAMING_SNAKE_CASE     ( self:       Any					):
         lowercase     :Optional[Any]   						=    "Hello my name is Jasmine and I am a newbie to the"
         lowercase     :int   						=    RwkvForCausalLM.from_pretrained(self.model_id					).to(_lowerCAmelCase					)

         lowercase     :List[str]   						=    self.tokenizer("Hello my name is"							,       return_tensors="pt"					).input_ids.to(_lowerCAmelCase					)
         lowercase     :Dict   						=    model.generate(_lowerCAmelCase							,       max_new_tokens=10					)
         lowercase     :List[Any]   						=    self.tokenizer.decode(output[0].tolist()					)

         self.assertEqual(_lowerCAmelCase							,       _lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       List[Any]					):
         lowercase     :List[str]   						=    "Hello my name is Jasmine and I am a newbie to the"

         lowercase     :Optional[int]   						=    self.tokenizer("Hello my name is"							,       return_tensors="pt"					).input_ids.to(_lowerCAmelCase					)
         lowercase     :Optional[int]   						=    RwkvForCausalLM.from_pretrained(self.model_id							,       torch_dtype=torch.bfloataa					).to(_lowerCAmelCase					)

         lowercase     :Dict   						=    model.generate(_lowerCAmelCase							,       max_new_tokens=10					)
         lowercase     :List[Any]   						=    self.tokenizer.decode(output[0].tolist()					)

         self.assertEqual(_lowerCAmelCase							,       _lowerCAmelCase					)









import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()




class  __lowerCAmelCase							(				unittest.TestCase):
       _a			       =    '''hf-internal-testing/tiny-stable-diffusion-pipe'''
       _a			       =    1
       _a			       =    77
       _a			       =    32
       _a			       =    4
       _a			       =    _a			       =    64
       _a			       =    torch.manual_seed(0)






       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]							,       _lowerCAmelCase:       int=0.99_99					):
         lowercase     :Optional[int]   						=    UNetaDConditionModel.from_pretrained(self.model_id							,       subfolder="unet"					)
         lowercase     :Union[str, Any]   						=    unet.to(_lowerCAmelCase					)
         lowercase     :Any   						=    EMAModel(unet.parameters()							,       decay=_lowerCAmelCase							,       model_cls=_lowerCAmelCase							,       model_config=unet.config					)
         return unet, ema_unet






       def        SCREAMING_SNAKE_CASE     ( self:       List[str]					):
         lowercase     :Optional[Any]   						=    torch.randn(
             self.batch_size							,       self.num_in_channels							,       self.latent_height							,       self.latent_width							,       generator=self.generator					).to(_lowerCAmelCase					)
         lowercase     :Dict   						=    torch.randint(0							,       10_00							,       size=(self.batch_size,)							,       generator=self.generator					).to(_lowerCAmelCase					)
         lowercase     :List[Any]   						=    torch.randn(
             self.batch_size							,       self.prompt_length							,       self.text_encoder_hidden_dim							,       generator=self.generator					).to(_lowerCAmelCase					)
         return noisy_latents, timesteps, encoder_hidden_states






       def        SCREAMING_SNAKE_CASE     ( self:       Optional[int]							,       _lowerCAmelCase:       Optional[Any]					):
         lowercase     :Any   						=    {}
         for k, param in unet.state_dict().items():
           lowercase     :Optional[int]   						=    torch.randn_like(_lowerCAmelCase					) + (param * torch.randn_like(_lowerCAmelCase					))
           updated_state_dict.update({k: updated_param}					)
         unet.load_state_dict(_lowerCAmelCase					)
         return unet






       def        SCREAMING_SNAKE_CASE     ( self:       int					):
         lowercase						,     lowercase     :List[str]   						=    self.get_models()
         # Take the first (hypothetical) EMA step.
         ema_unet.step(unet.parameters()					)
         assert ema_unet.optimization_step == 1

         # Take two more.
         for _ in range(2					):
           ema_unet.step(unet.parameters()					)
         assert ema_unet.optimization_step == 3






       def        SCREAMING_SNAKE_CASE     ( self:       Any					):
         lowercase						,     lowercase     :str   						=    self.get_models()
         # Since the `unet` is not being updated (i.e., backprop'd)
         # there won't be any difference between the `params` of `unet`
         # and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
         ema_unet.step(unet.parameters()					)
         lowercase     :Tuple   						=    list(unet.parameters()					)
         for s_param, param in zip(ema_unet.shadow_params							,       _lowerCAmelCase					):
           assert torch.allclose(_lowerCAmelCase							,       _lowerCAmelCase					)

         # The above holds true even if we call `ema.step()` multiple times since
         # `unet` params are still not being updated.
         for _ in range(4					):
           ema_unet.step(unet.parameters()					)
         for s_param, param in zip(ema_unet.shadow_params							,       _lowerCAmelCase					):
           assert torch.allclose(_lowerCAmelCase							,       _lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       Dict					):
         lowercase						,     lowercase     :Optional[int]   						=    self.get_models()
         # Here we simulate the parameter updates for `unet`. Since there might
         # be some parameters which are initialized to zero we take extra care to
         # initialize their values to something non-zero before the multiplication.
         lowercase     :List[str]   						=    self.simulate_backprop(_lowerCAmelCase					)

         # Take the EMA step.
         ema_unet.step(unet_pseudo_updated_step_one.parameters()					)

         # Now the EMA'd parameters won't be equal to the original model parameters.
         lowercase     :str   						=    list(unet_pseudo_updated_step_one.parameters()					)
         for s_param, param in zip(ema_unet.shadow_params							,       _lowerCAmelCase					):
           assert ~torch.allclose(_lowerCAmelCase							,       _lowerCAmelCase					)

         # Ensure this is the case when we take multiple EMA steps.
         for _ in range(4					):
           ema_unet.step(unet.parameters()					)
         for s_param, param in zip(ema_unet.shadow_params							,       _lowerCAmelCase					):
           assert ~torch.allclose(_lowerCAmelCase							,       _lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       Optional[int]					):
         # If we call EMA step after a backpropagation consecutively for two times,
         # the shadow params from those two steps should be different.
         lowercase						,     lowercase     :Optional[Any]   						=    self.get_models()

         # First backprop + EMA
         lowercase     :Union[str, Any]   						=    self.simulate_backprop(_lowerCAmelCase					)
         ema_unet.step(unet_step_one.parameters()					)
         lowercase     :int   						=    ema_unet.shadow_params

         # Second backprop + EMA
         lowercase     :Optional[int]   						=    self.simulate_backprop(_lowerCAmelCase					)
         ema_unet.step(unet_step_two.parameters()					)
         lowercase     :Optional[int]   						=    ema_unet.shadow_params

         for step_one, step_two in zip(_lowerCAmelCase							,       _lowerCAmelCase					):
           assert ~torch.allclose(_lowerCAmelCase							,       _lowerCAmelCase					)






       def        SCREAMING_SNAKE_CASE     ( self:       Union[str, Any]					):
         # If there's no decay even if there are backprops, EMA steps
         # won't take any effect i.e., the shadow params would remain the
         # same.
         lowercase						,     lowercase     :Any   						=    self.get_models(decay=0.0					)
         lowercase     :List[Any]   						=    self.simulate_backprop(_lowerCAmelCase					)
         ema_unet.step(unet_step_one.parameters()					)
         lowercase     :List[str]   						=    ema_unet.shadow_params

         lowercase     :str   						=    self.simulate_backprop(_lowerCAmelCase					)
         ema_unet.step(unet_step_two.parameters()					)
         lowercase     :List[str]   						=    ema_unet.shadow_params

         for step_one, step_two in zip(_lowerCAmelCase							,       _lowerCAmelCase					):
           assert torch.allclose(_lowerCAmelCase							,       _lowerCAmelCase					)






       @skip_mps
       def        SCREAMING_SNAKE_CASE     ( self:       List[Any]					):
         lowercase						,     lowercase     :Optional[Any]   						=    self.get_models()
         lowercase						,     lowercase						,     lowercase     :str   						=    self.get_dummy_inputs()

         with tempfile.TemporaryDirectory() as tmpdir:
           ema_unet.save_pretrained(_lowerCAmelCase					)
           lowercase     :Optional[int]   						=    UNetaDConditionModel.from_pretrained(_lowerCAmelCase							,       model_cls=_lowerCAmelCase					)
           lowercase     :Tuple   						=    loaded_unet.to(unet.device					)

         # Since no EMA step has been performed the outputs should match.
         lowercase     :Dict   						=    unet(_lowerCAmelCase							,       _lowerCAmelCase							,       _lowerCAmelCase					).sample
         lowercase     :Any   						=    loaded_unet(_lowerCAmelCase							,       _lowerCAmelCase							,       _lowerCAmelCase					).sample

         assert torch.allclose(_lowerCAmelCase							,       _lowerCAmelCase							,       atol=1e-4					)



