
"""simple docstring"""


import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class  __lowerCamelCase  ( lowerCAmelCase						,	lowerCAmelCase						,	unittest.TestCase			):
				a__:      Optional[int]          =    StableDiffusionDiffEditPipeline
				a__:      Any          =    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'height', 'width', 'image'} | {'image_latents'}
				a__:      Optional[int]          =    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {'image'} | {'image_latents'}
				a__:      str          =    frozenset(
				    []			)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
				a__:      Dict          =    frozenset([]			)




				def        UpperCAmelCase__						(				self   ):
							torch.manual_seed(0   )
							lowerCamelCase_					=  UNetaDConditionModel(
							    block_out_channels=(32, 64)							,						layers_per_block=2							,						sample_size=32							,						in_channels=4							,						out_channels=4							,						down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')							,						up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')							,						cross_attention_dim=32							,						attention_head_dim=(2, 4)							,						use_linear_projection=UpperCAmelCase							,						)
							lowerCamelCase_					=  DDIMScheduler(
							    beta_start=0.0_0_0_8_5							,						beta_end=0.0_1_2							,						beta_schedule='''scaled_linear'''							,						clip_sample=UpperCAmelCase							,						set_alpha_to_one=UpperCAmelCase							,						)
							lowerCamelCase_					=  DDIMInverseScheduler(
							    beta_start=0.0_0_0_8_5							,						beta_end=0.0_1_2							,						beta_schedule='''scaled_linear'''							,						clip_sample=UpperCAmelCase							,						set_alpha_to_zero=UpperCAmelCase							,						)
							torch.manual_seed(0   )
							lowerCamelCase_					=  AutoencoderKL(
							    block_out_channels=[32, 64]							,						in_channels=3							,						out_channels=3							,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']							,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']							,						latent_channels=4							,						sample_size=128							,						)
							torch.manual_seed(0   )
							lowerCamelCase_					=  CLIPTextConfig(
							    bos_token_id=0							,						eos_token_id=2							,						hidden_size=32							,						intermediate_size=37							,						layer_norm_eps=1e-0_5							,						num_attention_heads=4							,						num_hidden_layers=5							,						pad_token_id=1							,						vocab_size=1000							,						hidden_act='''gelu'''							,						projection_dim=512							,						)
							lowerCamelCase_					=  CLIPTextModel(UpperCAmelCase   )
							lowerCamelCase_					=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

							lowerCamelCase_					=  {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''inverse_scheduler''': inverse_scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}

							return components




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase=0   ):
							lowerCamelCase_					=  floats_tensor((1, 16, 16)							,						rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
							lowerCamelCase_					=  floats_tensor((1, 2, 4, 16, 16)							,						rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
							if str(UpperCAmelCase   ).startswith('''mps'''   ):
										lowerCamelCase_					=  torch.manual_seed(UpperCAmelCase   )
							else:
										lowerCamelCase_					=  torch.Generator(device=UpperCAmelCase   ).manual_seed(UpperCAmelCase   )
							lowerCamelCase_					=  {
							    '''prompt''': '''a dog and a newt''',
							    '''mask_image''': mask,
							    '''image_latents''': latents,
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''inpaint_strength''': 1.0,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}

							return inputs




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase=0   ):
							lowerCamelCase_					=  floats_tensor((1, 3, 32, 32)							,						rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
							lowerCamelCase_					=  image.cpu().permute(0							,						2							,						3							,						1   )[0]
							lowerCamelCase_					=  Image.fromarray(np.uinta(UpperCAmelCase   )   ).convert('''RGB'''   )
							if str(UpperCAmelCase   ).startswith('''mps'''   ):
										lowerCamelCase_					=  torch.manual_seed(UpperCAmelCase   )
							else:
										lowerCamelCase_					=  torch.Generator(device=UpperCAmelCase   ).manual_seed(UpperCAmelCase   )
							lowerCamelCase_					=  {
							    '''image''': image,
							    '''source_prompt''': '''a cat and a frog''',
							    '''target_prompt''': '''a dog and a newt''',
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''num_maps_per_mask''': 2,
							    '''mask_encode_strength''': 1.0,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}

							return inputs




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase=0   ):
							lowerCamelCase_					=  floats_tensor((1, 3, 32, 32)							,						rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
							lowerCamelCase_					=  image.cpu().permute(0							,						2							,						3							,						1   )[0]
							lowerCamelCase_					=  Image.fromarray(np.uinta(UpperCAmelCase   )   ).convert('''RGB'''   )
							if str(UpperCAmelCase   ).startswith('''mps'''   ):
										lowerCamelCase_					=  torch.manual_seed(UpperCAmelCase   )
							else:
										lowerCamelCase_					=  torch.Generator(device=UpperCAmelCase   ).manual_seed(UpperCAmelCase   )
							lowerCamelCase_					=  {
							    '''image''': image,
							    '''prompt''': '''a cat and a frog''',
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''inpaint_strength''': 1.0,
							    '''guidance_scale''': 6.0,
							    '''decode_latents''': True,
							    '''output_type''': '''numpy''',
							}
							return inputs




				def        UpperCAmelCase__						(				self   ):
							if not hasattr(self.pipeline_class							,						'''_optional_components'''   ):
										return

							lowerCamelCase_					=  self.get_dummy_components()
							lowerCamelCase_					=  self.pipeline_class(**UpperCAmelCase   )
							pipe.to(UpperCAmelCase   )
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							# set all optional components to None and update pipeline config accordingly
							for optional_component in pipe._optional_components:
										setattr(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   )
							pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}   )

							lowerCamelCase_					=  self.get_dummy_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe(**UpperCAmelCase   )[0]

							with tempfile.TemporaryDirectory() as tmpdir:
										pipe.save_pretrained(UpperCAmelCase   )
										lowerCamelCase_					=  self.pipeline_class.from_pretrained(UpperCAmelCase   )
										pipe_loaded.to(UpperCAmelCase   )
										pipe_loaded.set_progress_bar_config(disable=UpperCAmelCase   )

							for optional_component in pipe._optional_components:
										self.assertTrue(
										    getattr(UpperCAmelCase							,						UpperCAmelCase   ) is None							,						f"`{optional_component}` did not stay set to None after loading."							,						)

							lowerCamelCase_					=  self.get_dummy_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe_loaded(**UpperCAmelCase   )[0]

							lowerCamelCase_					=  np.abs(output - output_loaded   ).max()
							self.assertLess(UpperCAmelCase							,						1e-4   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  '''cpu'''

							lowerCamelCase_					=  self.get_dummy_components()
							lowerCamelCase_					=  self.pipeline_class(**UpperCAmelCase   )
							pipe.to(UpperCAmelCase   )
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  self.get_dummy_mask_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe.generate_mask(**UpperCAmelCase   )
							lowerCamelCase_					=  mask[0, -3:, -3:]

							self.assertEqual(mask.shape							,						(1, 16, 16)   )
							lowerCamelCase_					=  np.array([0] * 9   )
							lowerCamelCase_					=  np.abs(mask_slice.flatten() - expected_slice   ).max()
							self.assertLessEqual(UpperCAmelCase							,						1e-3   )
							self.assertEqual(mask[0, -3, -4]							,						0   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  '''cpu'''

							lowerCamelCase_					=  self.get_dummy_components()
							lowerCamelCase_					=  self.pipeline_class(**UpperCAmelCase   )
							pipe.to(UpperCAmelCase   )
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  self.get_dummy_inversion_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe.invert(**UpperCAmelCase   ).images
							lowerCamelCase_					=  image[0, -1, -3:, -3:]

							self.assertEqual(image.shape							,						(2, 32, 32, 3)   )
							lowerCamelCase_					=  np.array(
							    [0.5_1_5_0, 0.5_1_3_4, 0.5_0_4_3, 0.5_3_7_6, 0.4_6_9_4, 0.5_1_0_5_0, 0.5_0_1_5, 0.4_4_0_7, 0.4_7_9_9]							,						)
							lowerCamelCase_					=  np.abs(image_slice.flatten() - expected_slice   ).max()
							self.assertLessEqual(UpperCAmelCase							,						1e-3   )




				def        UpperCAmelCase__						(				self   ):
							super().test_inference_batch_single_identical(expected_max_diff=5e-3   )





				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  '''cpu'''

							lowerCamelCase_					=  self.get_dummy_components()

							lowerCamelCase_					=  {'''beta_start''': 0.0_0_0_8_5, '''beta_end''': 0.0_1_2, '''beta_schedule''': '''scaled_linear'''}
							lowerCamelCase_					=  DPMSolverMultistepScheduler(**UpperCAmelCase   )
							lowerCamelCase_					=  DPMSolverMultistepInverseScheduler(**UpperCAmelCase   )

							lowerCamelCase_					=  self.pipeline_class(**UpperCAmelCase   )
							pipe.to(UpperCAmelCase   )
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  self.get_dummy_inversion_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe.invert(**UpperCAmelCase   ).images
							lowerCamelCase_					=  image[0, -1, -3:, -3:]

							self.assertEqual(image.shape							,						(2, 32, 32, 3)   )
							lowerCamelCase_					=  np.array(
							    [0.5_1_5_0, 0.5_1_3_4, 0.5_0_4_3, 0.5_3_7_6, 0.4_6_9_4, 0.5_1_0_5_0, 0.5_0_1_5, 0.4_4_0_7, 0.4_7_9_9]							,						)
							lowerCamelCase_					=  np.abs(image_slice.flatten() - expected_slice   ).max()
							self.assertLessEqual(UpperCAmelCase							,						1e-3   )




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




				def        UpperCAmelCase__						(				self   ):
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




				@classmethod
				def        UpperCAmelCase__						(				cls   ):
							lowerCamelCase_					=  load_image(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'''   )

							lowerCamelCase_					=  raw_image.convert('''RGB'''   ).resize((768, 768)   )

							lowerCamelCase_					=  raw_image




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  torch.manual_seed(0   )

							lowerCamelCase_					=  StableDiffusionDiffEditPipeline.from_pretrained(
							    '''stabilityai/stable-diffusion-2-1'''							,						safety_checker=UpperCAmelCase							,						torch_dtype=torch.floataa   )
							lowerCamelCase_					=  DDIMScheduler.from_config(pipe.scheduler.config   )
							lowerCamelCase_					=  DDIMInverseScheduler.from_config(pipe.scheduler.config   )
							pipe.enable_model_cpu_offload()
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  '''a bowl of fruit'''
							lowerCamelCase_					=  '''a bowl of pears'''

							lowerCamelCase_					=  pipe.generate_mask(
							    image=self.raw_image							,						source_prompt=UpperCAmelCase							,						target_prompt=UpperCAmelCase							,						generator=UpperCAmelCase							,						)

							lowerCamelCase_					=  pipe.invert(
							    prompt=UpperCAmelCase							,						image=self.raw_image							,						inpaint_strength=0.7							,						generator=UpperCAmelCase   ).latents

							lowerCamelCase_					=  pipe(
							    prompt=UpperCAmelCase							,						mask_image=UpperCAmelCase							,						image_latents=UpperCAmelCase							,						generator=UpperCAmelCase							,						negative_prompt=UpperCAmelCase							,						inpaint_strength=0.7							,						output_type='''numpy'''							,						).images[0]

							lowerCamelCase_					=  (
							    np.array(
							        load_image(
							            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							            '''/diffedit/pears.png'''   ).resize((768, 768)   )   )
							    / 255
							)
							assert np.abs((expected_image - image).max()   ) < 5e-1





				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  torch.manual_seed(0   )

							lowerCamelCase_					=  StableDiffusionDiffEditPipeline.from_pretrained(
							    '''stabilityai/stable-diffusion-2-1'''							,						safety_checker=UpperCAmelCase							,						torch_dtype=torch.floataa   )
							lowerCamelCase_					=  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config   )
							lowerCamelCase_					=  DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config   )
							pipe.enable_model_cpu_offload()
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  '''a bowl of fruit'''
							lowerCamelCase_					=  '''a bowl of pears'''

							lowerCamelCase_					=  pipe.generate_mask(
							    image=self.raw_image							,						source_prompt=UpperCAmelCase							,						target_prompt=UpperCAmelCase							,						generator=UpperCAmelCase							,						)

							lowerCamelCase_					=  pipe.invert(
							    prompt=UpperCAmelCase							,						image=self.raw_image							,						inpaint_strength=0.7							,						generator=UpperCAmelCase							,						num_inference_steps=25							,						).latents

							lowerCamelCase_					=  pipe(
							    prompt=UpperCAmelCase							,						mask_image=UpperCAmelCase							,						image_latents=UpperCAmelCase							,						generator=UpperCAmelCase							,						negative_prompt=UpperCAmelCase							,						inpaint_strength=0.7							,						num_inference_steps=25							,						output_type='''numpy'''							,						).images[0]

							lowerCamelCase_					=  (
							    np.array(
							        load_image(
							            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							            '''/diffedit/pears.png'''   ).resize((768, 768)   )   )
							    / 255
							)
							assert np.abs((expected_image - image).max()   ) < 5e-1





"""simple docstring"""


import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
						import jax
						import jax.numpy as jnp

						from transformers.generation import (
						    FlaxForcedBOSTokenLogitsProcessor,
						    FlaxForcedEOSTokenLogitsProcessor,
						    FlaxLogitsProcessorList,
						    FlaxMinLengthLogitsProcessor,
						    FlaxTemperatureLogitsWarper,
						    FlaxTopKLogitsWarper,
						    FlaxTopPLogitsWarper,
						)




@require_flax
class  __lowerCamelCase  ( unittest.TestCase			):




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase   ):
							lowerCamelCase_					=  jnp.ones((batch_size, length)   ) / length
							return scores




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  None
							lowerCamelCase_					=  20

							lowerCamelCase_					=  self._get_uniform_logits(batch_size=2							,						length=UpperCAmelCase   )

							# tweak scores to not be uniform anymore
							lowerCamelCase_					=  scores.at[1, 5].set((1 / length) + 0.1   )  # peak, 1st batch
							lowerCamelCase_					=  scores.at[1, 10].set((1 / length) - 0.4   )  # valley, 1st batch

							# compute softmax
							lowerCamelCase_					=  jax.nn.softmax(UpperCAmelCase							,						axis=-1   )

							lowerCamelCase_					=  FlaxTemperatureLogitsWarper(temperature=0.5   )
							lowerCamelCase_					=  FlaxTemperatureLogitsWarper(temperature=1.3   )

							lowerCamelCase_					=  jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase							,						scores.copy()							,						cur_len=UpperCAmelCase   )							,						axis=-1   )
							lowerCamelCase_					=  jax.nn.softmax(temp_dist_warper_smoother(UpperCAmelCase							,						scores.copy()							,						cur_len=UpperCAmelCase   )							,						axis=-1   )

							# uniform distribution stays uniform
							self.assertTrue(jnp.allclose(probs[0, :]							,						warped_prob_sharp[0, :]							,						atol=1e-3   )   )
							self.assertTrue(jnp.allclose(probs[0, :]							,						warped_prob_smooth[0, :]							,						atol=1e-3   )   )

							# sharp peaks get higher, valleys get lower
							self.assertLess(probs[1, :].max()							,						warped_prob_sharp[1, :].max()   )
							self.assertGreater(probs[1, :].min()							,						warped_prob_sharp[1, :].min()   )

							# smooth peaks get lower, valleys get higher
							self.assertGreater(probs[1, :].max()							,						warped_prob_smooth[1, :].max()   )
							self.assertLess(probs[1, :].min()							,						warped_prob_smooth[1, :].min()   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  None
							lowerCamelCase_					=  10
							lowerCamelCase_					=  2

							# create ramp distribution
							lowerCamelCase_					=  np.broadcast_to(np.arange(UpperCAmelCase   )[None, :]							,						(batch_size, vocab_size)   ).copy()
							lowerCamelCase_					=  ramp_logits[1:, : vocab_size // 2] + vocab_size

							lowerCamelCase_					=  FlaxTopKLogitsWarper(3   )

							lowerCamelCase_					=  top_k_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )

							# check that correct tokens are filtered
							self.assertListEqual(jnp.isinf(scores[0]   ).tolist()							,						7 * [True] + 3 * [False]   )
							self.assertListEqual(jnp.isinf(scores[1]   ).tolist()							,						2 * [True] + 3 * [False] + 5 * [True]   )

							# check special case
							lowerCamelCase_					=  5
							lowerCamelCase_					=  FlaxTopKLogitsWarper(top_k=1							,						filter_value=0.0							,						min_tokens_to_keep=3   )

							lowerCamelCase_					=  np.broadcast_to(np.arange(UpperCAmelCase   )[None, :]							,						(batch_size, length)   ).copy()
							lowerCamelCase_					=  top_k_warp_safety_check(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )

							# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
							self.assertListEqual((scores == 0.0).sum(axis=-1   ).tolist()							,						[2, 2]   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  None
							lowerCamelCase_					=  10
							lowerCamelCase_					=  2

							# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
							lowerCamelCase_					=  np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.1_5, 0.3, 0.3, 0.2_5]]   )   )

							lowerCamelCase_					=  FlaxTopPLogitsWarper(0.8   )
							lowerCamelCase_					=  np.exp(top_p_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )   )

							# dist should be filtered to keep min num values so that sum is >= top_p
							# exp (-inf) => 0
							lowerCamelCase_					=  np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]   )
							self.assertTrue(np.allclose(UpperCAmelCase							,						UpperCAmelCase							,						atol=1e-3   )   )

							# check edge cases with negative and extreme logits
							lowerCamelCase_					=  np.broadcast_to(np.arange(UpperCAmelCase   )[None, :]							,						(batch_size, vocab_size)   ).copy() - (
							    vocab_size // 2
							)

							# make ramp_logits more extreme
							lowerCamelCase_					=  ramp_logits[1] * 1_0_0.0

							# make sure at least 2 tokens are kept
							lowerCamelCase_					=  FlaxTopPLogitsWarper(0.9							,						min_tokens_to_keep=2							,						filter_value=0.0   )
							lowerCamelCase_					=  top_p_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )

							# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
							self.assertListEqual((filtered_dist != 0.0).sum(axis=-1   ).tolist()							,						[3, 2]   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  20
							lowerCamelCase_					=  4
							lowerCamelCase_					=  0

							lowerCamelCase_					=  FlaxMinLengthLogitsProcessor(min_length=10							,						eos_token_id=UpperCAmelCase   )

							# check that min length is applied at length 5
							lowerCamelCase_					=  ids_tensor((batch_size, 20)							,						vocab_size=20   )
							lowerCamelCase_					=  5
							lowerCamelCase_					=  self._get_uniform_logits(UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  min_dist_processor(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()							,						4 * [-float('''inf'''   )]   )

							# check that min length is not applied anymore at length 15
							lowerCamelCase_					=  self._get_uniform_logits(UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  15
							lowerCamelCase_					=  min_dist_processor(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							self.assertFalse(jnp.isinf(UpperCAmelCase   ).any()   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  20
							lowerCamelCase_					=  4
							lowerCamelCase_					=  0

							lowerCamelCase_					=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase   )

							# check that all scores are -inf except the bos_token_id score
							lowerCamelCase_					=  ids_tensor((batch_size, 1)							,						vocab_size=20   )
							lowerCamelCase_					=  1
							lowerCamelCase_					=  self._get_uniform_logits(UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  logits_processor(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]   ).all()   )
							self.assertListEqual(scores[:, bos_token_id].tolist()							,						4 * [0]   )  # score for bos_token_id shold be zero

							# check that bos_token_id is not forced if current length is greater than 1
							lowerCamelCase_					=  3
							lowerCamelCase_					=  self._get_uniform_logits(UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  logits_processor(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							self.assertFalse(jnp.isinf(UpperCAmelCase   ).any()   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  20
							lowerCamelCase_					=  4
							lowerCamelCase_					=  0
							lowerCamelCase_					=  5

							lowerCamelCase_					=  FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase							,						eos_token_id=UpperCAmelCase   )

							# check that all scores are -inf except the eos_token_id when max_length is reached
							lowerCamelCase_					=  ids_tensor((batch_size, 4)							,						vocab_size=20   )
							lowerCamelCase_					=  4
							lowerCamelCase_					=  self._get_uniform_logits(UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  logits_processor(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]   ).all()   )
							self.assertListEqual(scores[:, eos_token_id].tolist()							,						4 * [0]   )  # score for eos_token_id should be zero

							# check that eos_token_id is not forced if max_length is not reached
							lowerCamelCase_					=  3
							lowerCamelCase_					=  self._get_uniform_logits(UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  logits_processor(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							self.assertFalse(jnp.isinf(UpperCAmelCase   ).any()   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  4
							lowerCamelCase_					=  10
							lowerCamelCase_					=  15
							lowerCamelCase_					=  2
							lowerCamelCase_					=  1
							lowerCamelCase_					=  15

							# dummy input_ids and scores
							lowerCamelCase_					=  ids_tensor((batch_size, sequence_length)							,						UpperCAmelCase   )
							lowerCamelCase_					=  input_ids.copy()

							lowerCamelCase_					=  self._get_uniform_logits(UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  scores.copy()

							# instantiate all dist processors
							lowerCamelCase_					=  FlaxTemperatureLogitsWarper(temperature=0.5   )
							lowerCamelCase_					=  FlaxTopKLogitsWarper(3   )
							lowerCamelCase_					=  FlaxTopPLogitsWarper(0.8   )

							# instantiate all logits processors
							lowerCamelCase_					=  FlaxMinLengthLogitsProcessor(min_length=10							,						eos_token_id=UpperCAmelCase   )
							lowerCamelCase_					=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase   )
							lowerCamelCase_					=  FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase							,						eos_token_id=UpperCAmelCase   )

							lowerCamelCase_					=  10

							# no processor list
							lowerCamelCase_					=  temp_dist_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							lowerCamelCase_					=  top_k_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							lowerCamelCase_					=  top_p_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							lowerCamelCase_					=  min_dist_proc(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							lowerCamelCase_					=  bos_dist_proc(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
							lowerCamelCase_					=  eos_dist_proc(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )

							# with processor list
							lowerCamelCase_					=  FlaxLogitsProcessorList(
							    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]   )
							lowerCamelCase_					=  processor(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )

							# scores should be equal
							self.assertTrue(jnp.allclose(UpperCAmelCase							,						UpperCAmelCase							,						atol=1e-3   )   )

							# input_ids should never be changed
							self.assertListEqual(input_ids.tolist()							,						input_ids_comp.tolist()   )





				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  4
							lowerCamelCase_					=  10
							lowerCamelCase_					=  15
							lowerCamelCase_					=  2
							lowerCamelCase_					=  1
							lowerCamelCase_					=  15

							# dummy input_ids and scores
							lowerCamelCase_					=  ids_tensor((batch_size, sequence_length)							,						UpperCAmelCase   )
							lowerCamelCase_					=  input_ids.copy()

							lowerCamelCase_					=  self._get_uniform_logits(UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  scores.copy()

							# instantiate all dist processors
							lowerCamelCase_					=  FlaxTemperatureLogitsWarper(temperature=0.5   )
							lowerCamelCase_					=  FlaxTopKLogitsWarper(3   )
							lowerCamelCase_					=  FlaxTopPLogitsWarper(0.8   )

							# instantiate all logits processors
							lowerCamelCase_					=  FlaxMinLengthLogitsProcessor(min_length=10							,						eos_token_id=UpperCAmelCase   )
							lowerCamelCase_					=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase   )
							lowerCamelCase_					=  FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase							,						eos_token_id=UpperCAmelCase   )

							lowerCamelCase_					=  10

							# no processor list
							def run_no_processor_list(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   ):
										lowerCamelCase_					=  temp_dist_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
										lowerCamelCase_					=  top_k_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
										lowerCamelCase_					=  top_p_warp(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
										lowerCamelCase_					=  min_dist_proc(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
										lowerCamelCase_					=  bos_dist_proc(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
										lowerCamelCase_					=  eos_dist_proc(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
										return scores

							# with processor list
							def run_processor_list(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   ):
										lowerCamelCase_					=  FlaxLogitsProcessorList(
										    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]   )
										lowerCamelCase_					=  processor(UpperCAmelCase							,						UpperCAmelCase							,						cur_len=UpperCAmelCase   )
										return scores

							lowerCamelCase_					=  jax.jit(UpperCAmelCase   )
							lowerCamelCase_					=  jax.jit(UpperCAmelCase   )

							lowerCamelCase_					=  jitted_run_no_processor_list(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   )
							lowerCamelCase_					=  jitted_run_processor_list(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   )

							# scores should be equal
							self.assertTrue(jnp.allclose(UpperCAmelCase							,						UpperCAmelCase							,						atol=1e-3   )   )

							# input_ids should never be changed
							self.assertListEqual(input_ids.tolist()							,						input_ids_comp.tolist()   )



