





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


lowercase_		:				Tuple				=						get_logger(__name__)


lowercase_		:				Union[str, Any]				=						r'''
    Args:
        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
            search or log softmax for each vocabulary token when using beam search
        kwargs (`Dict[str, Any]`, *optional*):
            Additional logits processor specific kwargs.

    Return:
        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.

'''


class 			lowercase				:



						"""simple docstring"""

						@add_start_docstrings(lowerCamelCase_					)
						def __call__(    self			:		Any ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray					):





										'''simple docstring'''







										raise NotImplementedError(
										    f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''					)



class 			lowercase				:



						"""simple docstring"""

						@add_start_docstrings(lowerCamelCase_					)
						def __call__(    self			:		Dict ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray					):





										'''simple docstring'''







										raise NotImplementedError(
										    f'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''					)



class 			lowercase				(							a_     ):



						"""simple docstring"""

						@add_start_docstrings(lowerCamelCase_					)
						def __call__(    self			:		Tuple ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int ,				**lowerCamelCase_			:		Dict					):





										'''simple docstring'''







										for processor in self:
														_snake_case :			List[Any]  						=   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.'''					)
																		_snake_case :			Dict  						=   processor(lowerCamelCase_ ,				lowerCamelCase_ ,				lowerCamelCase_ ,				**lowerCamelCase_					)
														else:
																		_snake_case :			List[Any]  						=   processor(lowerCamelCase_ ,				lowerCamelCase_ ,				lowerCamelCase_					)
										return scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		int ,				lowerCamelCase_			:		float					):





										'''simple docstring'''







										if not isinstance(lowerCamelCase_ ,				lowerCamelCase_					) or not (temperature > 0):
														raise ValueError(f'''`temperature` has to be a strictly positive float, but is {temperature}'''					)

										_snake_case :			Tuple  						=   temperature







						def __call__(    self			:		Union[str, Any] ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case :			int  						=   scores / self.temperature
										return scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		Dict ,				lowerCamelCase_			:		float ,				lowerCamelCase_			:		float = -float('Inf'					) ,				lowerCamelCase_			:		int = 1					):





										'''simple docstring'''







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

										_snake_case :			str  						=   top_p
										_snake_case :			List[Any]  						=   filter_value
										_snake_case :			Optional[int]  						=   min_tokens_to_keep







						def __call__(    self			:		Dict ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case						,			_snake_case :			Dict  						=   lax.top_k(lowerCamelCase_ ,				scores.shape[-1]					)

										_snake_case :			List[Any]  						=   jnp.full_like(lowerCamelCase_ ,				self.filter_value					)
										_snake_case :			str  						=   jax.nn.softmax(lowerCamelCase_ ,				axis=-1					).cumsum(axis=-1					)
										_snake_case :			List[str]  						=   cumulative_probs < self.top_p

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

										# min tokens to keep
										_snake_case :			int  						=   score_mask.at[:, : self.min_tokens_to_keep].set(lowerCamelCase_					)

										_snake_case :			Optional[int]  						=   jnp.where(lowerCamelCase_ ,				lowerCamelCase_ ,				lowerCamelCase_					)
										_snake_case :			Optional[Any]  						=   jax.lax.sort_key_val(lowerCamelCase_ ,				lowerCamelCase_					)[-1]

										return next_scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		Any ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		float = -float('Inf'					) ,				lowerCamelCase_			:		int = 1					):





										'''simple docstring'''







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

										_snake_case :			List[Any]  						=   max(lowerCamelCase_ ,				lowerCamelCase_					)
										_snake_case :			Union[str, Any]  						=   filter_value







						def __call__(    self			:		str ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case						,			_snake_case :			str  						=   scores.shape
										_snake_case :			Any  						=   jnp.full(batch_size * vocab_size ,				self.filter_value					)

										_snake_case :			Dict  						=   min(self.top_k ,				scores.shape[-1]					)  # Safety check
										_snake_case						,			_snake_case :			Optional[Any]  						=   lax.top_k(lowerCamelCase_ ,				lowerCamelCase_					)
										_snake_case :			Optional[Any]  						=   jnp.broadcast_to((jnp.arange(lowerCamelCase_					) * vocab_size)[:, None] ,				(batch_size, topk)					).flatten()
										_snake_case :			List[str]  						=   topk_scores.flatten()
										_snake_case :			Any  						=   topk_indices.flatten() + shift

										_snake_case :			Tuple  						=   next_scores_flat.at[topk_indices_flat].set(lowerCamelCase_					)
										_snake_case :			List[Any]  						=   next_scores_flat.reshape(lowerCamelCase_ ,				lowerCamelCase_					)
										return next_scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		Union[str, Any] ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case :			List[Any]  						=   bos_token_id







						def __call__(    self			:		str ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case :			List[Any]  						=   jnp.full(scores.shape ,				-float('inf'					)					)

										_snake_case :			Optional[int]  						=   1 - jnp.bool_(cur_len - 1					)

										_snake_case :			Any  						=   jnp.where(lowerCamelCase_ ,				new_scores.at[:, self.bos_token_id].set(0					) ,				lowerCamelCase_					)

										return scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		Dict ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case :			Optional[int]  						=   max_length
										_snake_case :			Union[str, Any]  						=   eos_token_id







						def __call__(    self			:		Union[str, Any] ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case :			Any  						=   jnp.full(scores.shape ,				-float('inf'					)					)

										_snake_case :			str  						=   1 - jnp.bool_(cur_len - self.max_length + 1					)

										_snake_case :			str  						=   jnp.where(lowerCamelCase_ ,				new_scores.at[:, self.eos_token_id].set(0					) ,				lowerCamelCase_					)

										return scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		str ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







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

										_snake_case :			str  						=   min_length
										_snake_case :			Optional[Any]  						=   eos_token_id







						def __call__(    self			:		Union[str, Any] ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case :			List[Any]  						=   1 - jnp.clip(cur_len - self.min_length ,				0 ,				1					)

										_snake_case :			Dict  						=   jnp.where(lowerCamelCase_ ,				scores.at[:, self.eos_token_id].set(-float('inf'					)					) ,				lowerCamelCase_					)

										return scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		Optional[Any] ,				lowerCamelCase_			:		List[Any] ,				lowerCamelCase_			:		Optional[Any]					):





										'''simple docstring'''







										_snake_case :			Optional[int]  						=   list(lowerCamelCase_					)
										_snake_case :			int  						=   begin_index







						def __call__(    self			:		Tuple ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case :			Any  						=   1 - jnp.bool_(cur_len - self.begin_index					)

										_snake_case :			Union[str, Any]  						=   jnp.where(lowerCamelCase_ ,				scores.at[:, self.begin_suppress_tokens].set(-float('inf'					)					) ,				lowerCamelCase_					)

										return scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		Optional[Any] ,				lowerCamelCase_			:		list					):





										'''simple docstring'''







										_snake_case :			Optional[Any]  						=   list(lowerCamelCase_					)







						def __call__(    self			:		int ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										_snake_case :			Optional[int]  						=   scores.at[..., self.suppress_tokens].set(-float('inf'					)					)

										return scores



class 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		int ,				lowerCamelCase_			:		Optional[int]					):





										'''simple docstring'''







										_snake_case :			List[Any]  						=   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.
										_snake_case :			Dict  						=   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:
																		_snake_case :			Any  						=   force_token_array.at[index].set(lowerCamelCase_					)
										_snake_case :			Union[str, Any]  						=   jnp.intaa(lowerCamelCase_					)







						def __call__(    self			:		int ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		jnp.ndarray ,				lowerCamelCase_			:		int					):





										'''simple docstring'''







										def _force_token(lowerCamelCase_			:		Any					):
														_snake_case :			int  						=   scores.shape[0]
														_snake_case :			List[str]  						=   self.force_token_array[generation_idx]

														_snake_case :			List[Any]  						=   jnp.ones_like(lowerCamelCase_ ,				dtype=scores.dtype					) * -float('inf'					)
														_snake_case :			int  						=   jnp.zeros((batch_size, 1) ,				dtype=scores.dtype					)
														_snake_case :			List[Any]  						=   lax.dynamic_update_slice(lowerCamelCase_ ,				lowerCamelCase_ ,				(0, current_token)					)
														return new_scores

										_snake_case :			Optional[Any]  						=   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 			lowercase				(							a_     ):



						"""simple docstring"""

						def __init__(    self			:		Any ,				lowerCamelCase_			:		List[Any] ,				lowerCamelCase_			:		Optional[int] ,				lowerCamelCase_			:		Optional[Any]					):





										'''simple docstring'''







										_snake_case :			Dict  						=   generate_config.eos_token_id
										_snake_case :			Any  						=   generate_config.no_timestamps_token_id
										_snake_case :			Optional[Any]  						=   generate_config.no_timestamps_token_id + 1

										_snake_case :			Dict  						=   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'					):
														_snake_case :			str  						=   generate_config.max_initial_timestamp_index
										else:
														_snake_case :			Union[str, Any]  						=   model_config.vocab_size
										if self.max_initial_timestamp_index is None:
														_snake_case :			List[Any]  						=   model_config.vocab_size







						def __call__(    self			:		List[str] ,				lowerCamelCase_			:		List[Any] ,				lowerCamelCase_			:		Dict ,				lowerCamelCase_			:		Dict					):





										'''simple docstring'''







										_snake_case :			List[str]  						=   scores.at[:, self.no_timestamps_token_id].set(-float('inf'					)					)

										def handle_pairs(lowerCamelCase_			:		List[Any] ,				lowerCamelCase_			:		List[Any]					):
														_snake_case :			Optional[int]  						=   jnp.where((cur_len - self.begin_index) >= 1 ,				lowerCamelCase_ ,				lowerCamelCase_					)
														_snake_case :			Tuple  						=   jnp.where(
														    input_ids_k[cur_len - 1] >= self.timestamp_begin ,				True and last_was_timestamp ,				lowerCamelCase_ ,				)

														_snake_case :			Any  						=   jnp.where((cur_len - self.begin_index) < 2 ,				lowerCamelCase_ ,				lowerCamelCase_					)
														_snake_case :			Tuple  						=   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_ ,				)

										_snake_case :			Optional[Any]  						=   jax.vmap(lowerCamelCase_					)(lowerCamelCase_ ,				lowerCamelCase_					)

										_snake_case :			List[str]  						=   jnp.where(cur_len == self.begin_index ,				lowerCamelCase_ ,				lowerCamelCase_					)
										_snake_case :			Optional[int]  						=   jnp.where(
										    self.max_initial_timestamp_index is not None ,				True and apply_max_initial_timestamp ,				lowerCamelCase_ ,				)

										_snake_case :			Any  						=   self.timestamp_begin + self.max_initial_timestamp_index

										_snake_case :			str  						=   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
										_snake_case :			List[Any]  						=   jax.nn.log_softmax(lowerCamelCase_ ,				axis=-1					)

										def handle_cumulative_probs(lowerCamelCase_			:		Union[str, Any] ,				lowerCamelCase_			:		List[Any]					):
														_snake_case :			List[str]  						=   jax.nn.logsumexp(logprobs_k[self.timestamp_begin :] ,				axis=-1					)
														_snake_case :			Tuple  						=   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_ ,				)

										_snake_case :			Optional[int]  						=   jax.vmap(lowerCamelCase_					)(lowerCamelCase_ ,				lowerCamelCase_					)

										return scores













import functools

def       A__(		__lowerCAmelCase	,			__lowerCAmelCase ):

				# Validation
				if not isinstance(__lowerCAmelCase	,			__lowerCAmelCase ) or not all(isinstance(__lowerCAmelCase	,			__lowerCAmelCase ) for day in days ):
								raise ValueError('The parameter days should be a list of integers' )

				if len(__lowerCAmelCase ) != 3 or not all(isinstance(__lowerCAmelCase	,			__lowerCAmelCase ) for cost in costs ):
								raise ValueError('The parameter costs should be a list of three integers' )

				if len(__lowerCAmelCase ) == 0:
								return 0

				if min(__lowerCAmelCase ) <= 0:
								raise ValueError('All days elements should be greater than 0' )

				if max(__lowerCAmelCase ) >= 3_66:
								raise ValueError('All days elements should be less than 366' )

				_snake_case :			Optional[int]  						=   set(__lowerCAmelCase )

				@functools.cache
				def dynamic_programming(__lowerCAmelCase ) -> int:
								if index > 3_65:
												return 0

								if index not in days_set:
												return dynamic_programming(index + 1 )

								return min(
								    costs[0] + dynamic_programming(index + 1 )	,			costs[1] + dynamic_programming(index + 7 )	,			costs[2] + dynamic_programming(index + 30 )	,			)

				return dynamic_programming(1 )


if __name__ == "__main__":
		import doctest

		doctest.testmod()






