



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

    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int):







      """simple docstring"""







      lowercase_								= jnp.ones((batch_size, length)) / length
      return scores

    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      lowercase_								= None
      lowercase_								= 2_0

      lowercase_								= self._get_uniform_logits(batch_size=2 ,  length=lowerCAmelCase_)

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

      # compute softmax
      lowercase_								= jax.nn.softmax(lowerCAmelCase_ ,  axis=-1)

      lowercase_								= FlaxTemperatureLogitsWarper(temperature=0.5)
      lowercase_								= FlaxTemperatureLogitsWarper(temperature=1.3)

      lowercase_								= jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase_ ,  scores.copy() ,  cur_len=lowerCAmelCase_) ,  axis=-1)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      lowercase_								= None
      lowercase_								= 1_0
      lowercase_								= 2

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

      lowercase_								= FlaxTopKLogitsWarper(3)

      lowercase_								= 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
      lowercase_								= 5
      lowercase_								= FlaxTopKLogitsWarper(top_k=1 ,  filter_value=0.0 ,  min_tokens_to_keep=3)

      lowercase_								= np.broadcast_to(np.arange(lowerCAmelCase_)[None, :] ,  (batch_size, length)).copy()
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= None
      lowercase_								= 1_0
      lowercase_								= 2

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

      lowercase_								= FlaxTopPLogitsWarper(0.8)
      lowercase_								= 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
      lowercase_								= 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
      lowercase_								= np.broadcast_to(np.arange(lowerCAmelCase_)[None, :] ,  (batch_size, vocab_size)).copy() - (
          vocab_size // 2
      )

      # make ramp_logits more extreme
      lowercase_								= ramp_logits[1] * 100.0

      # make sure at least 2 tokens are kept
      lowercase_								= FlaxTopPLogitsWarper(0.9 ,  min_tokens_to_keep=2 ,  filter_value=0.0)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= 2_0
      lowercase_								= 4
      lowercase_								= 0

      lowercase_								= FlaxMinLengthLogitsProcessor(min_length=1_0 ,  eos_token_id=lowerCAmelCase_)

      # check that min length is applied at length 5
      lowercase_								= ids_tensor((batch_size, 2_0) ,  vocab_size=2_0)
      lowercase_								= 5
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= 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
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= 1_5
      lowercase_								= min_dist_processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      self.assertFalse(jnp.isinf(lowerCAmelCase_).any())

    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= 2_0
      lowercase_								= 4
      lowercase_								= 0

      lowercase_								= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_)

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

    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= 2_0
      lowercase_								= 4
      lowercase_								= 0
      lowercase_								= 5

      lowercase_								= FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_ ,  eos_token_id=lowerCAmelCase_)

      # check that all scores are -inf except the eos_token_id when max_length is reached
      lowercase_								= ids_tensor((batch_size, 4) ,  vocab_size=2_0)
      lowercase_								= 4
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= 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
      lowercase_								= 3
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= logits_processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      self.assertFalse(jnp.isinf(lowerCAmelCase_).any())

    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      lowercase_								= 4
      lowercase_								= 1_0
      lowercase_								= 1_5
      lowercase_								= 2
      lowercase_								= 1
      lowercase_								= 1_5

      # dummy input_ids and scores
      lowercase_								= ids_tensor((batch_size, sequence_length) ,  lowerCAmelCase_)
      lowercase_								= input_ids.copy()

      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= scores.copy()

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

      # instantiate all logits processors
      lowercase_								= FlaxMinLengthLogitsProcessor(min_length=1_0 ,  eos_token_id=lowerCAmelCase_)
      lowercase_								= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_)
      lowercase_								= FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_ ,  eos_token_id=lowerCAmelCase_)

      lowercase_								= 1_0

      # no processor list
      lowercase_								= temp_dist_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= top_k_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= top_p_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= min_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= bos_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= eos_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)

      # with processor list
      lowercase_								= FlaxLogitsProcessorList(
          [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= 4
      lowercase_								= 1_0
      lowercase_								= 1_5
      lowercase_								= 2
      lowercase_								= 1
      lowercase_								= 1_5

      # dummy input_ids and scores
      lowercase_								= ids_tensor((batch_size, sequence_length) ,  lowerCAmelCase_)
      lowercase_								= input_ids.copy()

      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= scores.copy()

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

      # instantiate all logits processors
      lowercase_								= FlaxMinLengthLogitsProcessor(min_length=1_0 ,  eos_token_id=lowerCAmelCase_)
      lowercase_								= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_)
      lowercase_								= FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_ ,  eos_token_id=lowerCAmelCase_)

      lowercase_								= 1_0

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

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

      lowercase_								= jax.jit(lowerCAmelCase_)
      lowercase_								= jax.jit(lowerCAmelCase_)

      lowercase_								= jitted_run_no_processor_list(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= 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())






"""simple docstring"""
def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase=2_81_23							) ->						Optional[Any]:







  '''simple docstring'''


  lowercase_								= [1] * (limit + 1)

  for i in range(2 ,  int(limit**0.5							) + 1							):
    sum_divs[i * i] += i
    for k in range(i + 1 ,  limit // i + 1							):
      sum_divs[k * i] += k + i

  lowercase_								= set()
  lowercase_								= 0

  for n in range(1 ,  limit + 1							):
    if sum_divs[n] > n:
      abundants.add(__lowerCAmelCase							)

    if not any((n - a in abundants) for a in abundants							):
      res += n

  return res


if __name__ == "__main__":
       print(solution())

