





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


				def _snake_case    (						self				,	UpperCAmelCase_				,	UpperCAmelCase_							):
								lowerCamelCase			=jnp.ones((batch_size, length)							) / length
								return scores


				def _snake_case    (						self							):
								lowerCamelCase			=None
								lowerCamelCase			=20

								lowerCamelCase			=self._get_uniform_logits(batch_size=2				,	length=UpperCAmelCase_							)

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

								# compute softmax
								lowerCamelCase			=jax.nn.softmax(UpperCAmelCase_				,	axis=-1							)

								lowerCamelCase			=FlaxTemperatureLogitsWarper(temperature=0.5							)
								lowerCamelCase			=FlaxTemperatureLogitsWarper(temperature=1.3							)

								lowerCamelCase			=jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase_				,	scores.copy()				,	cur_len=UpperCAmelCase_							)				,	axis=-1							)
								lowerCamelCase			=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							):
								lowerCamelCase			=None
								lowerCamelCase			=10
								lowerCamelCase			=2

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

								lowerCamelCase			=FlaxTopKLogitsWarper(3							)

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

								lowerCamelCase			=np.broadcast_to(np.arange(UpperCAmelCase_							)[None, :]				,	(batch_size, length)							).copy()
								lowerCamelCase			=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							):
								lowerCamelCase			=None
								lowerCamelCase			=10
								lowerCamelCase			=2

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

								lowerCamelCase			=FlaxTopPLogitsWarper(0.8							)
								lowerCamelCase			=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
								lowerCamelCase			=np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]							)
								self.assertTrue(np.allclose(UpperCAmelCase_				,	UpperCAmelCase_				,	atol=1E-3							)							)

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

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

								# make sure at least 2 tokens are kept
								lowerCamelCase			=FlaxTopPLogitsWarper(0.9				,	min_tokens_to_keep=2				,	filter_value=0.0							)
								lowerCamelCase			=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							):
								lowerCamelCase			=20
								lowerCamelCase			=4
								lowerCamelCase			=0

								lowerCamelCase			=FlaxMinLengthLogitsProcessor(min_length=10				,	eos_token_id=UpperCAmelCase_							)

								# check that min length is applied at length 5
								lowerCamelCase			=ids_tensor((batch_size, 20)				,	vocab_size=20							)
								lowerCamelCase			=5
								lowerCamelCase			=self._get_uniform_logits(UpperCAmelCase_				,	UpperCAmelCase_							)
								lowerCamelCase			=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
								lowerCamelCase			=self._get_uniform_logits(UpperCAmelCase_				,	UpperCAmelCase_							)
								lowerCamelCase			=15
								lowerCamelCase			=min_dist_processor(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
								self.assertFalse(jnp.isinf(UpperCAmelCase_							).any()							)


				def _snake_case    (						self							):
								lowerCamelCase			=20
								lowerCamelCase			=4
								lowerCamelCase			=0

								lowerCamelCase			=FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase_							)

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


				def _snake_case    (						self							):
								lowerCamelCase			=20
								lowerCamelCase			=4
								lowerCamelCase			=0
								lowerCamelCase			=5

								lowerCamelCase			=FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase_				,	eos_token_id=UpperCAmelCase_							)

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


				def _snake_case    (						self							):
								lowerCamelCase			=4
								lowerCamelCase			=10
								lowerCamelCase			=15
								lowerCamelCase			=2
								lowerCamelCase			=1
								lowerCamelCase			=15

								# dummy input_ids and scores
								lowerCamelCase			=ids_tensor((batch_size, sequence_length)				,	UpperCAmelCase_							)
								lowerCamelCase			=input_ids.copy()

								lowerCamelCase			=self._get_uniform_logits(UpperCAmelCase_				,	UpperCAmelCase_							)
								lowerCamelCase			=scores.copy()

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

								# instantiate all logits processors
								lowerCamelCase			=FlaxMinLengthLogitsProcessor(min_length=10				,	eos_token_id=UpperCAmelCase_							)
								lowerCamelCase			=FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase_							)
								lowerCamelCase			=FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase_				,	eos_token_id=UpperCAmelCase_							)

								lowerCamelCase			=10

								# no processor list
								lowerCamelCase			=temp_dist_warp(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
								lowerCamelCase			=top_k_warp(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
								lowerCamelCase			=top_p_warp(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
								lowerCamelCase			=min_dist_proc(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
								lowerCamelCase			=bos_dist_proc(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
								lowerCamelCase			=eos_dist_proc(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)

								# with processor list
								lowerCamelCase			=FlaxLogitsProcessorList(
								    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
								lowerCamelCase			=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							):
								lowerCamelCase			=4
								lowerCamelCase			=10
								lowerCamelCase			=15
								lowerCamelCase			=2
								lowerCamelCase			=1
								lowerCamelCase			=15

								# dummy input_ids and scores
								lowerCamelCase			=ids_tensor((batch_size, sequence_length)				,	UpperCAmelCase_							)
								lowerCamelCase			=input_ids.copy()

								lowerCamelCase			=self._get_uniform_logits(UpperCAmelCase_				,	UpperCAmelCase_							)
								lowerCamelCase			=scores.copy()

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

								# instantiate all logits processors
								lowerCamelCase			=FlaxMinLengthLogitsProcessor(min_length=10				,	eos_token_id=UpperCAmelCase_							)
								lowerCamelCase			=FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase_							)
								lowerCamelCase			=FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase_				,	eos_token_id=UpperCAmelCase_							)

								lowerCamelCase			=10

								# no processor list
								def run_no_processor_list(UpperCAmelCase_				,	UpperCAmelCase_				,	UpperCAmelCase_							):
												lowerCamelCase			=temp_dist_warp(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
												lowerCamelCase			=top_k_warp(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
												lowerCamelCase			=top_p_warp(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
												lowerCamelCase			=min_dist_proc(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
												lowerCamelCase			=bos_dist_proc(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
												lowerCamelCase			=eos_dist_proc(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
												return scores

								# with processor list
								def run_processor_list(UpperCAmelCase_				,	UpperCAmelCase_				,	UpperCAmelCase_							):
												lowerCamelCase			=FlaxLogitsProcessorList(
												    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
												lowerCamelCase			=processor(UpperCAmelCase_				,	UpperCAmelCase_				,	cur_len=UpperCAmelCase_							)
												return scores

								lowerCamelCase			=jax.jit(UpperCAmelCase_							)
								lowerCamelCase			=jax.jit(UpperCAmelCase_							)

								lowerCamelCase			=jitted_run_no_processor_list(UpperCAmelCase_				,	UpperCAmelCase_				,	UpperCAmelCase_							)
								lowerCamelCase			=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()							)












from math import cos, sin, sqrt, tau

from audio_filters.iir_filter import IIRFilter





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase = 1 / sqrt(2			)			)  ->					IIRFilter:
				lowerCamelCase			=tau * frequency / samplerate
				lowerCamelCase			=sin(_UpperCAmelCase			)
				lowerCamelCase			=cos(_UpperCAmelCase			)
				lowerCamelCase			=_sin / (2 * q_factor)

				lowerCamelCase			=(1 - _cos) / 2
				lowerCamelCase			=1 - _cos

				lowerCamelCase			=1 + alpha
				lowerCamelCase			=-2 * _cos
				lowerCamelCase			=1 - alpha

				lowerCamelCase			=IIRFilter(2			)
				filt.set_coefficients([aa, aa, aa]     , [ba, ba, ba]			)
				return filt





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase = 1 / sqrt(2			)			)  ->					IIRFilter:
				lowerCamelCase			=tau * frequency / samplerate
				lowerCamelCase			=sin(_UpperCAmelCase			)
				lowerCamelCase			=cos(_UpperCAmelCase			)
				lowerCamelCase			=_sin / (2 * q_factor)

				lowerCamelCase			=(1 + _cos) / 2
				lowerCamelCase			=-1 - _cos

				lowerCamelCase			=1 + alpha
				lowerCamelCase			=-2 * _cos
				lowerCamelCase			=1 - alpha

				lowerCamelCase			=IIRFilter(2			)
				filt.set_coefficients([aa, aa, aa]     , [ba, ba, ba]			)
				return filt





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase = 1 / sqrt(2			)			)  ->					IIRFilter:
				lowerCamelCase			=tau * frequency / samplerate
				lowerCamelCase			=sin(_UpperCAmelCase			)
				lowerCamelCase			=cos(_UpperCAmelCase			)
				lowerCamelCase			=_sin / (2 * q_factor)

				lowerCamelCase			=_sin / 2
				lowerCamelCase			=0
				lowerCamelCase			=-ba

				lowerCamelCase			=1 + alpha
				lowerCamelCase			=-2 * _cos
				lowerCamelCase			=1 - alpha

				lowerCamelCase			=IIRFilter(2			)
				filt.set_coefficients([aa, aa, aa]     , [ba, ba, ba]			)
				return filt





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase = 1 / sqrt(2			)			)  ->					IIRFilter:
				lowerCamelCase			=tau * frequency / samplerate
				lowerCamelCase			=sin(_UpperCAmelCase			)
				lowerCamelCase			=cos(_UpperCAmelCase			)
				lowerCamelCase			=_sin / (2 * q_factor)

				lowerCamelCase			=1 - alpha
				lowerCamelCase			=-2 * _cos
				lowerCamelCase			=1 + alpha

				lowerCamelCase			=IIRFilter(2			)
				filt.set_coefficients([ba, ba, ba]     , [ba, ba, ba]			)
				return filt





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase = 1 / sqrt(2			)     , )  ->					IIRFilter:
				lowerCamelCase			=tau * frequency / samplerate
				lowerCamelCase			=sin(_UpperCAmelCase			)
				lowerCamelCase			=cos(_UpperCAmelCase			)
				lowerCamelCase			=_sin / (2 * q_factor)
				lowerCamelCase			=10 ** (gain_db / 40)

				lowerCamelCase			=1 + alpha * big_a
				lowerCamelCase			=-2 * _cos
				lowerCamelCase			=1 - alpha * big_a
				lowerCamelCase			=1 + alpha / big_a
				lowerCamelCase			=-2 * _cos
				lowerCamelCase			=1 - alpha / big_a

				lowerCamelCase			=IIRFilter(2			)
				filt.set_coefficients([aa, aa, aa]     , [ba, ba, ba]			)
				return filt





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase = 1 / sqrt(2			)     , )  ->					IIRFilter:
				lowerCamelCase			=tau * frequency / samplerate
				lowerCamelCase			=sin(_UpperCAmelCase			)
				lowerCamelCase			=cos(_UpperCAmelCase			)
				lowerCamelCase			=_sin / (2 * q_factor)
				lowerCamelCase			=10 ** (gain_db / 40)
				lowerCamelCase			=(big_a + 1) - (big_a - 1) * _cos
				lowerCamelCase			=(big_a + 1) + (big_a - 1) * _cos
				lowerCamelCase			=(big_a - 1) - (big_a + 1) * _cos
				lowerCamelCase			=(big_a - 1) + (big_a + 1) * _cos
				lowerCamelCase			=2 * sqrt(_UpperCAmelCase			) * alpha

				lowerCamelCase			=big_a * (pmc + aaa)
				lowerCamelCase			=2 * big_a * mpc
				lowerCamelCase			=big_a * (pmc - aaa)
				lowerCamelCase			=ppmc + aaa
				lowerCamelCase			=-2 * pmpc
				lowerCamelCase			=ppmc - aaa

				lowerCamelCase			=IIRFilter(2			)
				filt.set_coefficients([aa, aa, aa]     , [ba, ba, ba]			)
				return filt





def        _lowercase   (  _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase     , _UpperCAmelCase = 1 / sqrt(2			)     , )  ->					IIRFilter:
				lowerCamelCase			=tau * frequency / samplerate
				lowerCamelCase			=sin(_UpperCAmelCase			)
				lowerCamelCase			=cos(_UpperCAmelCase			)
				lowerCamelCase			=_sin / (2 * q_factor)
				lowerCamelCase			=10 ** (gain_db / 40)
				lowerCamelCase			=(big_a + 1) - (big_a - 1) * _cos
				lowerCamelCase			=(big_a + 1) + (big_a - 1) * _cos
				lowerCamelCase			=(big_a - 1) - (big_a + 1) * _cos
				lowerCamelCase			=(big_a - 1) + (big_a + 1) * _cos
				lowerCamelCase			=2 * sqrt(_UpperCAmelCase			) * alpha

				lowerCamelCase			=big_a * (ppmc + aaa)
				lowerCamelCase			=-2 * big_a * pmpc
				lowerCamelCase			=big_a * (ppmc - aaa)
				lowerCamelCase			=pmc + aaa
				lowerCamelCase			=2 * mpc
				lowerCamelCase			=pmc - aaa

				lowerCamelCase			=IIRFilter(2			)
				filt.set_coefficients([aa, aa, aa]     , [ba, ba, ba]			)
				return filt





