


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




     def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase):
        UpperCAmelCase__					:   Optional[int]       				=       jnp.ones((batch_size, length)) / length
        return scores




     def 			snake_case__  ( self):
        UpperCAmelCase__					:   str       				=       None
        UpperCAmelCase__					:   int       				=       20

        UpperCAmelCase__					:   str       				=       self._get_uniform_logits(batch_size=2							,							length=UpperCAmelCase__)

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

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

        UpperCAmelCase__					:   Any       				=       FlaxTemperatureLogitsWarper(temperature=0.5)
        UpperCAmelCase__					:   str       				=       FlaxTemperatureLogitsWarper(temperature=1.3)

        UpperCAmelCase__					:   List[Any]       				=       jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase__							,							scores.copy()							,							cur_len=UpperCAmelCase__)							,							axis=-1)
        UpperCAmelCase__					:   Any       				=       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 			snake_case__  ( self):
        UpperCAmelCase__					:   Any       				=       None
        UpperCAmelCase__					:   List[Any]       				=       10
        UpperCAmelCase__					:   Union[str, Any]       				=       2

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

        UpperCAmelCase__					:   Dict       				=       FlaxTopKLogitsWarper(3)

        UpperCAmelCase__					:   Union[str, Any]       				=       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
        UpperCAmelCase__					:   Dict       				=       5
        UpperCAmelCase__					:   Optional[int]       				=       FlaxTopKLogitsWarper(top_k=1							,							filter_value=0.0							,							min_tokens_to_keep=3)

        UpperCAmelCase__					:   int       				=       np.broadcast_to(np.arange(UpperCAmelCase__)[None, :]							,							(batch_size, length)).copy()
        UpperCAmelCase__					:   Union[str, Any]       				=       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 			snake_case__  ( self):
        UpperCAmelCase__					:   Dict       				=       None
        UpperCAmelCase__					:   Any       				=       10
        UpperCAmelCase__					:   List[str]       				=       2

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

        UpperCAmelCase__					:   Optional[int]       				=       FlaxTopPLogitsWarper(0.8)
        UpperCAmelCase__					:   Union[str, Any]       				=       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
        UpperCAmelCase__					:   Any       				=       np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]])
        self.assertTrue(np.allclose(UpperCAmelCase__							,							UpperCAmelCase__							,							atol=1e-3))

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

        # make ramp_logits more extreme
        UpperCAmelCase__					:   List[str]       				=       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__					:   Union[str, Any]       				=       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 			snake_case__  ( self):
        UpperCAmelCase__					:   Optional[int]       				=       20
        UpperCAmelCase__					:   List[Any]       				=       4
        UpperCAmelCase__					:   int       				=       0

        UpperCAmelCase__					:   Any       				=       FlaxMinLengthLogitsProcessor(min_length=10							,							eos_token_id=UpperCAmelCase__)

        # check that min length is applied at length 5
        UpperCAmelCase__					:   Any       				=       ids_tensor((batch_size, 20)							,							vocab_size=20)
        UpperCAmelCase__					:   Dict       				=       5
        UpperCAmelCase__					:   Union[str, Any]       				=       self._get_uniform_logits(UpperCAmelCase__							,							UpperCAmelCase__)
        UpperCAmelCase__					:   List[Any]       				=       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
        UpperCAmelCase__					:   Optional[Any]       				=       self._get_uniform_logits(UpperCAmelCase__							,							UpperCAmelCase__)
        UpperCAmelCase__					:   Dict       				=       15
        UpperCAmelCase__					:   Optional[int]       				=       min_dist_processor(UpperCAmelCase__							,							UpperCAmelCase__							,							cur_len=UpperCAmelCase__)
        self.assertFalse(jnp.isinf(UpperCAmelCase__).any())




     def 			snake_case__  ( self):
        UpperCAmelCase__					:   Tuple       				=       20
        UpperCAmelCase__					:   Any       				=       4
        UpperCAmelCase__					:   Any       				=       0

        UpperCAmelCase__					:   str       				=       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__)

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




     def 			snake_case__  ( self):
        UpperCAmelCase__					:   str       				=       20
        UpperCAmelCase__					:   str       				=       4
        UpperCAmelCase__					:   Dict       				=       0
        UpperCAmelCase__					:   Optional[int]       				=       5

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

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




     def 			snake_case__  ( self):
        UpperCAmelCase__					:   Optional[Any]       				=       4
        UpperCAmelCase__					:   Tuple       				=       10
        UpperCAmelCase__					:   Any       				=       15
        UpperCAmelCase__					:   List[str]       				=       2
        UpperCAmelCase__					:   str       				=       1
        UpperCAmelCase__					:   Tuple       				=       15

        # dummy input_ids and scores
        UpperCAmelCase__					:   Optional[int]       				=       ids_tensor((batch_size, sequence_length)							,							UpperCAmelCase__)
        UpperCAmelCase__					:   List[str]       				=       input_ids.copy()

        UpperCAmelCase__					:   int       				=       self._get_uniform_logits(UpperCAmelCase__							,							UpperCAmelCase__)
        UpperCAmelCase__					:   Optional[Any]       				=       scores.copy()

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

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

        UpperCAmelCase__					:   int       				=       10

        # no processor list
        UpperCAmelCase__					:   str       				=       temp_dist_warp(UpperCAmelCase__							,							UpperCAmelCase__							,							cur_len=UpperCAmelCase__)
        UpperCAmelCase__					:   Optional[Any]       				=       top_k_warp(UpperCAmelCase__							,							UpperCAmelCase__							,							cur_len=UpperCAmelCase__)
        UpperCAmelCase__					:   int       				=       top_p_warp(UpperCAmelCase__							,							UpperCAmelCase__							,							cur_len=UpperCAmelCase__)
        UpperCAmelCase__					:   Tuple       				=       min_dist_proc(UpperCAmelCase__							,							UpperCAmelCase__							,							cur_len=UpperCAmelCase__)
        UpperCAmelCase__					:   Optional[Any]       				=       bos_dist_proc(UpperCAmelCase__							,							UpperCAmelCase__							,							cur_len=UpperCAmelCase__)
        UpperCAmelCase__					:   str       				=       eos_dist_proc(UpperCAmelCase__							,							UpperCAmelCase__							,							cur_len=UpperCAmelCase__)

        # with processor list
        UpperCAmelCase__					:   int       				=       FlaxLogitsProcessorList(
            [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
        UpperCAmelCase__					:   Optional[Any]       				=       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 			snake_case__  ( self):
        UpperCAmelCase__					:   Tuple       				=       4
        UpperCAmelCase__					:   int       				=       10
        UpperCAmelCase__					:   Any       				=       15
        UpperCAmelCase__					:   str       				=       2
        UpperCAmelCase__					:   Optional[int]       				=       1
        UpperCAmelCase__					:   Optional[Any]       				=       15

        # dummy input_ids and scores
        UpperCAmelCase__					:   Dict       				=       ids_tensor((batch_size, sequence_length)							,							UpperCAmelCase__)
        UpperCAmelCase__					:   Tuple       				=       input_ids.copy()

        UpperCAmelCase__					:   Tuple       				=       self._get_uniform_logits(UpperCAmelCase__							,							UpperCAmelCase__)
        UpperCAmelCase__					:   str       				=       scores.copy()

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

        # instantiate all logits processors
        UpperCAmelCase__					:   int       				=       FlaxMinLengthLogitsProcessor(min_length=10							,							eos_token_id=UpperCAmelCase__)
        UpperCAmelCase__					:   List[Any]       				=       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__)
        UpperCAmelCase__					:   Optional[Any]       				=       FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__							,							eos_token_id=UpperCAmelCase__)

        UpperCAmelCase__					:   List[str]       				=       10

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

        # with processor list
        def run_processor_list(_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
           UpperCAmelCase__					:   int       				=       FlaxLogitsProcessorList(
               [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
           UpperCAmelCase__					:   Dict       				=       processor(UpperCAmelCase__							,							UpperCAmelCase__							,							cur_len=UpperCAmelCase__)
           return scores

        UpperCAmelCase__					:   Tuple       				=       jax.jit(UpperCAmelCase__)
        UpperCAmelCase__					:   List[str]       				=       jax.jit(UpperCAmelCase__)

        UpperCAmelCase__					:   Tuple       				=       jitted_run_no_processor_list(UpperCAmelCase__							,							UpperCAmelCase__							,							UpperCAmelCase__)
        UpperCAmelCase__					:   Union[str, Any]       				=       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())



'''simple docstring'''



def       _UpperCamelCase (						UpperCamelCase__ = 4_0_0_0_0_0_0   ):
			UpperCAmelCase__					:   List[str]       				=       [0, 1]
			UpperCAmelCase__					:   Any       				=       0
			while fib[i] <= n:
						fib.append(fib[i] + fib[i + 1]   )
						if fib[i + 2] > n:
									break
						i += 1
			UpperCAmelCase__					:   str       				=       0
			for j in range(len(UpperCamelCase__   ) - 1   ):
						if fib[j] % 2 == 0:
									total += fib[j]

			return total


if __name__ == "__main__":
		print(f"""{solution() = }""")