






import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
     import jax.numpy as jnp

     from transformers.models.roberta_prelayernorm.modeling_flax_roberta_prelayernorm import (
         FlaxRobertaPreLayerNormForCausalLM,
         FlaxRobertaPreLayerNormForMaskedLM,
         FlaxRobertaPreLayerNormForMultipleChoice,
         FlaxRobertaPreLayerNormForQuestionAnswering,
         FlaxRobertaPreLayerNormForSequenceClassification,
         FlaxRobertaPreLayerNormForTokenClassification,
         FlaxRobertaPreLayerNormModel,
     )




class       snake_case__ (				unittest.TestCase):

       '''simple docstring'''



       def __init__(				self      ,    a__      ,    a__=13      ,    a__=7      ,    a__=True      ,    a__=True      ,    a__=True      ,    a__=True      ,    a__=99      ,    a__=32      ,    a__=5      ,    a__=4      ,    a__=37      ,    a__="gelu"      ,    a__=0.1      ,    a__=0.1      ,    a__=5_12      ,    a__=16      ,    a__=2      ,    a__=0.02      ,    a__=4      ,    )			->   int:






            '''simple docstring'''





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



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






            '''simple docstring'''





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

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

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

            __snake_case      :str 							=    RobertaPreLayerNormConfig(
                vocab_size=self.vocab_size      ,    hidden_size=self.hidden_size      ,    num_hidden_layers=self.num_hidden_layers      ,    num_attention_heads=self.num_attention_heads      ,    intermediate_size=self.intermediate_size      ,    hidden_act=self.hidden_act      ,    hidden_dropout_prob=self.hidden_dropout_prob      ,    attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,    max_position_embeddings=self.max_position_embeddings      ,    type_vocab_size=self.type_vocab_size      ,    is_decoder=a__      ,    initializer_range=self.initializer_range      ,    )

            return config, input_ids, token_type_ids, attention_mask



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :str 							=    self.prepare_config_and_inputs()
            __snake_case							,			__snake_case							,			__snake_case							,			__snake_case      :Optional[int] 							=    config_and_inputs
            __snake_case      :List[Any] 							=    {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
            return config, inputs_dict



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






            '''simple docstring'''





            __snake_case      :Union[str, Any] 							=    self.prepare_config_and_inputs()
            __snake_case							,			__snake_case							,			__snake_case							,			__snake_case      :int 							=    config_and_inputs

            __snake_case      :Optional[Any] 							=    True
            __snake_case      :List[Any] 							=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
            __snake_case      :List[str] 							=    ids_tensor([self.batch_size, self.seq_length]      ,    vocab_size=2     )

            return (
                config,
                input_ids,
                token_type_ids,
                encoder_hidden_states,
                encoder_attention_mask,
            )




@require_flax
# Copied from tests.models.roberta.test_modelling_flax_roberta.FlaxRobertaPreLayerNormModelTest with ROBERTA->ROBERTA_PRELAYERNORM,Roberta->RobertaPreLayerNorm,roberta-base->andreasmadsen/efficient_mlm_m0.40
class       snake_case__ (				lowercase_						,     unittest.TestCase):

       '''simple docstring'''
       lowerCamelCase				:					Tuple								=			True

       lowerCamelCase				:					Optional[int]								=			(
           (
               FlaxRobertaPreLayerNormModel,
               FlaxRobertaPreLayerNormForCausalLM,
               FlaxRobertaPreLayerNormForMaskedLM,
               FlaxRobertaPreLayerNormForSequenceClassification,
               FlaxRobertaPreLayerNormForTokenClassification,
               FlaxRobertaPreLayerNormForMultipleChoice,
               FlaxRobertaPreLayerNormForQuestionAnswering,
           )
           if is_flax_available()
           else ()
       )



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :Dict 							=    FlaxRobertaPreLayerNormModelTester(self     )



       @slow
       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            for model_class_name in self.all_model_classes:
                 __snake_case      :List[str] 							=    model_class_name.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""      ,    from_pt=a__     )
                 __snake_case      :Union[str, Any] 							=    model(np.ones((1, 1)     )     )
                 self.assertIsNotNone(a__     )




@require_flax
class       snake_case__ (				unittest.TestCase):

       '''simple docstring'''



       @slow
       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :Dict 							=    FlaxRobertaPreLayerNormForMaskedLM.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""      ,    from_pt=a__     )

            __snake_case      :List[Any] 							=    np.array([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]      ,    dtype=jnp.intaa     )
            __snake_case      :int 							=    model(a__     )[0]
            __snake_case      :str 							=    [1, 11, 5_02_65]
            self.assertEqual(list(output.shape     )      ,    a__     )
            # compare the actual values for a slice.
            __snake_case      :Optional[Any] 							=    np.array(
                [[[40.48_80, 18.01_99, -5.23_67], [-1.88_77, -4.08_85, 10.70_85], [-2.26_13, -5.61_10, 7.26_65]]]      ,    dtype=np.floataa     )
            self.assertTrue(np.allclose(output[:, :3, :3]      ,    a__      ,    atol=1e-4     )     )



       @slow
       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    FlaxRobertaPreLayerNormModel.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""      ,    from_pt=a__     )

            __snake_case      :Tuple 							=    np.array([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]      ,    dtype=jnp.intaa     )
            __snake_case      :Any 							=    model(a__     )[0]
            # compare the actual values for a slice.
            __snake_case      :str 							=    np.array(
                [[[0.02_08, -0.03_56, 0.02_37], [-0.15_69, -0.04_11, -0.26_26], [0.18_79, 0.01_25, -0.00_89]]]      ,    dtype=np.floataa     )
            self.assertTrue(np.allclose(output[:, :3, :3]      ,    a__      ,    atol=1e-4     )     )










