
'''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 ):



    def   lowerCamelCase_  (	self      :       List[Any],__A      :       int,__A      :       int     ):
          _lowerCamelCase				:   Dict								=   jnp.ones((batch_size, length)     ) / length
          return scores



    def   lowerCamelCase_  (	self      :       int     ):
          _lowerCamelCase				:   int								=   None
          _lowerCamelCase				:   Any								=   2_0

          _lowerCamelCase				:   Optional[int]								=   self._get_uniform_logits(batch_size=2,length=__A     )

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

          # compute softmax
          _lowerCamelCase				:   Union[str, Any]								=   jax.nn.softmax(__A,axis=-1     )

          _lowerCamelCase				:   List[str]								=   FlaxTemperatureLogitsWarper(temperature=0.5     )
          _lowerCamelCase				:   Dict								=   FlaxTemperatureLogitsWarper(temperature=1.3     )

          _lowerCamelCase				:   Optional[Any]								=   jax.nn.softmax(temp_dist_warper_sharper(__A,scores.copy(),cur_len=__A     ),axis=-1     )
          _lowerCamelCase				:   int								=   jax.nn.softmax(temp_dist_warper_smoother(__A,scores.copy(),cur_len=__A     ),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   lowerCamelCase_  (	self      :       List[Any]     ):
          _lowerCamelCase				:   Dict								=   None
          _lowerCamelCase				:   Union[str, Any]								=   1_0
          _lowerCamelCase				:   List[Any]								=   2

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

          _lowerCamelCase				:   List[str]								=   FlaxTopKLogitsWarper(3     )

          _lowerCamelCase				:   Dict								=   top_k_warp(__A,__A,cur_len=__A     )

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

          _lowerCamelCase				:   str								=   np.broadcast_to(np.arange(__A     )[None, :],(batch_size, length)     ).copy()
          _lowerCamelCase				:   str								=   top_k_warp_safety_check(__A,__A,cur_len=__A     )

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



    def   lowerCamelCase_  (	self      :       Tuple     ):
          _lowerCamelCase				:   List[Any]								=   None
          _lowerCamelCase				:   int								=   1_0
          _lowerCamelCase				:   Any								=   2

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

          _lowerCamelCase				:   Optional[int]								=   FlaxTopPLogitsWarper(0.8     )
          _lowerCamelCase				:   List[Any]								=   np.exp(top_p_warp(__A,__A,cur_len=__A     )     )

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

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

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

          # make sure at least 2 tokens are kept
          _lowerCamelCase				:   Tuple								=   FlaxTopPLogitsWarper(0.9,min_tokens_to_keep=2,filter_value=0.0     )
          _lowerCamelCase				:   Optional[Any]								=   top_p_warp(__A,__A,cur_len=__A     )

          # 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   lowerCamelCase_  (	self      :       Any     ):
          _lowerCamelCase				:   Union[str, Any]								=   2_0
          _lowerCamelCase				:   Optional[Any]								=   4
          _lowerCamelCase				:   Union[str, Any]								=   0

          _lowerCamelCase				:   Any								=   FlaxMinLengthLogitsProcessor(min_length=1_0,eos_token_id=__A     )

          # check that min length is applied at length 5
          _lowerCamelCase				:   Dict								=   ids_tensor((batch_size, 2_0),vocab_size=2_0     )
          _lowerCamelCase				:   int								=   5
          _lowerCamelCase				:   List[Any]								=   self._get_uniform_logits(__A,__A     )
          _lowerCamelCase				:   Dict								=   min_dist_processor(__A,__A,cur_len=__A     )
          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				:   List[str]								=   self._get_uniform_logits(__A,__A     )
          _lowerCamelCase				:   Any								=   1_5
          _lowerCamelCase				:   List[str]								=   min_dist_processor(__A,__A,cur_len=__A     )
          self.assertFalse(jnp.isinf(__A     ).any()     )



    def   lowerCamelCase_  (	self      :       Union[str, Any]     ):
          _lowerCamelCase				:   List[Any]								=   2_0
          _lowerCamelCase				:   Optional[int]								=   4
          _lowerCamelCase				:   Any								=   0

          _lowerCamelCase				:   str								=   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__A     )

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



    def   lowerCamelCase_  (	self      :       Any     ):
          _lowerCamelCase				:   str								=   2_0
          _lowerCamelCase				:   List[str]								=   4
          _lowerCamelCase				:   List[str]								=   0
          _lowerCamelCase				:   Any								=   5

          _lowerCamelCase				:   Union[str, Any]								=   FlaxForcedEOSTokenLogitsProcessor(max_length=__A,eos_token_id=__A     )

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



    def   lowerCamelCase_  (	self      :       int     ):
          _lowerCamelCase				:   List[str]								=   4
          _lowerCamelCase				:   Any								=   1_0
          _lowerCamelCase				:   int								=   1_5
          _lowerCamelCase				:   List[Any]								=   2
          _lowerCamelCase				:   Tuple								=   1
          _lowerCamelCase				:   str								=   1_5

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

          _lowerCamelCase				:   Dict								=   self._get_uniform_logits(__A,__A     )
          _lowerCamelCase				:   Optional[Any]								=   scores.copy()

          # instantiate all dist processors
          _lowerCamelCase				:   List[Any]								=   FlaxTemperatureLogitsWarper(temperature=0.5     )
          _lowerCamelCase				:   Union[str, Any]								=   FlaxTopKLogitsWarper(3     )
          _lowerCamelCase				:   List[Any]								=   FlaxTopPLogitsWarper(0.8     )

          # instantiate all logits processors
          _lowerCamelCase				:   Optional[Any]								=   FlaxMinLengthLogitsProcessor(min_length=1_0,eos_token_id=__A     )
          _lowerCamelCase				:   List[Any]								=   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__A     )
          _lowerCamelCase				:   Dict								=   FlaxForcedEOSTokenLogitsProcessor(max_length=__A,eos_token_id=__A     )

          _lowerCamelCase				:   Tuple								=   1_0

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

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

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

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



    def   lowerCamelCase_  (	self      :       Optional[int]     ):
          _lowerCamelCase				:   Dict								=   4
          _lowerCamelCase				:   List[str]								=   1_0
          _lowerCamelCase				:   Any								=   1_5
          _lowerCamelCase				:   Dict								=   2
          _lowerCamelCase				:   Optional[Any]								=   1
          _lowerCamelCase				:   Any								=   1_5

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

          _lowerCamelCase				:   Optional[Any]								=   self._get_uniform_logits(__A,__A     )
          _lowerCamelCase				:   Optional[Any]								=   scores.copy()

          # instantiate all dist processors
          _lowerCamelCase				:   Any								=   FlaxTemperatureLogitsWarper(temperature=0.5     )
          _lowerCamelCase				:   List[Any]								=   FlaxTopKLogitsWarper(3     )
          _lowerCamelCase				:   Dict								=   FlaxTopPLogitsWarper(0.8     )

          # instantiate all logits processors
          _lowerCamelCase				:   List[Any]								=   FlaxMinLengthLogitsProcessor(min_length=1_0,eos_token_id=__A     )
          _lowerCamelCase				:   List[str]								=   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__A     )
          _lowerCamelCase				:   List[str]								=   FlaxForcedEOSTokenLogitsProcessor(max_length=__A,eos_token_id=__A     )

          _lowerCamelCase				:   List[str]								=   1_0

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

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

          _lowerCamelCase				:   str								=   jax.jit(__A     )
          _lowerCamelCase				:   Optional[Any]								=   jax.jit(__A     )

          _lowerCamelCase				:   Optional[Any]								=   jitted_run_no_processor_list(__A,__A,__A     )
          _lowerCamelCase				:   Any								=   jitted_run_processor_list(__A,__A,__A     )

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

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


'''simple docstring'''






import math







def 		A_       (	_lowerCAmelCase   :    int							):



      """simple docstring"""




      _lowerCamelCase				:   Optional[int]								=   math.loga(math.sqrt(4 * positive_integer + 1							) / 2 + 1 / 2							)

      return exponent == int(_lowerCAmelCase							)







def 		A_       (	_lowerCAmelCase   :    float = 1 / 12345							):



      """simple docstring"""




      _lowerCamelCase				:   Optional[int]								=   0
      _lowerCamelCase				:   int								=   0

      _lowerCamelCase				:   List[str]								=   3
      while True:
            _lowerCamelCase				:   List[Any]								=   (integer**2 - 1) / 4
            # if candidate is an integer, then there is a partition for k
            if partition_candidate == int(_lowerCAmelCase							):
                  _lowerCamelCase				:   Any								=   int(_lowerCAmelCase							)
                  total_partitions += 1
                  if check_partition_perfect(_lowerCAmelCase							):
                        perfect_partitions += 1
            if perfect_partitions > 0:
                  if perfect_partitions / total_partitions < max_proportion:
                        return int(_lowerCAmelCase							)
            integer += 1


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