


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







    """simple docstring"""




    def 	__snake_case							(					self				:    Union[str, Any]				, __UpperCAmelCase				:    int				, __UpperCAmelCase				:    int):
      a						:		str      							=			jnp.ones((batch_size, length)) / length
      return scores






    def 	__snake_case							(					self				:    Optional[Any]):
      a						:		Optional[int]      							=			None
      a						:		Optional[Any]      							=			20

      a						:		Optional[int]      							=			self._get_uniform_logits(batch_size=2				, length=__UpperCAmelCase)

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

      # compute softmax
      a						:		Dict      							=			jax.nn.softmax(__UpperCAmelCase				, axis=-1)

      a						:		Optional[Any]      							=			FlaxTemperatureLogitsWarper(temperature=0.5)
      a						:		Dict      							=			FlaxTemperatureLogitsWarper(temperature=1.3)

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

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

      a						:		str      							=			FlaxTopKLogitsWarper(3)

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

      a						:		Dict      							=			np.broadcast_to(np.arange(__UpperCAmelCase)[None, :]				, (batch_size, length)).copy()
      a						:		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				:    Tuple):
      a						:		List[Any]      							=			None
      a						:		int      							=			10
      a						:		Dict      							=			2

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

      a						:		List[Any]      							=			FlaxTopPLogitsWarper(0.8)
      a						:		List[str]      							=			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
      a						:		str      							=			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
      a						:		int      							=			np.broadcast_to(np.arange(__UpperCAmelCase)[None, :]				, (batch_size, vocab_size)).copy() - (
          vocab_size // 2
      )

      # make ramp_logits more extreme
      a						:		Optional[int]      							=			ramp_logits[1] * 100.0

      # make sure at least 2 tokens are kept
      a						:		int      							=			FlaxTopPLogitsWarper(0.9				, min_tokens_to_keep=2				, filter_value=0.0)
      a						:		str      							=			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				:    List[str]):
      a						:		Union[str, Any]      							=			20
      a						:		Any      							=			4
      a						:		Any      							=			0

      a						:		List[str]      							=			FlaxMinLengthLogitsProcessor(min_length=10				, eos_token_id=__UpperCAmelCase)

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






    def 	__snake_case							(					self				:    str):
      a						:		Dict      							=			20
      a						:		Any      							=			4
      a						:		str      							=			0

      a						:		Optional[int]      							=			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__UpperCAmelCase)

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






    def 	__snake_case							(					self				:    Dict):
      a						:		List[Any]      							=			20
      a						:		Tuple      							=			4
      a						:		List[Any]      							=			0
      a						:		str      							=			5

      a						:		List[str]      							=			FlaxForcedEOSTokenLogitsProcessor(max_length=__UpperCAmelCase				, eos_token_id=__UpperCAmelCase)

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






    def 	__snake_case							(					self				:    Optional[Any]):
      a						:		Optional[Any]      							=			4
      a						:		int      							=			10
      a						:		int      							=			15
      a						:		int      							=			2
      a						:		List[str]      							=			1
      a						:		Dict      							=			15

      # dummy input_ids and scores
      a						:		Any      							=			ids_tensor((batch_size, sequence_length)				, __UpperCAmelCase)
      a						:		List[Any]      							=			input_ids.copy()

      a						:		Optional[int]      							=			self._get_uniform_logits(__UpperCAmelCase				, __UpperCAmelCase)
      a						:		Dict      							=			scores.copy()

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

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

      a						:		Tuple      							=			10

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

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

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

      a						:		Union[str, Any]      							=			self._get_uniform_logits(__UpperCAmelCase				, __UpperCAmelCase)
      a						:		Optional[int]      							=			scores.copy()

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

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

      a						:		Optional[int]      							=			10

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

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

      a						:		List[Any]      							=			jax.jit(__UpperCAmelCase)
      a						:		Union[str, Any]      							=			jax.jit(__UpperCAmelCase)

      a						:		List[Any]      							=			jitted_run_no_processor_list(__UpperCAmelCase				, __UpperCAmelCase				, __UpperCAmelCase)
      a						:		List[str]      							=			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 			snake_case_   (  SCREAMING_SNAKE_CASE__							,      SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				return int((input_a, input_a).count(0     ) != 0     )







def 			snake_case_   (  ):







				"""simple docstring"""





				assert nand_gate(0							,      0     ) == 1
				assert nand_gate(0							,      1     ) == 1
				assert nand_gate(1							,      0     ) == 1
				assert nand_gate(1							,      1     ) == 0


if __name__ == "__main__":
				print(nand_gate(0, 0))
				print(nand_gate(0, 1))
				print(nand_gate(1, 0))
				print(nand_gate(1, 1))




