import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
     from transformers.models.electra.modeling_flax_electra import (
         FlaxElectraForCausalLM,
         FlaxElectraForMaskedLM,
         FlaxElectraForMultipleChoice,
         FlaxElectraForPreTraining,
         FlaxElectraForQuestionAnswering,
         FlaxElectraForSequenceClassification,
         FlaxElectraForTokenClassification,
         FlaxElectraModel,
     )







class      SCREAMING_SNAKE_CASE__					(  unittest.TestCase    ):

       def __init__(							self					:     Tuple			,			_lowerCAmelCase					:     List[str]			,			_lowerCAmelCase					:     Any=13			,			_lowerCAmelCase					:     Union[str, Any]=7			,			_lowerCAmelCase					:     Optional[Any]=True			,			_lowerCAmelCase					:     List[Any]=True			,			_lowerCAmelCase					:     List[str]=True			,			_lowerCAmelCase					:     Optional[int]=True			,			_lowerCAmelCase					:     str=99			,			_lowerCAmelCase					:     List[str]=24			,			_lowerCAmelCase					:     Optional[Any]=32			,			_lowerCAmelCase					:     Any=5			,			_lowerCAmelCase					:     Tuple=4			,			_lowerCAmelCase					:     str=37			,			_lowerCAmelCase					:     Optional[int]="gelu"			,			_lowerCAmelCase					:     Dict=0.1			,			_lowerCAmelCase					:     List[str]=0.1			,			_lowerCAmelCase					:     int=5_12			,			_lowerCAmelCase					:     int=16			,			_lowerCAmelCase					:     Dict=2			,			_lowerCAmelCase					:     str=0.02			,			_lowerCAmelCase					:     List[Any]=4			,			):
           __snake_case       :				str 			=    parent
           __snake_case       :				List[str] 			=    batch_size
           __snake_case       :				str 			=    seq_length
           __snake_case       :				int 			=    is_training
           __snake_case       :				Optional[Any] 			=    use_attention_mask
           __snake_case       :				str 			=    use_token_type_ids
           __snake_case       :				List[str] 			=    use_labels
           __snake_case       :				List[Any] 			=    vocab_size
           __snake_case       :				Optional[Any] 			=    hidden_size
           __snake_case       :				Any 			=    embedding_size
           __snake_case       :				Optional[Any] 			=    num_hidden_layers
           __snake_case       :				List[Any] 			=    num_attention_heads
           __snake_case       :				str 			=    intermediate_size
           __snake_case       :				str 			=    hidden_act
           __snake_case       :				int 			=    hidden_dropout_prob
           __snake_case       :				Dict 			=    attention_probs_dropout_prob
           __snake_case       :				Dict 			=    max_position_embeddings
           __snake_case       :				List[str] 			=    type_vocab_size
           __snake_case       :				Optional[Any] 			=    type_sequence_label_size
           __snake_case       :				Optional[int] 			=    initializer_range
           __snake_case       :				Optional[int] 			=    num_choices

       def 			snake_case__    (							self					:     Any			):
           __snake_case       :				Optional[int] 			=    ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size			)

           __snake_case       :				Union[str, Any] 			=    None
           if self.use_attention_mask:
               __snake_case       :				Any 			=    random_attention_mask([self.batch_size, self.seq_length]			)

           __snake_case       :				List[str] 			=    None
           if self.use_token_type_ids:
               __snake_case       :				Dict 			=    ids_tensor([self.batch_size, self.seq_length]			,			self.type_vocab_size			)

           __snake_case       :				Any 			=    ElectraConfig(
               vocab_size=self.vocab_size			,			hidden_size=self.hidden_size			,			embedding_size=self.embedding_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			,			initializer_range=self.initializer_range			,			)

           return config, input_ids, token_type_ids, attention_mask





       def 			snake_case__    (							self					:     Dict			):
           __snake_case       :				Any 			=    self.prepare_config_and_inputs()
           __snake_case			, __snake_case			, __snake_case			, __snake_case       :				Optional[Any] 			=    config_and_inputs
           __snake_case       :				Optional[Any] 			=    {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
           return config, inputs_dict







@require_flax
class      SCREAMING_SNAKE_CASE__					(  __UpperCamelCase   ,							unittest.TestCase    ):
       A							:	int          =	True

       A							:	Tuple          =	(
           (
               FlaxElectraModel,
               FlaxElectraForCausalLM,
               FlaxElectraForMaskedLM,
               FlaxElectraForPreTraining,
               FlaxElectraForTokenClassification,
               FlaxElectraForQuestionAnswering,
               FlaxElectraForMultipleChoice,
               FlaxElectraForSequenceClassification,
           )
           if is_flax_available()
           else ()
       )

       def 			snake_case__    (							self					:     Dict			):
           __snake_case       :				Dict 			=    FlaxElectraModelTester(self			)





       @slow
       def 			snake_case__    (							self					:     Optional[Any]			):
           for model_class_name in self.all_model_classes:
               if model_class_name == FlaxElectraForMaskedLM:
                   __snake_case       :				Tuple 			=    model_class_name.from_pretrained("""google/electra-small-generator"""			)
               else:
                   __snake_case       :				str 			=    model_class_name.from_pretrained("""google/electra-small-discriminator"""			)
               __snake_case       :				int 			=    model(np.ones((1, 1)			)			)
               self.assertIsNotNone(_lowerCAmelCase			)




import tempfile
import unittest

from diffusers import (
    DDIMScheduler,
    DDPMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    PNDMScheduler,
    logging,
)
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils.testing_utils import CaptureLogger







class      SCREAMING_SNAKE_CASE__					(  __UpperCamelCase    ):
       A							:	str          =	"config.json"

       @register_to_config
       def __init__(							self					:     Any			,			_lowerCAmelCase					:     Tuple=2			,			_lowerCAmelCase					:     List[Any]=5			,			_lowerCAmelCase					:     int=(2, 5)			,			_lowerCAmelCase					:     Dict="for diffusion"			,			_lowerCAmelCase					:     List[str]=[1, 3]			,			):
           pass







class      SCREAMING_SNAKE_CASE__					(  __UpperCamelCase    ):
       A							:	Optional[int]          =	"config.json"

       @register_to_config
       def __init__(							self					:     int			,			_lowerCAmelCase					:     Dict=2			,			_lowerCAmelCase					:     Union[str, Any]=5			,			_lowerCAmelCase					:     Any=(2, 5)			,			_lowerCAmelCase					:     Optional[int]="for diffusion"			,			_lowerCAmelCase					:     int=[1, 3]			,			):
           pass







class      SCREAMING_SNAKE_CASE__					(  __UpperCamelCase    ):
       A							:	str          =	"config.json"

       @register_to_config
       def __init__(							self					:     Any			,			_lowerCAmelCase					:     Optional[int]=2			,			_lowerCAmelCase					:     Optional[int]=5			,			_lowerCAmelCase					:     int=(2, 5)			,			_lowerCAmelCase					:     Tuple="for diffusion"			,			_lowerCAmelCase					:     Any=[1, 3]			,			_lowerCAmelCase					:     str=[1, 3]			,			):
           pass







class      SCREAMING_SNAKE_CASE__					(  __UpperCamelCase    ):
       A							:	Union[str, Any]          =	"config.json"

       @register_to_config
       def __init__(							self					:     Any			,			_lowerCAmelCase					:     Dict=2			,			_lowerCAmelCase					:     Optional[int]=5			,			_lowerCAmelCase					:     Optional[Any]=(2, 5)			,			_lowerCAmelCase					:     int="for diffusion"			,			_lowerCAmelCase					:     Dict=[1, 5]			,			_lowerCAmelCase					:     Optional[Any]=[5, 4]			,			):
           pass







class      SCREAMING_SNAKE_CASE__					(  unittest.TestCase    ):

       def 			snake_case__    (							self					:     List[str]			):
           with self.assertRaises(_lowerCAmelCase			):
               ConfigMixin.load_config("""dummy_path"""			)

       def 			snake_case__    (							self					:     List[Any]			):
           __snake_case       :				List[Any] 			=    SampleObject()
           __snake_case       :				Optional[Any] 			=    obj.config
           assert config["a"] == 2
           assert config["b"] == 5
           assert config["c"] == (2, 5)
           assert config["d"] == "for diffusion"
           assert config["e"] == [1, 3]

           # init ignore private arguments
           __snake_case       :				Optional[Any] 			=    SampleObject(_name_or_path="""lalala"""			)
           __snake_case       :				List[Any] 			=    obj.config
           assert config["a"] == 2
           assert config["b"] == 5
           assert config["c"] == (2, 5)
           assert config["d"] == "for diffusion"
           assert config["e"] == [1, 3]

           # can override default
           __snake_case       :				int 			=    SampleObject(c=6			)
           __snake_case       :				Union[str, Any] 			=    obj.config
           assert config["a"] == 2
           assert config["b"] == 5
           assert config["c"] == 6
           assert config["d"] == "for diffusion"
           assert config["e"] == [1, 3]

           # can use positional arguments.
           __snake_case       :				List[str] 			=    SampleObject(1			,			c=6			)
           __snake_case       :				str 			=    obj.config
           assert config["a"] == 1
           assert config["b"] == 5
           assert config["c"] == 6
           assert config["d"] == "for diffusion"
           assert config["e"] == [1, 3]

       def 			snake_case__    (							self					:     Union[str, Any]			):
           __snake_case       :				Union[str, Any] 			=    SampleObject()
           __snake_case       :				Optional[Any] 			=    obj.config

           assert config["a"] == 2
           assert config["b"] == 5
           assert config["c"] == (2, 5)
           assert config["d"] == "for diffusion"
           assert config["e"] == [1, 3]

           with tempfile.TemporaryDirectory() as tmpdirname:
               obj.save_config(_lowerCAmelCase			)
               __snake_case       :				List[Any] 			=    SampleObject.from_config(SampleObject.load_config(_lowerCAmelCase			)			)
               __snake_case       :				Union[str, Any] 			=    new_obj.config

           # unfreeze configs
           __snake_case       :				Optional[Any] 			=    dict(_lowerCAmelCase			)
           __snake_case       :				Optional[Any] 			=    dict(_lowerCAmelCase			)

           assert config.pop("""c"""			) == (2, 5)  # instantiated as tuple
           assert new_config.pop("""c"""			) == [2, 5]  # saved & loaded as list because of json
           config.pop("""_use_default_values"""			)
           assert config == new_config

       def 			snake_case__    (							self					:     Tuple			):
           __snake_case       :				Optional[Any] 			=    logging.get_logger("""diffusers.configuration_utils"""			)
           # 30 for warning
           logger.setLevel(30			)

           with CaptureLogger(_lowerCAmelCase			) as cap_logger:
               __snake_case       :				List[Any] 			=    DDIMScheduler.from_pretrained(
                   """hf-internal-testing/tiny-stable-diffusion-torch"""			,			subfolder="""scheduler"""			)

           assert ddim.__class__ == DDIMScheduler
           # no warning should be thrown
           assert cap_logger.out == ""

       def 			snake_case__    (							self					:     Optional[int]			):
           __snake_case       :				int 			=    logging.get_logger("""diffusers.configuration_utils"""			)
           # 30 for warning
           logger.setLevel(30			)

           with CaptureLogger(_lowerCAmelCase			) as cap_logger:
               __snake_case       :				Optional[int] 			=    EulerDiscreteScheduler.from_pretrained(
                   """hf-internal-testing/tiny-stable-diffusion-torch"""			,			subfolder="""scheduler"""			)

           assert euler.__class__ == EulerDiscreteScheduler
           # no warning should be thrown
           assert cap_logger.out == ""

       def 			snake_case__    (							self					:     Tuple			):
           __snake_case       :				int 			=    logging.get_logger("""diffusers.configuration_utils"""			)
           # 30 for warning
           logger.setLevel(30			)

           with CaptureLogger(_lowerCAmelCase			) as cap_logger:
               __snake_case       :				Any 			=    EulerAncestralDiscreteScheduler.from_pretrained(
                   """hf-internal-testing/tiny-stable-diffusion-torch"""			,			subfolder="""scheduler"""			)

           assert euler.__class__ == EulerAncestralDiscreteScheduler
           # no warning should be thrown
           assert cap_logger.out == ""

       def 			snake_case__    (							self					:     List[Any]			):
           __snake_case       :				str 			=    logging.get_logger("""diffusers.configuration_utils"""			)
           # 30 for warning
           logger.setLevel(30			)

           with CaptureLogger(_lowerCAmelCase			) as cap_logger:
               __snake_case       :				Optional[int] 			=    PNDMScheduler.from_pretrained(
                   """hf-internal-testing/tiny-stable-diffusion-torch"""			,			subfolder="""scheduler"""			)

           assert pndm.__class__ == PNDMScheduler
           # no warning should be thrown
           assert cap_logger.out == ""

       def 			snake_case__    (							self					:     Tuple			):
           __snake_case       :				int 			=    logging.get_logger("""diffusers.configuration_utils"""			)
           # 30 for warning
           logger.setLevel(30			)

           with CaptureLogger(_lowerCAmelCase			) as cap_logger:
               __snake_case       :				List[str] 			=    DDPMScheduler.from_pretrained(
                   """hf-internal-testing/tiny-stable-diffusion-torch"""			,			subfolder="""scheduler"""			,			prediction_type="""sample"""			,			beta_end=8			,			)

           with CaptureLogger(_lowerCAmelCase			) as cap_logger_a:
               __snake_case       :				Optional[Any] 			=    DDPMScheduler.from_pretrained("""google/ddpm-celebahq-256"""			,			beta_start=88			)

           assert ddpm.__class__ == DDPMScheduler
           assert ddpm.config.prediction_type == "sample"
           assert ddpm.config.beta_end == 8
           assert ddpm_a.config.beta_start == 88

           # no warning should be thrown
           assert cap_logger.out == ""
           assert cap_logger_a.out == ""

       def 			snake_case__    (							self					:     List[Any]			):
           __snake_case       :				List[Any] 			=    logging.get_logger("""diffusers.configuration_utils"""			)
           # 30 for warning
           logger.setLevel(30			)

           with CaptureLogger(_lowerCAmelCase			) as cap_logger:
               __snake_case       :				str 			=    DPMSolverMultistepScheduler.from_pretrained(
                   """hf-internal-testing/tiny-stable-diffusion-torch"""			,			subfolder="""scheduler"""			)

           assert dpm.__class__ == DPMSolverMultistepScheduler
           # no warning should be thrown
           assert cap_logger.out == ""





       def 			snake_case__    (							self					:     List[Any]			):
           # let's first save a config that should be in the form
           #    a=2,
           #    b=5,
           #    c=(2, 5),
           #    d="for diffusion",
           #    e=[1, 3],

           __snake_case       :				Optional[Any] 			=    SampleObject()

           __snake_case       :				Union[str, Any] 			=    {k: v for k, v in config.config.items() if not k.startswith("""_"""			)}

           # make sure that default config has all keys in `_use_default_values`
           assert set(config_dict.keys()			) == set(config.config._use_default_values			)

           with tempfile.TemporaryDirectory() as tmpdirname:
               config.save_config(_lowerCAmelCase			)

               # now loading it with SampleObject2 should put f into `_use_default_values`
               __snake_case       :				Union[str, Any] 			=    SampleObjecta.from_config(_lowerCAmelCase			)

               assert "f" in config._use_default_values
               assert config.f == [1, 3]

           # now loading the config, should **NOT** use [1, 3] for `f`, but the default [1, 4] value
           # **BECAUSE** it is part of `config._use_default_values`
           __snake_case       :				Optional[int] 			=    SampleObjecta.from_config(config.config			)
           assert new_config.f == [5, 4]

           config.config._use_default_values.pop()
           __snake_case       :				Optional[Any] 			=    SampleObjecta.from_config(config.config			)
           assert new_config_a.f == [1, 3]

           # Nevertheless "e" should still be correctly loaded to [1, 3] from SampleObject2 instead of defaulting to [1, 5]
           assert new_config_a.e == [1, 3]



