





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






       """simple docstring"""

       def      A_  (      self				:     Tuple							,	a				:     int							,	a				:     int    )	->Dict:
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jnp.ones((batch_size, length)    ) / length
             return scores

       def      A_  (      self				:     Tuple    )	->List[str]:
             SCREAMING_SNAKE_CASE__     :				List[Any]   =					None
             SCREAMING_SNAKE_CASE__     :				Tuple   =					20

             SCREAMING_SNAKE_CASE__     :				Optional[int]   =					self._get_uniform_logits(batch_size=2							,	length=lowerCAmelCase_    )

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

             # compute softmax
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					jax.nn.softmax(lowerCAmelCase_							,	axis=-1    )

             SCREAMING_SNAKE_CASE__     :				List[Any]   =					FlaxTemperatureLogitsWarper(temperature=0.5    )
             SCREAMING_SNAKE_CASE__     :				Optional[int]   =					FlaxTemperatureLogitsWarper(temperature=1.3    )

             SCREAMING_SNAKE_CASE__     :				str   =					jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase_							,	scores.copy()							,	cur_len=lowerCAmelCase_    )							,	axis=-1    )
             SCREAMING_SNAKE_CASE__     :				Any   =					jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase_							,	scores.copy()							,	cur_len=lowerCAmelCase_    )							,	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      A_  (      self				:     List[Any]    )	->Optional[Any]:
             SCREAMING_SNAKE_CASE__     :				Dict   =					None
             SCREAMING_SNAKE_CASE__     :				Optional[int]   =					10
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					2

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

             SCREAMING_SNAKE_CASE__     :				str   =					FlaxTopKLogitsWarper(3    )

             SCREAMING_SNAKE_CASE__     :				str   =					top_k_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )

             # 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
             SCREAMING_SNAKE_CASE__     :				Tuple   =					5
             SCREAMING_SNAKE_CASE__     :				int   =					FlaxTopKLogitsWarper(top_k=1							,	filter_value=0.0							,	min_tokens_to_keep=3    )

             SCREAMING_SNAKE_CASE__     :				List[Any]   =					np.broadcast_to(np.arange(lowerCAmelCase_    )[None, :]							,	(batch_size, length)    ).copy()
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					top_k_warp_safety_check(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )

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

       def      A_  (      self				:     int    )	->Dict:
             SCREAMING_SNAKE_CASE__     :				List[Any]   =					None
             SCREAMING_SNAKE_CASE__     :				Tuple   =					10
             SCREAMING_SNAKE_CASE__     :				str   =					2

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

             SCREAMING_SNAKE_CASE__     :				int   =					FlaxTopPLogitsWarper(0.8    )
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					np.exp(top_p_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )    )

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

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

             # make ramp_logits more extreme
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					ramp_logits[1] * 100.0

             # make sure at least 2 tokens are kept
             SCREAMING_SNAKE_CASE__     :				Tuple   =					FlaxTopPLogitsWarper(0.9							,	min_tokens_to_keep=2							,	filter_value=0.0    )
             SCREAMING_SNAKE_CASE__     :				Dict   =					top_p_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )

             # 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      A_  (      self				:     Dict    )	->Union[str, Any]:
             SCREAMING_SNAKE_CASE__     :				List[str]   =					20
             SCREAMING_SNAKE_CASE__     :				str   =					4
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					0

             SCREAMING_SNAKE_CASE__     :				Any   =					FlaxMinLengthLogitsProcessor(min_length=10							,	eos_token_id=lowerCAmelCase_    )

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

             # check that min length is not applied anymore at length 15
             SCREAMING_SNAKE_CASE__     :				List[str]   =					self._get_uniform_logits(lowerCAmelCase_							,	lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				List[Any]   =					15
             SCREAMING_SNAKE_CASE__     :				int   =					min_dist_processor(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
             self.assertFalse(jnp.isinf(lowerCAmelCase_    ).any()    )

       def      A_  (      self				:     str    )	->List[Any]:
             SCREAMING_SNAKE_CASE__     :				Tuple   =					20
             SCREAMING_SNAKE_CASE__     :				Tuple   =					4
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					0

             SCREAMING_SNAKE_CASE__     :				Optional[int]   =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_    )

             # check that all scores are -inf except the bos_token_id score
             SCREAMING_SNAKE_CASE__     :				Tuple   =					ids_tensor((batch_size, 1)							,	vocab_size=20    )
             SCREAMING_SNAKE_CASE__     :				str   =					1
             SCREAMING_SNAKE_CASE__     :				Tuple   =					self._get_uniform_logits(lowerCAmelCase_							,	lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Tuple   =					logits_processor(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
             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
             SCREAMING_SNAKE_CASE__     :				Tuple   =					3
             SCREAMING_SNAKE_CASE__     :				Dict   =					self._get_uniform_logits(lowerCAmelCase_							,	lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Dict   =					logits_processor(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
             self.assertFalse(jnp.isinf(lowerCAmelCase_    ).any()    )

       def      A_  (      self				:     Optional[Any]    )	->List[Any]:
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					20
             SCREAMING_SNAKE_CASE__     :				Dict   =					4
             SCREAMING_SNAKE_CASE__     :				List[Any]   =					0
             SCREAMING_SNAKE_CASE__     :				Dict   =					5

             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_							,	eos_token_id=lowerCAmelCase_    )

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

       def      A_  (      self				:     List[Any]    )	->Optional[Any]:
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					4
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					10
             SCREAMING_SNAKE_CASE__     :				str   =					15
             SCREAMING_SNAKE_CASE__     :				int   =					2
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					1
             SCREAMING_SNAKE_CASE__     :				Any   =					15

             # dummy input_ids and scores
             SCREAMING_SNAKE_CASE__     :				str   =					ids_tensor((batch_size, sequence_length)							,	lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				List[str]   =					input_ids.copy()

             SCREAMING_SNAKE_CASE__     :				Tuple   =					self._get_uniform_logits(lowerCAmelCase_							,	lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					scores.copy()

             # instantiate all dist processors
             SCREAMING_SNAKE_CASE__     :				Tuple   =					FlaxTemperatureLogitsWarper(temperature=0.5    )
             SCREAMING_SNAKE_CASE__     :				Optional[int]   =					FlaxTopKLogitsWarper(3    )
             SCREAMING_SNAKE_CASE__     :				Optional[int]   =					FlaxTopPLogitsWarper(0.8    )

             # instantiate all logits processors
             SCREAMING_SNAKE_CASE__     :				int   =					FlaxMinLengthLogitsProcessor(min_length=10							,	eos_token_id=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				str   =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				List[str]   =					FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_							,	eos_token_id=lowerCAmelCase_    )

             SCREAMING_SNAKE_CASE__     :				List[str]   =					10

             # no processor list
             SCREAMING_SNAKE_CASE__     :				List[Any]   =					temp_dist_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Dict   =					top_k_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Tuple   =					top_p_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				List[Any]   =					min_dist_proc(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Tuple   =					bos_dist_proc(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				List[Any]   =					eos_dist_proc(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )

             # with processor list
             SCREAMING_SNAKE_CASE__     :				int   =					FlaxLogitsProcessorList(
                 [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					processor(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )

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

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






       def      A_  (      self				:     int    )	->List[Any]:
             SCREAMING_SNAKE_CASE__     :				str   =					4
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					10
             SCREAMING_SNAKE_CASE__     :				List[Any]   =					15
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					2
             SCREAMING_SNAKE_CASE__     :				Optional[int]   =					1
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					15

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

             SCREAMING_SNAKE_CASE__     :				int   =					self._get_uniform_logits(lowerCAmelCase_							,	lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				List[str]   =					scores.copy()

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

             # instantiate all logits processors
             SCREAMING_SNAKE_CASE__     :				int   =					FlaxMinLengthLogitsProcessor(min_length=10							,	eos_token_id=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				str   =					FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_							,	eos_token_id=lowerCAmelCase_    )

             SCREAMING_SNAKE_CASE__     :				Dict   =					10

             # no processor list
             def run_no_processor_list(a				:     int							,	a				:     Dict							,	a				:     List[Any]    ):
                   SCREAMING_SNAKE_CASE__     :				Any   =					temp_dist_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
                   SCREAMING_SNAKE_CASE__     :				List[Any]   =					top_k_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
                   SCREAMING_SNAKE_CASE__     :				Tuple   =					top_p_warp(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
                   SCREAMING_SNAKE_CASE__     :				str   =					min_dist_proc(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
                   SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					bos_dist_proc(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
                   SCREAMING_SNAKE_CASE__     :				Tuple   =					eos_dist_proc(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
                   return scores

             # with processor list
             def run_processor_list(a				:     Union[str, Any]							,	a				:     int							,	a				:     Optional[Any]    ):
                   SCREAMING_SNAKE_CASE__     :				Tuple   =					FlaxLogitsProcessorList(
                       [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
                   SCREAMING_SNAKE_CASE__     :				Optional[int]   =					processor(lowerCAmelCase_							,	lowerCAmelCase_							,	cur_len=lowerCAmelCase_    )
                   return scores

             SCREAMING_SNAKE_CASE__     :				List[Any]   =					jax.jit(lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					jax.jit(lowerCAmelCase_    )

             SCREAMING_SNAKE_CASE__     :				str   =					jitted_run_no_processor_list(lowerCAmelCase_							,	lowerCAmelCase_							,	lowerCAmelCase_    )
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jitted_run_processor_list(lowerCAmelCase_							,	lowerCAmelCase_							,	lowerCAmelCase_    )

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

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






from __future__ import annotations







def       UpperCAmelCase							(     _lowerCamelCase      :  list[int | float]			,   _lowerCamelCase      :  int			,   _lowerCamelCase      :  int     ):





						'''simple docstring'''


						if len(_lowerCamelCase     ) == 0:
												raise ValueError("find_max() arg is an empty sequence"     )
						if (
						    left >= len(_lowerCamelCase     )
						    or left < -len(_lowerCamelCase     )
						    or right >= len(_lowerCamelCase     )
						    or right < -len(_lowerCamelCase     )
						):
												raise IndexError("list index out of range"     )
						if left == right:
												return nums[left]
						SCREAMING_SNAKE_CASE__     :				Optional[int]   =					(left + right) >> 1  # the middle
						SCREAMING_SNAKE_CASE__     :				List[Any]   =					find_max(_lowerCamelCase			,   _lowerCamelCase			,   _lowerCamelCase     )  # find max in range[left, mid]
						SCREAMING_SNAKE_CASE__     :				Optional[int]   =					find_max(_lowerCamelCase			,   mid + 1			,   _lowerCamelCase     )  # find max in range[mid + 1, right]

						return left_max if left_max >= right_max else right_max


if __name__ == "__main__":
				import doctest

				doctest.testmod(verbose=True)