

import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


__A			:					Dict						=			get_logger(__name__)


__A			:					Tuple						=			r'\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n'





class 			_SCREAMING_SNAKE_CASE :







   '''simple docstring'''




   @add_start_docstrings(__lowerCamelCase			)
   def __call__(       self					:			Union[str, Any]     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray			):
       raise NotImplementedError(
           f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."			)





class 			_SCREAMING_SNAKE_CASE :







   '''simple docstring'''




   @add_start_docstrings(__lowerCamelCase			)
   def __call__(       self					:			List[Any]     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray			):
       raise NotImplementedError(
           f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."			)





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   @add_start_docstrings(__lowerCamelCase			)
   def __call__(       self					:			Optional[Any]     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int     , **__lowerCamelCase					:			Optional[Any]			):
       for processor in self:
           SCREAMING_SNAKE_CASE										=   inspect.signature(processor.__call__			).parameters
           if len(__lowerCamelCase			) > 3:
               if not all(arg in kwargs for arg in list(function_args.keys()			)[2:]			):
                   raise ValueError(
                       f"Make sure that all the required parameters: {list(function_args.keys()			)} for "
                       f"{processor.__class__} are passed to the logits processor."			)
               SCREAMING_SNAKE_CASE										=   processor(__lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , **__lowerCamelCase			)
           else:
               SCREAMING_SNAKE_CASE										=   processor(__lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase			)
       return scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			Tuple     , __lowerCamelCase					:			float			):
       if not isinstance(__lowerCamelCase     , __lowerCamelCase			) or not (temperature > 0):
           raise ValueError(f"`temperature` has to be a strictly positive float, but is {temperature}"			)

       SCREAMING_SNAKE_CASE										=   temperature





   def __call__(       self					:			Any     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE										=   scores / self.temperature
       return scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			Optional[int]     , __lowerCamelCase					:			float     , __lowerCamelCase					:			float = -float("Inf"			)     , __lowerCamelCase					:			int = 1			):
       if not isinstance(__lowerCamelCase     , __lowerCamelCase			) or (top_p < 0 or top_p > 1.0):
           raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}"			)
       if not isinstance(__lowerCamelCase     , __lowerCamelCase			) or (min_tokens_to_keep < 1):
           raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}"			)

       SCREAMING_SNAKE_CASE										=   top_p
       SCREAMING_SNAKE_CASE										=   filter_value
       SCREAMING_SNAKE_CASE										=   min_tokens_to_keep





   def __call__(       self					:			Tuple     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   lax.top_k(__lowerCamelCase     , scores.shape[-1]			)

       SCREAMING_SNAKE_CASE										=   jnp.full_like(__lowerCamelCase     , self.filter_value			)
       SCREAMING_SNAKE_CASE										=   jax.nn.softmax(__lowerCamelCase     , axis=-1			).cumsum(axis=-1			)
       SCREAMING_SNAKE_CASE										=   cumulative_probs < self.top_p

       # include the token that is higher than top_p as well
       SCREAMING_SNAKE_CASE										=   jnp.roll(__lowerCamelCase     , 1			)
       score_mask |= score_mask.at[:, 0].set(__lowerCamelCase			)

       # min tokens to keep
       SCREAMING_SNAKE_CASE										=   score_mask.at[:, : self.min_tokens_to_keep].set(__lowerCamelCase			)

       SCREAMING_SNAKE_CASE										=   jnp.where(__lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   jax.lax.sort_key_val(__lowerCamelCase     , __lowerCamelCase			)[-1]

       return next_scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			List[Any]     , __lowerCamelCase					:			int     , __lowerCamelCase					:			float = -float("Inf"			)     , __lowerCamelCase					:			int = 1			):
       if not isinstance(__lowerCamelCase     , __lowerCamelCase			) or top_k <= 0:
           raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}"			)

       SCREAMING_SNAKE_CASE										=   max(__lowerCamelCase     , __lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   filter_value





   def __call__(       self					:			Union[str, Any]     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   scores.shape
       SCREAMING_SNAKE_CASE										=   jnp.full(batch_size * vocab_size     , self.filter_value			)

       SCREAMING_SNAKE_CASE										=   min(self.top_k     , scores.shape[-1]			)  # Safety check
       SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   lax.top_k(__lowerCamelCase     , __lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   jnp.broadcast_to((jnp.arange(__lowerCamelCase			) * vocab_size)[:, None]     , (batch_size, topk)			).flatten()
       SCREAMING_SNAKE_CASE										=   topk_scores.flatten()
       SCREAMING_SNAKE_CASE										=   topk_indices.flatten() + shift

       SCREAMING_SNAKE_CASE										=   next_scores_flat.at[topk_indices_flat].set(__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   next_scores_flat.reshape(__lowerCamelCase     , __lowerCamelCase			)
       return next_scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			Optional[Any]     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE										=   bos_token_id





   def __call__(       self					:			str     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE										=   jnp.full(scores.shape     , -float("inf"			)			)

       SCREAMING_SNAKE_CASE										=   1 - jnp.bool_(cur_len - 1			)

       SCREAMING_SNAKE_CASE										=   jnp.where(__lowerCamelCase     , new_scores.at[:, self.bos_token_id].set(0			)     , __lowerCamelCase			)

       return scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			int     , __lowerCamelCase					:			int     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE										=   max_length
       SCREAMING_SNAKE_CASE										=   eos_token_id





   def __call__(       self					:			List[str]     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE										=   jnp.full(scores.shape     , -float("inf"			)			)

       SCREAMING_SNAKE_CASE										=   1 - jnp.bool_(cur_len - self.max_length + 1			)

       SCREAMING_SNAKE_CASE										=   jnp.where(__lowerCamelCase     , new_scores.at[:, self.eos_token_id].set(0			)     , __lowerCamelCase			)

       return scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			Union[str, Any]     , __lowerCamelCase					:			int     , __lowerCamelCase					:			int			):
       if not isinstance(__lowerCamelCase     , __lowerCamelCase			) or min_length < 0:
           raise ValueError(f"`min_length` has to be a positive integer, but is {min_length}"			)

       if not isinstance(__lowerCamelCase     , __lowerCamelCase			) or eos_token_id < 0:
           raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}"			)

       SCREAMING_SNAKE_CASE										=   min_length
       SCREAMING_SNAKE_CASE										=   eos_token_id





   def __call__(       self					:			List[str]     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int			):
       # create boolean flag to decide if min length penalty should be applied
       SCREAMING_SNAKE_CASE										=   1 - jnp.clip(cur_len - self.min_length     , 0     , 1			)

       SCREAMING_SNAKE_CASE										=   jnp.where(__lowerCamelCase     , scores.at[:, self.eos_token_id].set(-float("inf"			)			)     , __lowerCamelCase			)

       return scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			str     , __lowerCamelCase					:			List[str]     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE										=   list(__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   begin_index





   def __call__(       self					:			Any     , __lowerCamelCase					:			int     , __lowerCamelCase					:			List[Any]     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE										=   1 - jnp.bool_(cur_len - self.begin_index			)

       SCREAMING_SNAKE_CASE										=   jnp.where(__lowerCamelCase     , scores.at[:, self.begin_suppress_tokens].set(-float("inf"			)			)     , __lowerCamelCase			)

       return scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			Optional[Any]     , __lowerCamelCase					:			list			):
       SCREAMING_SNAKE_CASE										=   list(__lowerCamelCase			)





   def __call__(       self					:			List[str]     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int			):
       SCREAMING_SNAKE_CASE										=   scores.at[..., self.suppress_tokens].set(-float("inf"			)			)

       return scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			int     , __lowerCamelCase					:			Union[str, Any]			):
       SCREAMING_SNAKE_CASE										=   dict(__lowerCamelCase			)
       # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
       # index of the array corresponds to the index of the token to be forced, for XLA compatibility.
       # Indexes without forced tokens will have a negative value.
       SCREAMING_SNAKE_CASE										=   jnp.ones((max(force_token_map.keys()			) + 1)     , dtype=jnp.intaa			) * -1
       for index, token in force_token_map.items():
           if token is not None:
               SCREAMING_SNAKE_CASE										=   force_token_array.at[index].set(__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   jnp.intaa(__lowerCamelCase			)





   def __call__(       self					:			str     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			jnp.ndarray     , __lowerCamelCase					:			int			):
       def _force_token(__lowerCamelCase					:			str			):
           SCREAMING_SNAKE_CASE										=   scores.shape[0]
           SCREAMING_SNAKE_CASE										=   self.force_token_array[generation_idx]

           SCREAMING_SNAKE_CASE										=   jnp.ones_like(__lowerCamelCase     , dtype=scores.dtype			) * -float("inf"			)
           SCREAMING_SNAKE_CASE										=   jnp.zeros((batch_size, 1)     , dtype=scores.dtype			)
           SCREAMING_SNAKE_CASE										=   lax.dynamic_update_slice(__lowerCamelCase     , __lowerCamelCase     , (0, current_token)			)
           return new_scores

       SCREAMING_SNAKE_CASE										=   lax.cond(
           cur_len >= self.force_token_array.shape[0]     , lambda: scores     , lambda: lax.cond(
               self.force_token_array[cur_len] >= 0     , lambda: _force_token(__lowerCamelCase			)     , lambda: scores     , )     , )
       return scores





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			Union[str, Any]     , __lowerCamelCase					:			List[str]     , __lowerCamelCase					:			Tuple     , __lowerCamelCase					:			List[Any]			):
       SCREAMING_SNAKE_CASE										=   generate_config.eos_token_id
       SCREAMING_SNAKE_CASE										=   generate_config.no_timestamps_token_id
       SCREAMING_SNAKE_CASE										=   generate_config.no_timestamps_token_id + 1

       SCREAMING_SNAKE_CASE										=   decoder_input_length + 1

       if generate_config.is_multilingual:
           # room for language token and task token
           self.begin_index += 2
       if hasattr(__lowerCamelCase     , "max_initial_timestamp_index"			):
           SCREAMING_SNAKE_CASE										=   generate_config.max_initial_timestamp_index
       else:
           SCREAMING_SNAKE_CASE										=   model_config.vocab_size
       if self.max_initial_timestamp_index is None:
           SCREAMING_SNAKE_CASE										=   model_config.vocab_size





   def __call__(       self					:			Union[str, Any]     , __lowerCamelCase					:			Union[str, Any]     , __lowerCamelCase					:			Dict     , __lowerCamelCase					:			int			):
       # suppress <|notimestamps|> which is handled by without_timestamps
       SCREAMING_SNAKE_CASE										=   scores.at[:, self.no_timestamps_token_id].set(-float("inf"			)			)

       def handle_pairs(__lowerCamelCase					:			Tuple     , __lowerCamelCase					:			List[str]			):
           SCREAMING_SNAKE_CASE										=   jnp.where((cur_len - self.begin_index) >= 1     , __lowerCamelCase     , __lowerCamelCase			)
           SCREAMING_SNAKE_CASE										=   jnp.where(
               input_ids_k[cur_len - 1] >= self.timestamp_begin     , True and last_was_timestamp     , __lowerCamelCase     , )

           SCREAMING_SNAKE_CASE										=   jnp.where((cur_len - self.begin_index) < 2     , __lowerCamelCase     , __lowerCamelCase			)
           SCREAMING_SNAKE_CASE										=   jnp.where(
               input_ids_k[cur_len - 2] >= self.timestamp_begin     , __lowerCamelCase     , __lowerCamelCase     , )

           return jnp.where(
               __lowerCamelCase     , jnp.where(
                   penultimate_was_timestamp > 0     , scores_k.at[self.timestamp_begin :].set(-float("inf"			)			)     , scores_k.at[: self.eos_token_id].set(-float("inf"			)			)     , )     , __lowerCamelCase     , )

       SCREAMING_SNAKE_CASE										=   jax.vmap(__lowerCamelCase			)(__lowerCamelCase     , __lowerCamelCase			)

       SCREAMING_SNAKE_CASE										=   jnp.where(cur_len == self.begin_index     , __lowerCamelCase     , __lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   jnp.where(
           self.max_initial_timestamp_index is not None     , True and apply_max_initial_timestamp     , __lowerCamelCase     , )

       SCREAMING_SNAKE_CASE										=   self.timestamp_begin + self.max_initial_timestamp_index

       SCREAMING_SNAKE_CASE										=   jnp.where(
           __lowerCamelCase     , scores.at[:, last_allowed + 1 :].set(-float("inf"			)			)     , __lowerCamelCase     , )

       # if sum of probability over timestamps is above any other token, sample timestamp
       SCREAMING_SNAKE_CASE										=   jax.nn.log_softmax(__lowerCamelCase     , axis=-1			)

       def handle_cumulative_probs(__lowerCamelCase					:			Tuple     , __lowerCamelCase					:			List[str]			):
           SCREAMING_SNAKE_CASE										=   jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]     , axis=-1			)
           SCREAMING_SNAKE_CASE										=   jnp.max(logprobs_k[: self.timestamp_begin]			)
           return jnp.where(
               timestamp_logprob > max_text_token_logprob     , scores_k.at[: self.timestamp_begin].set(-float("inf"			)			)     , __lowerCamelCase     , )

       SCREAMING_SNAKE_CASE										=   jax.vmap(__lowerCamelCase			)(__lowerCamelCase     , __lowerCamelCase			)

       return scores




'''simple docstring'''



import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
							import tensorflow as tf
							from tensorflow.python.eager import context
							from tensorflow.python.framework import ops

							from transformers import GradientAccumulator, create_optimizer
@require_tf
class 						lowerCamelCase_				(       unittest.TestCase      ):







						def        __lowercase     (	self :   List[Any]					,     lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   Dict					,     lowerCAmelCase__ :   int					):

													"""simple docstring"""
													self.assertEqual(len(lowerCAmelCase__					)					,     len(lowerCAmelCase__					)					)
													for a, b in zip(lowerCAmelCase__					,     lowerCAmelCase__					):
																				self.assertAlmostEqual(lowerCAmelCase__					,     lowerCAmelCase__					,     delta=lowerCAmelCase__					)







						def        __lowercase     (	self :   Dict					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[Any]         =      GradientAccumulator()
													accumulator([tf.constant([1.0, 2.0]					)]					)
													accumulator([tf.constant([-2.0, 1.0]					)]					)
													accumulator([tf.constant([-1.0, 2.0]					)]					)
													with self.assertRaises(lowerCAmelCase__					):
																				accumulator([tf.constant([1.0, 1.0]					), tf.constant([2.0, 2.0]					)]					)
													self.assertEqual(accumulator.step					,     3					)
													self.assertEqual(len(accumulator.gradients					)					,     1					)
													self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()					,     [-2.0, 5.0]					,     tol=1e-2					)
													accumulator.reset()
													self.assertEqual(accumulator.step					,     0					)
													self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()					,     [0.0, 0.0]					,     tol=1e-2					)







						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      None
													ops.enable_eager_execution_internal()
													SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.config.list_physical_devices('''CPU'''					)
													if len(lowerCAmelCase__					) == 1:
																				tf.config.set_logical_device_configuration(
																				    physical_devices[0]					,     [tf.config.LogicalDeviceConfiguration(), tf.config.LogicalDeviceConfiguration()]					)
													SCREAMING_SNAKE_CASE       :   Tuple         =      tf.config.list_logical_devices(device_type='''CPU'''					)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.distribute.MirroredStrategy(devices=devices[:2]					)

													with strategy.scope():
																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      GradientAccumulator()
																				SCREAMING_SNAKE_CASE       :   Dict         =      tf.Variable([4.0, 3.0]					)
																				SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   Optional[int]         =      create_optimizer(5e-5					,     10					,     5					)
																				SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.Variable([0.0, 0.0]					,     trainable=lowerCAmelCase__					)

													def accumulate_on_replica(lowerCAmelCase__ :   List[str]					):
																				accumulator([gradient]					)

													def apply_on_replica():
																				optimizer.apply_gradients(list(zip(accumulator.gradients					,     [variable]					)					)					)

													@tf.function
													def accumulate(lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   List[Any]					):
																				with strategy.scope():
																											SCREAMING_SNAKE_CASE       :   Optional[Any]         =      strategy.experimental_local_results(lowerCAmelCase__					)
																											local_variables[0].assign(lowerCAmelCase__					)
																											local_variables[1].assign(lowerCAmelCase__					)
																											strategy.run(lowerCAmelCase__					,     args=(gradient_placeholder,)					)

													@tf.function
													def apply_grad():
																				with strategy.scope():
																											strategy.run(lowerCAmelCase__					)

													def _check_local_values(lowerCAmelCase__ :   Dict					,     lowerCAmelCase__ :   int					):
																				SCREAMING_SNAKE_CASE       :   Tuple         =      strategy.experimental_local_results(accumulator._gradients[0]					)
																				self.assertListAlmostEqual(values[0].value()					,     lowerCAmelCase__					,     tol=1e-2					)
																				self.assertListAlmostEqual(values[1].value()					,     lowerCAmelCase__					,     tol=1e-2					)

													accumulate([1.0, 2.0]					,     [-1.0, 1.0]					)
													accumulate([3.0, -1.0]					,     [-1.0, -1.0]					)
													accumulate([-2.0, 2.0]					,     [3.0, -2.0]					)
													self.assertEqual(accumulator.step					,     3					)
													_check_local_values([2.0, 3.0]					,     [1.0, -2.0]					)
													apply_grad()
													self.assertListAlmostEqual(variable.value()					,     [4.0, 3.0]					,     tol=1e-2					)
													accumulator.reset()
													self.assertEqual(accumulator.step					,     0					)
													_check_local_values([0.0, 0.0]					,     [0.0, 0.0]					)