import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class       snake_case__ (				lowercase_						,     unittest.TestCase):

       '''simple docstring'''
       lowerCamelCase				:					Tuple								=			KandinskyVaaPriorPipeline
       lowerCamelCase				:					Optional[Any]								=			["prompt"]
       lowerCamelCase				:					Any								=			["prompt", "negative_prompt"]
       lowerCamelCase				:					Tuple								=			[
           "num_images_per_prompt",
           "generator",
           "num_inference_steps",
           "latents",
           "negative_prompt",
           "guidance_scale",
           "output_type",
           "return_dict",
       ]
       lowerCamelCase				:					Any								=			False



       @property
       def        __lowercase    (				self     )			->   Union[str, Any]:






            '''simple docstring'''





            return 32



       @property
       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            return 32



       @property
       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            return self.time_input_dim



       @property
       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            return self.time_input_dim * 4



       @property
       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            return 1_00



       @property
       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :Tuple 							=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )
            return tokenizer



       @property
       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            torch.manual_seed(0     )
            __snake_case      :Dict 							=    CLIPTextConfig(
                bos_token_id=0      ,    eos_token_id=2      ,    hidden_size=self.text_embedder_hidden_size      ,    projection_dim=self.text_embedder_hidden_size      ,    intermediate_size=37      ,    layer_norm_eps=1e-05      ,    num_attention_heads=4      ,    num_hidden_layers=5      ,    pad_token_id=1      ,    vocab_size=10_00      ,    )
            return CLIPTextModelWithProjection(a__     )



       @property
       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            torch.manual_seed(0     )

            __snake_case      :List[Any] 							=    {
                """num_attention_heads""": 2,
                """attention_head_dim""": 12,
                """embedding_dim""": self.text_embedder_hidden_size,
                """num_layers""": 1,
            }

            __snake_case      :Union[str, Any] 							=    PriorTransformer(**a__     )
            # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
            __snake_case      :int 							=    nn.Parameter(torch.ones(model.clip_std.shape     )     )
            return model



       @property
       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            torch.manual_seed(0     )
            __snake_case      :List[Any] 							=    CLIPVisionConfig(
                hidden_size=self.text_embedder_hidden_size      ,    image_size=2_24      ,    projection_dim=self.text_embedder_hidden_size      ,    intermediate_size=37      ,    num_attention_heads=4      ,    num_channels=3      ,    num_hidden_layers=5      ,    patch_size=14      ,    )

            __snake_case      :Tuple 							=    CLIPVisionModelWithProjection(a__     )
            return model



       @property
       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    CLIPImageProcessor(
                crop_size=2_24      ,    do_center_crop=a__      ,    do_normalize=a__      ,    do_resize=a__      ,    image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]      ,    image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]      ,    resample=3      ,    size=2_24      ,    )

            return image_processor



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






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    self.dummy_prior
            __snake_case      :Tuple 							=    self.dummy_image_encoder
            __snake_case      :List[Any] 							=    self.dummy_text_encoder
            __snake_case      :List[str] 							=    self.dummy_tokenizer
            __snake_case      :int 							=    self.dummy_image_processor

            __snake_case      :Union[str, Any] 							=    UnCLIPScheduler(
                variance_type="""fixed_small_log"""      ,    prediction_type="""sample"""      ,    num_train_timesteps=10_00      ,    clip_sample=a__      ,    clip_sample_range=10.0      ,    )

            __snake_case      :Union[str, Any] 							=    {
                """prior""": prior,
                """image_encoder""": image_encoder,
                """text_encoder""": text_encoder,
                """tokenizer""": tokenizer,
                """scheduler""": scheduler,
                """image_processor""": image_processor,
            }

            return components



       def        __lowercase    (				self      ,    a__      ,    a__=0     )			->   Optional[Any]:






            '''simple docstring'''





            if str(a__     ).startswith("""mps"""     ):
                 __snake_case      :Any 							=    torch.manual_seed(a__     )
            else:
                 __snake_case      :Any 							=    torch.Generator(device=a__     ).manual_seed(a__     )
            __snake_case      :Union[str, Any] 							=    {
                """prompt""": """horse""",
                """generator""": generator,
                """guidance_scale""": 4.0,
                """num_inference_steps""": 2,
                """output_type""": """np""",
            }
            return inputs



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :str 							=    """cpu"""

            __snake_case      :Optional[Any] 							=    self.get_dummy_components()

            __snake_case      :List[Any] 							=    self.pipeline_class(**a__     )
            __snake_case      :Union[str, Any] 							=    pipe.to(a__     )

            pipe.set_progress_bar_config(disable=a__     )

            __snake_case      :List[Any] 							=    pipe(**self.get_dummy_inputs(a__     )     )
            __snake_case      :Any 							=    output.image_embeds

            __snake_case      :Any 							=    pipe(
                **self.get_dummy_inputs(a__     )      ,    return_dict=a__      ,    )[0]

            __snake_case      :Any 							=    image[0, -10:]
            __snake_case      :Any 							=    image_from_tuple[0, -10:]

            assert image.shape == (1, 32)

            __snake_case      :List[str] 							=    np.array(
                [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]     )

            assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2
            assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 1e-2



       @skip_mps
       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    torch_device == """cpu"""
            __snake_case      :List[str] 							=    True
            __snake_case      :Optional[Any] 							=    False

            self._test_inference_batch_single_identical(
                test_max_difference=a__      ,    relax_max_difference=a__      ,    test_mean_pixel_difference=a__      ,    )



       @skip_mps
       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            __snake_case      :Any 							=    torch_device == """cpu"""
            __snake_case      :str 							=    False

            self._test_attention_slicing_forward_pass(
                test_max_difference=a__      ,    test_mean_pixel_difference=a__      ,    )


