





'''simple docstring'''






import gc
import random
import unittest

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

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class 					UpperCAmelCase	(     _lowercase				,						unittest.TestCase  ):
      _lowercase:     Optional[Any]			=  UnCLIPImageVariationPipeline
      _lowercase:     Dict			=  IMAGE_VARIATION_PARAMS - {'''height''', '''width''', '''guidance_scale'''}
      _lowercase:     List[str]			=  IMAGE_VARIATION_BATCH_PARAMS

      _lowercase:     Union[str, Any]			=  [
          '''generator''',
          '''return_dict''',
          '''decoder_num_inference_steps''',
          '''super_res_num_inference_steps''',
      ]
      _lowercase:     Any			=  False




      @property
      def 							lowercase__       (     self     :					Optional[Any]    )							-> Union[str, Any]:
            return 32




      @property
      def 							lowercase__       (     self     :					int    )							-> Dict:
            return 32




      @property
      def 							lowercase__       (     self     :					int    )							-> Tuple:
            return self.time_input_dim




      @property
      def 							lowercase__       (     self     :					Optional[int]    )							-> Tuple:
            return self.time_input_dim * 4




      @property
      def 							lowercase__       (     self     :					int    )							-> str:
            return 1_00




      @property
      def 							lowercase__       (     self     :					Optional[Any]    )							-> Dict:
            _lowerCAmelCase					=				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )
            return tokenizer




      @property
      def 							lowercase__       (     self     :					Optional[int]    )							-> int:
            torch.manual_seed(0    )
            _lowerCAmelCase					=				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-0_5     ,     num_attention_heads=4     ,     num_hidden_layers=5     ,     pad_token_id=1     ,     vocab_size=10_00     ,     )
            return CLIPTextModelWithProjection(_snake_case    )




      @property
      def 							lowercase__       (     self     :					Union[str, Any]    )							-> List[Any]:
            torch.manual_seed(0    )
            _lowerCAmelCase					=				CLIPVisionConfig(
                hidden_size=self.text_embedder_hidden_size     ,     projection_dim=self.text_embedder_hidden_size     ,     num_hidden_layers=5     ,     num_attention_heads=4     ,     image_size=32     ,     intermediate_size=37     ,     patch_size=1     ,     )
            return CLIPVisionModelWithProjection(_snake_case    )




      @property
      def 							lowercase__       (     self     :					Dict    )							-> Tuple:
            torch.manual_seed(0    )

            _lowerCAmelCase					=				{
                """clip_embeddings_dim""": self.text_embedder_hidden_size,
                """time_embed_dim""": self.time_embed_dim,
                """cross_attention_dim""": self.cross_attention_dim,
            }

            _lowerCAmelCase					=				UnCLIPTextProjModel(**_snake_case    )
            return model




      @property
      def 							lowercase__       (     self     :					Optional[int]    )							-> List[Any]:
            torch.manual_seed(0    )

            _lowerCAmelCase					=				{
                """sample_size""": 32,
                # RGB in channels
                """in_channels""": 3,
                # Out channels is double in channels because predicts mean and variance
                """out_channels""": 6,
                """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
                """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
                """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
                """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
                """layers_per_block""": 1,
                """cross_attention_dim""": self.cross_attention_dim,
                """attention_head_dim""": 4,
                """resnet_time_scale_shift""": """scale_shift""",
                """class_embed_type""": """identity""",
            }

            _lowerCAmelCase					=				UNetaDConditionModel(**_snake_case    )
            return model




      @property
      def 							lowercase__       (     self     :					str    )							-> Optional[int]:
            return {
                "sample_size": 64,
                "layers_per_block": 1,
                "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
                "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
                "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
                "in_channels": 6,
                "out_channels": 3,
            }




      @property
      def 							lowercase__       (     self     :					Dict    )							-> Tuple:
            torch.manual_seed(0    )

            _lowerCAmelCase					=				UNetaDModel(**self.dummy_super_res_kwargs    )
            return model




      @property
      def 							lowercase__       (     self     :					Tuple    )							-> Optional[int]:
            torch.manual_seed(1    )

            _lowerCAmelCase					=				UNetaDModel(**self.dummy_super_res_kwargs    )
            return model




      def 							lowercase__       (     self     :					Union[str, Any]    )							-> Optional[int]:
            _lowerCAmelCase					=				self.dummy_decoder
            _lowerCAmelCase					=				self.dummy_text_proj
            _lowerCAmelCase					=				self.dummy_text_encoder
            _lowerCAmelCase					=				self.dummy_tokenizer
            _lowerCAmelCase					=				self.dummy_super_res_first
            _lowerCAmelCase					=				self.dummy_super_res_last

            _lowerCAmelCase					=				UnCLIPScheduler(
                variance_type="""learned_range"""     ,     prediction_type="""epsilon"""     ,     num_train_timesteps=10_00     ,     )

            _lowerCAmelCase					=				UnCLIPScheduler(
                variance_type="""fixed_small_log"""     ,     prediction_type="""epsilon"""     ,     num_train_timesteps=10_00     ,     )

            _lowerCAmelCase					=				CLIPImageProcessor(crop_size=32     ,     size=32    )

            _lowerCAmelCase					=				self.dummy_image_encoder

            return {
                "decoder": decoder,
                "text_encoder": text_encoder,
                "tokenizer": tokenizer,
                "text_proj": text_proj,
                "feature_extractor": feature_extractor,
                "image_encoder": image_encoder,
                "super_res_first": super_res_first,
                "super_res_last": super_res_last,
                "decoder_scheduler": decoder_scheduler,
                "super_res_scheduler": super_res_scheduler,
            }




      def 							lowercase__       (     self     :					Optional[int]     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					str=0     ,     __snake_case     :					List[str]=True    )							-> Union[str, Any]:
            _lowerCAmelCase					=				floats_tensor((1, 3, 32, 32)     ,     rng=random.Random(_snake_case    )    ).to(_snake_case    )
            if str(_snake_case    ).startswith("""mps"""    ):
                  _lowerCAmelCase					=				torch.manual_seed(_snake_case    )
            else:
                  _lowerCAmelCase					=				torch.Generator(device=_snake_case    ).manual_seed(_snake_case    )

            if pil_image:
                  _lowerCAmelCase					=				input_image * 0.5 + 0.5
                  _lowerCAmelCase					=				input_image.clamp(0     ,     1    )
                  _lowerCAmelCase					=				input_image.cpu().permute(0     ,     2     ,     3     ,     1    ).float().numpy()
                  _lowerCAmelCase					=				DiffusionPipeline.numpy_to_pil(_snake_case    )[0]

            return {
                "image": input_image,
                "generator": generator,
                "decoder_num_inference_steps": 2,
                "super_res_num_inference_steps": 2,
                "output_type": "np",
            }




      def 							lowercase__       (     self     :					int    )							-> Optional[int]:
            _lowerCAmelCase					=				"""cpu"""

            _lowerCAmelCase					=				self.get_dummy_components()

            _lowerCAmelCase					=				self.pipeline_class(**_snake_case    )
            _lowerCAmelCase					=				pipe.to(_snake_case    )

            pipe.set_progress_bar_config(disable=_snake_case    )

            _lowerCAmelCase					=				self.get_dummy_inputs(_snake_case     ,     pil_image=_snake_case    )

            _lowerCAmelCase					=				pipe(**_snake_case    )
            _lowerCAmelCase					=				output.images

            _lowerCAmelCase					=				self.get_dummy_inputs(_snake_case     ,     pil_image=_snake_case    )

            _lowerCAmelCase					=				pipe(
                **_snake_case     ,     return_dict=_snake_case     ,     )[0]

            _lowerCAmelCase					=				image[0, -3:, -3:, -1]
            _lowerCAmelCase					=				image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)

            _lowerCAmelCase					=				np.array(
                [
                    0.99_97,
                    0.00_02,
                    0.99_97,
                    0.99_97,
                    0.99_69,
                    0.00_23,
                    0.99_97,
                    0.99_69,
                    0.99_70,
                ]    )

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




      def 							lowercase__       (     self     :					List[Any]    )							-> str:
            _lowerCAmelCase					=				"""cpu"""

            _lowerCAmelCase					=				self.get_dummy_components()

            _lowerCAmelCase					=				self.pipeline_class(**_snake_case    )
            _lowerCAmelCase					=				pipe.to(_snake_case    )

            pipe.set_progress_bar_config(disable=_snake_case    )

            _lowerCAmelCase					=				self.get_dummy_inputs(_snake_case     ,     pil_image=_snake_case    )

            _lowerCAmelCase					=				pipe(**_snake_case    )
            _lowerCAmelCase					=				output.images

            _lowerCAmelCase					=				self.get_dummy_inputs(_snake_case     ,     pil_image=_snake_case    )

            _lowerCAmelCase					=				pipe(
                **_snake_case     ,     return_dict=_snake_case     ,     )[0]

            _lowerCAmelCase					=				image[0, -3:, -3:, -1]
            _lowerCAmelCase					=				image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)

            _lowerCAmelCase					=				np.array([0.99_97, 0.00_03, 0.99_97, 0.99_97, 0.99_70, 0.00_24, 0.99_97, 0.99_71, 0.99_71]    )

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




      def 							lowercase__       (     self     :					Optional[Any]    )							-> Union[str, Any]:
            _lowerCAmelCase					=				"""cpu"""

            _lowerCAmelCase					=				self.get_dummy_components()

            _lowerCAmelCase					=				self.pipeline_class(**_snake_case    )
            _lowerCAmelCase					=				pipe.to(_snake_case    )

            pipe.set_progress_bar_config(disable=_snake_case    )

            _lowerCAmelCase					=				self.get_dummy_inputs(_snake_case     ,     pil_image=_snake_case    )
            _lowerCAmelCase					=				[
                pipeline_inputs["""image"""],
                pipeline_inputs["""image"""],
            ]

            _lowerCAmelCase					=				pipe(**_snake_case    )
            _lowerCAmelCase					=				output.images

            _lowerCAmelCase					=				self.get_dummy_inputs(_snake_case     ,     pil_image=_snake_case    )
            _lowerCAmelCase					=				[
                tuple_pipeline_inputs["""image"""],
                tuple_pipeline_inputs["""image"""],
            ]

            _lowerCAmelCase					=				pipe(
                **_snake_case     ,     return_dict=_snake_case     ,     )[0]

            _lowerCAmelCase					=				image[0, -3:, -3:, -1]
            _lowerCAmelCase					=				image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (2, 64, 64, 3)

            _lowerCAmelCase					=				np.array(
                [
                    0.99_97,
                    0.99_89,
                    0.00_08,
                    0.00_21,
                    0.99_60,
                    0.00_18,
                    0.00_14,
                    0.00_02,
                    0.99_33,
                ]    )

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




      def 							lowercase__       (     self     :					str    )							-> Any:
            _lowerCAmelCase					=				torch.device("""cpu"""    )




            class 					UpperCAmelCase	:
                  _lowercase:     Any			=  1




            _lowerCAmelCase					=				self.get_dummy_components()

            _lowerCAmelCase					=				self.pipeline_class(**_snake_case    )
            _lowerCAmelCase					=				pipe.to(_snake_case    )

            pipe.set_progress_bar_config(disable=_snake_case    )

            _lowerCAmelCase					=				torch.Generator(device=_snake_case    ).manual_seed(0    )
            _lowerCAmelCase					=				pipe.decoder.dtype
            _lowerCAmelCase					=				1

            _lowerCAmelCase					=				(
                batch_size,
                pipe.decoder.config.in_channels,
                pipe.decoder.config.sample_size,
                pipe.decoder.config.sample_size,
            )
            _lowerCAmelCase					=				pipe.prepare_latents(
                _snake_case     ,     dtype=_snake_case     ,     device=_snake_case     ,     generator=_snake_case     ,     latents=_snake_case     ,     scheduler=DummyScheduler()    )

            _lowerCAmelCase					=				(
                batch_size,
                pipe.super_res_first.config.in_channels // 2,
                pipe.super_res_first.config.sample_size,
                pipe.super_res_first.config.sample_size,
            )
            _lowerCAmelCase					=				pipe.prepare_latents(
                _snake_case     ,     dtype=_snake_case     ,     device=_snake_case     ,     generator=_snake_case     ,     latents=_snake_case     ,     scheduler=DummyScheduler()    )

            _lowerCAmelCase					=				self.get_dummy_inputs(_snake_case     ,     pil_image=_snake_case    )

            _lowerCAmelCase					=				pipe(
                **_snake_case     ,     decoder_latents=_snake_case     ,     super_res_latents=_snake_case    ).images

            _lowerCAmelCase					=				self.get_dummy_inputs(_snake_case     ,     pil_image=_snake_case    )
            # Don't pass image, instead pass embedding
            _lowerCAmelCase					=				pipeline_inputs.pop("""image"""    )
            _lowerCAmelCase					=				pipe.image_encoder(_snake_case    ).image_embeds

            _lowerCAmelCase					=				pipe(
                **_snake_case     ,     decoder_latents=_snake_case     ,     super_res_latents=_snake_case     ,     image_embeddings=_snake_case     ,     ).images

            # make sure passing text embeddings manually is identical
            assert np.abs(img_out_a - img_out_a    ).max() < 1E-4




      @skip_mps
      def 							lowercase__       (     self     :					Tuple    )							-> str:
            _lowerCAmelCase					=				torch_device == """cpu"""

            # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
            _lowerCAmelCase					=				1E-2

            self._test_attention_slicing_forward_pass(
                test_max_difference=_snake_case     ,     expected_max_diff=_snake_case    )




      @skip_mps
      def 							lowercase__       (     self     :					str    )							-> List[Any]:
            _lowerCAmelCase					=				torch_device == """cpu"""
            _lowerCAmelCase					=				True
            _lowerCAmelCase					=				[
                """decoder_num_inference_steps""",
                """super_res_num_inference_steps""",
            ]

            self._test_inference_batch_single_identical(
                test_max_difference=_snake_case     ,     relax_max_difference=_snake_case     ,     additional_params_copy_to_batched_inputs=_snake_case     ,     )




      def 							lowercase__       (     self     :					List[Any]    )							-> Optional[Any]:
            _lowerCAmelCase					=				[
                """decoder_num_inference_steps""",
                """super_res_num_inference_steps""",
            ]

            if torch_device == "mps":
                  # TODO: MPS errors with larger batch sizes
                  _lowerCAmelCase					=				[2, 3]
                  self._test_inference_batch_consistent(
                      batch_sizes=_snake_case     ,     additional_params_copy_to_batched_inputs=_snake_case     ,     )
            else:
                  self._test_inference_batch_consistent(
                      additional_params_copy_to_batched_inputs=_snake_case    )




      @skip_mps
      def 							lowercase__       (     self     :					Optional[Any]    )							-> int:
            return super().test_dict_tuple_outputs_equivalent()




      @skip_mps
      def 							lowercase__       (     self     :					Tuple    )							-> Union[str, Any]:
            return super().test_save_load_local()




      @skip_mps
      def 							lowercase__       (     self     :					Any    )							-> List[Any]:
            return super().test_save_load_optional_components()




