
'''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       UpperCAmelCase					(       unittest.TestCase			):





		"""simple docstring"""





		def _lowercase	(     self	,       _snake_case	,       _snake_case  )					->     str:
				_UpperCamelCase							:			Any     =   jnp.ones((batch_size, length)  ) / length
				return scores
		def _lowercase	(     self  )					->     List[str]:
				_UpperCamelCase							:			Any     =   None
				_UpperCamelCase							:			str     =   20

				_UpperCamelCase							:			Optional[int]     =   self._get_uniform_logits(batch_size=2	,       length=lowercase__  )

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

				# compute softmax
				_UpperCamelCase							:			Optional[int]     =   jax.nn.softmax(lowercase__	,       axis=-1  )

				_UpperCamelCase							:			List[Any]     =   FlaxTemperatureLogitsWarper(temperature=0.5  )
				_UpperCamelCase							:			Tuple     =   FlaxTemperatureLogitsWarper(temperature=1.3  )

				_UpperCamelCase							:			Optional[int]     =   jax.nn.softmax(temp_dist_warper_sharper(lowercase__	,       scores.copy()	,       cur_len=lowercase__  )	,       axis=-1  )
				_UpperCamelCase							:			Tuple     =   jax.nn.softmax(temp_dist_warper_smoother(lowercase__	,       scores.copy()	,       cur_len=lowercase__  )	,       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 _lowercase	(     self  )					->     Optional[Any]:
				_UpperCamelCase							:			Optional[int]     =   None
				_UpperCamelCase							:			Optional[Any]     =   10
				_UpperCamelCase							:			str     =   2

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

				_UpperCamelCase							:			List[Any]     =   FlaxTopKLogitsWarper(3  )

				_UpperCamelCase							:			Union[str, Any]     =   top_k_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )

				# 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
				_UpperCamelCase							:			Optional[int]     =   5
				_UpperCamelCase							:			List[Any]     =   FlaxTopKLogitsWarper(top_k=1	,       filter_value=0.0	,       min_tokens_to_keep=3  )

				_UpperCamelCase							:			Tuple     =   np.broadcast_to(np.arange(lowercase__  )[None, :]	,       (batch_size, length)  ).copy()
				_UpperCamelCase							:			Optional[Any]     =   top_k_warp_safety_check(lowercase__	,       lowercase__	,       cur_len=lowercase__  )

				# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
				self.assertListEqual((scores == 0.0).sum(axis=-1  ).tolist()	,       [2, 2]  )
		def _lowercase	(     self  )					->     Tuple:
				_UpperCamelCase							:			List[str]     =   None
				_UpperCamelCase							:			Dict     =   10
				_UpperCamelCase							:			str     =   2

				# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
				_UpperCamelCase							:			str     =   np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]  )  )

				_UpperCamelCase							:			List[Any]     =   FlaxTopPLogitsWarper(0.8  )
				_UpperCamelCase							:			Union[str, Any]     =   np.exp(top_p_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )  )

				# dist should be filtered to keep min num values so that sum is >= top_p
				# exp (-inf) => 0
				_UpperCamelCase							:			Optional[int]     =   np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]  )
				self.assertTrue(np.allclose(lowercase__	,       lowercase__	,       atol=1E-3  )  )

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

				# make ramp_logits more extreme
				_UpperCamelCase							:			Tuple     =   ramp_logits[1] * 100.0

				# make sure at least 2 tokens are kept
				_UpperCamelCase							:			Dict     =   FlaxTopPLogitsWarper(0.9	,       min_tokens_to_keep=2	,       filter_value=0.0  )
				_UpperCamelCase							:			int     =   top_p_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )

				# 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 _lowercase	(     self  )					->     int:
				_UpperCamelCase							:			int     =   20
				_UpperCamelCase							:			str     =   4
				_UpperCamelCase							:			Dict     =   0

				_UpperCamelCase							:			List[str]     =   FlaxMinLengthLogitsProcessor(min_length=10	,       eos_token_id=lowercase__  )

				# check that min length is applied at length 5
				_UpperCamelCase							:			Tuple     =   ids_tensor((batch_size, 20)	,       vocab_size=20  )
				_UpperCamelCase							:			Dict     =   5
				_UpperCamelCase							:			Union[str, Any]     =   self._get_uniform_logits(lowercase__	,       lowercase__  )
				_UpperCamelCase							:			Any     =   min_dist_processor(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()	,       4 * [-float('''inf'''  )]  )

				# check that min length is not applied anymore at length 15
				_UpperCamelCase							:			Optional[int]     =   self._get_uniform_logits(lowercase__	,       lowercase__  )
				_UpperCamelCase							:			Tuple     =   15
				_UpperCamelCase							:			str     =   min_dist_processor(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				self.assertFalse(jnp.isinf(lowercase__  ).any()  )
		def _lowercase	(     self  )					->     Optional[int]:
				_UpperCamelCase							:			Optional[int]     =   20
				_UpperCamelCase							:			List[str]     =   4
				_UpperCamelCase							:			Optional[int]     =   0

				_UpperCamelCase							:			List[str]     =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase__  )

				# check that all scores are -inf except the bos_token_id score
				_UpperCamelCase							:			List[str]     =   ids_tensor((batch_size, 1)	,       vocab_size=20  )
				_UpperCamelCase							:			Optional[int]     =   1
				_UpperCamelCase							:			Union[str, Any]     =   self._get_uniform_logits(lowercase__	,       lowercase__  )
				_UpperCamelCase							:			Any     =   logits_processor(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				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
				_UpperCamelCase							:			Optional[Any]     =   3
				_UpperCamelCase							:			List[Any]     =   self._get_uniform_logits(lowercase__	,       lowercase__  )
				_UpperCamelCase							:			str     =   logits_processor(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				self.assertFalse(jnp.isinf(lowercase__  ).any()  )
		def _lowercase	(     self  )					->     str:
				_UpperCamelCase							:			Tuple     =   20
				_UpperCamelCase							:			List[str]     =   4
				_UpperCamelCase							:			Any     =   0
				_UpperCamelCase							:			Tuple     =   5

				_UpperCamelCase							:			Optional[Any]     =   FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase__	,       eos_token_id=lowercase__  )

				# check that all scores are -inf except the eos_token_id when max_length is reached
				_UpperCamelCase							:			List[Any]     =   ids_tensor((batch_size, 4)	,       vocab_size=20  )
				_UpperCamelCase							:			List[Any]     =   4
				_UpperCamelCase							:			Any     =   self._get_uniform_logits(lowercase__	,       lowercase__  )
				_UpperCamelCase							:			List[Any]     =   logits_processor(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				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
				_UpperCamelCase							:			Dict     =   3
				_UpperCamelCase							:			Optional[Any]     =   self._get_uniform_logits(lowercase__	,       lowercase__  )
				_UpperCamelCase							:			Tuple     =   logits_processor(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				self.assertFalse(jnp.isinf(lowercase__  ).any()  )
		def _lowercase	(     self  )					->     Any:
				_UpperCamelCase							:			List[str]     =   4
				_UpperCamelCase							:			Optional[Any]     =   10
				_UpperCamelCase							:			List[str]     =   15
				_UpperCamelCase							:			int     =   2
				_UpperCamelCase							:			str     =   1
				_UpperCamelCase							:			Dict     =   15

				# dummy input_ids and scores
				_UpperCamelCase							:			Union[str, Any]     =   ids_tensor((batch_size, sequence_length)	,       lowercase__  )
				_UpperCamelCase							:			List[str]     =   input_ids.copy()

				_UpperCamelCase							:			Dict     =   self._get_uniform_logits(lowercase__	,       lowercase__  )
				_UpperCamelCase							:			List[Any]     =   scores.copy()

				# instantiate all dist processors
				_UpperCamelCase							:			str     =   FlaxTemperatureLogitsWarper(temperature=0.5  )
				_UpperCamelCase							:			Optional[Any]     =   FlaxTopKLogitsWarper(3  )
				_UpperCamelCase							:			Optional[Any]     =   FlaxTopPLogitsWarper(0.8  )

				# instantiate all logits processors
				_UpperCamelCase							:			Union[str, Any]     =   FlaxMinLengthLogitsProcessor(min_length=10	,       eos_token_id=lowercase__  )
				_UpperCamelCase							:			str     =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase__  )
				_UpperCamelCase							:			Dict     =   FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase__	,       eos_token_id=lowercase__  )

				_UpperCamelCase							:			List[Any]     =   10

				# no processor list
				_UpperCamelCase							:			List[str]     =   temp_dist_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				_UpperCamelCase							:			List[str]     =   top_k_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				_UpperCamelCase							:			Tuple     =   top_p_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				_UpperCamelCase							:			Tuple     =   min_dist_proc(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				_UpperCamelCase							:			Tuple     =   bos_dist_proc(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
				_UpperCamelCase							:			Tuple     =   eos_dist_proc(lowercase__	,       lowercase__	,       cur_len=lowercase__  )

				# with processor list
				_UpperCamelCase							:			List[Any]     =   FlaxLogitsProcessorList(
				    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
				_UpperCamelCase							:			str     =   processor(lowercase__	,       lowercase__	,       cur_len=lowercase__  )

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

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




		def _lowercase	(     self  )					->     List[Any]:
				_UpperCamelCase							:			Optional[int]     =   4
				_UpperCamelCase							:			Dict     =   10
				_UpperCamelCase							:			int     =   15
				_UpperCamelCase							:			str     =   2
				_UpperCamelCase							:			List[str]     =   1
				_UpperCamelCase							:			Optional[Any]     =   15

				# dummy input_ids and scores
				_UpperCamelCase							:			Optional[int]     =   ids_tensor((batch_size, sequence_length)	,       lowercase__  )
				_UpperCamelCase							:			Tuple     =   input_ids.copy()

				_UpperCamelCase							:			List[Any]     =   self._get_uniform_logits(lowercase__	,       lowercase__  )
				_UpperCamelCase							:			Union[str, Any]     =   scores.copy()

				# instantiate all dist processors
				_UpperCamelCase							:			Dict     =   FlaxTemperatureLogitsWarper(temperature=0.5  )
				_UpperCamelCase							:			Any     =   FlaxTopKLogitsWarper(3  )
				_UpperCamelCase							:			int     =   FlaxTopPLogitsWarper(0.8  )

				# instantiate all logits processors
				_UpperCamelCase							:			Any     =   FlaxMinLengthLogitsProcessor(min_length=10	,       eos_token_id=lowercase__  )
				_UpperCamelCase							:			List[str]     =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase__  )
				_UpperCamelCase							:			Tuple     =   FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase__	,       eos_token_id=lowercase__  )

				_UpperCamelCase							:			str     =   10

				# no processor list
				def run_no_processor_list(_snake_case	,       _snake_case	,       _snake_case  ):
						_UpperCamelCase							:			List[Any]     =   temp_dist_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
						_UpperCamelCase							:			List[Any]     =   top_k_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
						_UpperCamelCase							:			str     =   top_p_warp(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
						_UpperCamelCase							:			Tuple     =   min_dist_proc(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
						_UpperCamelCase							:			Tuple     =   bos_dist_proc(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
						_UpperCamelCase							:			Dict     =   eos_dist_proc(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
						return scores

				# with processor list
				def run_processor_list(_snake_case	,       _snake_case	,       _snake_case  ):
						_UpperCamelCase							:			Optional[int]     =   FlaxLogitsProcessorList(
						    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
						_UpperCamelCase							:			Optional[Any]     =   processor(lowercase__	,       lowercase__	,       cur_len=lowercase__  )
						return scores

				_UpperCamelCase							:			str     =   jax.jit(lowercase__  )
				_UpperCamelCase							:			Union[str, Any]     =   jax.jit(lowercase__  )

				_UpperCamelCase							:			Any     =   jitted_run_no_processor_list(lowercase__	,       lowercase__	,       lowercase__  )
				_UpperCamelCase							:			List[Any]     =   jitted_run_processor_list(lowercase__	,       lowercase__	,       lowercase__  )

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

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








from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
   if not (is_transformers_available() and is_torch_available() and is_transformers_version('''>=''', '''4.25.0''')):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import (
       VersatileDiffusionDualGuidedPipeline,
       VersatileDiffusionImageVariationPipeline,
       VersatileDiffusionPipeline,
       VersatileDiffusionTextToImagePipeline,
   )
else:
   from .modeling_text_unet import UNetFlatConditionModel
   from .pipeline_versatile_diffusion import VersatileDiffusionPipeline
   from .pipeline_versatile_diffusion_dual_guided import VersatileDiffusionDualGuidedPipeline
   from .pipeline_versatile_diffusion_image_variation import VersatileDiffusionImageVariationPipeline
   from .pipeline_versatile_diffusion_text_to_image import VersatileDiffusionTextToImagePipeline