@slow
@require_torch_gpu
class 					UpperCAmelCase	(     unittest.TestCase  ):




      def 							lowercase__       (     self     :					Union[str, Any]    )							-> Union[str, Any]:
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def 							lowercase__       (     self     :					Union[str, Any]    )							-> Dict:
            _lowerCAmelCase					=				load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"""    )
            _lowerCAmelCase					=				load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                """/unclip/karlo_v1_alpha_cat_variation_fp16.npy"""    )

            _lowerCAmelCase					=				UnCLIPImageVariationPipeline.from_pretrained(
                """kakaobrain/karlo-v1-alpha-image-variations"""     ,     torch_dtype=torch.floataa    )
            _lowerCAmelCase					=				pipeline.to(_snake_case    )
            pipeline.set_progress_bar_config(disable=_snake_case    )

            _lowerCAmelCase					=				torch.Generator(device="""cpu"""    ).manual_seed(0    )
            _lowerCAmelCase					=				pipeline(
                _snake_case     ,     generator=_snake_case     ,     output_type="""np"""     ,     )

            _lowerCAmelCase					=				output.images[0]

            assert image.shape == (2_56, 2_56, 3)

            assert_mean_pixel_difference(_snake_case     ,     _snake_case     ,     15    )








"""simple docstring"""





import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
		import os

		import jax.numpy as jnp
		from jax import jit

		from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
		from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

		UpperCamelCase_       :    List[Any]						      =			'''0.12'''  # assumed parallelism: 8


if is_torch_available():
		import torch
def 				A_  (__a					, __a					, __a=None					):

					'''simple docstring'''





					if rng is None:
										A_								=    random.Random()

					A_								=    1
					for dim in shape:
										total_dims *= dim

					A_								=    []
					for _ in range(__a					):
										values.append(rng.randint(0					, vocab_size - 1					)					)

					A_								=    np.array(__a					, dtype=jnp.intaa					).reshape(__a					)

					return output
def 				A_  (__a					, __a=None					):

					'''simple docstring'''





					A_								=    ids_tensor(__a					, vocab_size=2					, rng=__a					)
					# make sure that at least one token is attended to for each batch
					A_								=    1
					return attn_mask



@require_flax
class __lowerCAmelCase					:




			"""simple docstring"""







			snake_case						   =					None
			snake_case						   =					()




			def   lowerCamelCase__  (					self    :       Any						)   ->			Tuple:







								"""simple docstring"""
								A_			,     A_								=    self.model_tester.prepare_config_and_inputs_for_common()

								# cut to half length & take max batch_size 3
								A_								=    2
								A_								=    inputs["input_ids"].shape[-1] // 2
								A_								=    inputs["input_ids"][:max_batch_size, :sequence_length]

								A_								=    jnp.ones_like(_snake_case						)
								A_								=    attention_mask[:max_batch_size, :sequence_length]

								# generate max 5 tokens
								A_								=    input_ids.shape[-1] + 5
								if config.eos_token_id is not None and config.pad_token_id is None:
													# hack to allow generate for models such as GPT2 as is done in `generate()`
													A_								=    config.eos_token_id
								return config, input_ids, attention_mask, max_length




			@is_pt_flax_cross_test
			def   lowerCamelCase__  (					self    :       List[Any]						)   ->			List[Any]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    0

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
													A_								=    getattr(_snake_case							, _snake_case						)
													A_								=    pt_model_class(_snake_case						).eval()
													A_								=    load_flax_weights_in_pytorch_model(_snake_case							, flax_model.params						)

													A_								=    flax_model.generate(_snake_case						).sequences
													A_								=    pt_model.generate(torch.tensor(_snake_case							, dtype=torch.long						)						)

													if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																		A_								=    flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

													self.assertListEqual(pt_generation_outputs.numpy().tolist()							, flax_generation_outputs.tolist()						)




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







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			List[str]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    True
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    2

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       str						)   ->			Union[str, Any]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    2
								A_								=    2

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[0]							, input_ids.shape[0] * config.num_return_sequences						)




			def   lowerCamelCase__  (					self    :       Union[str, Any]						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    True
								A_								=    max_length
								A_								=    0.8
								A_								=    10
								A_								=    0.3
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			Any:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    max_length
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Optional[Any]						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    max_length
								A_								=    2
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Any:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    False
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			int:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    True
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)






			def   lowerCamelCase__  (					self    :       Dict						)   ->			int:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    2
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)





@require_flax
class __lowerCAmelCase					(       unittest.TestCase   ):




			"""simple docstring"""







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







								"""simple docstring"""
								A_								=    AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"						)
								A_								=    FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"						)

								A_								=    "Hello world"
								A_								=    tokenizer(_snake_case							, return_tensors="np"						).input_ids

								# typos are quickly detected (the correct argument is `do_sample`)
								with self.assertRaisesRegex(_snake_case							, "do_samples"						):
													model.generate(_snake_case							, do_samples=_snake_case						)

								# arbitrary arguments that will not be used anywhere are also not accepted
								with self.assertRaisesRegex(_snake_case							, "foo"						):
													A_								=    {"foo": "bar"}
													model.generate(_snake_case							, **_snake_case						)




