






"""simple docstring"""
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


lowerCamelCase_							      = get_logger(__name__)


lowerCamelCase_							      = 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      UpperCamelCase_			:





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







class      UpperCamelCase_			:





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







class      UpperCamelCase_			(a_	):





	@add_start_docstrings(UpperCamelCase_		)
	def __call__( self   :					Any      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int      ,      **lowerCAmelCase_   :					str		)		->     Any:

				for processor in self:
							UpperCAmelCase_	:					Tuple     	=   inspect.signature(processor.__call__		).parameters
							if len(UpperCamelCase_		) > 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."""		)
										UpperCAmelCase_	:					Tuple     	=   processor(UpperCamelCase_      ,      UpperCamelCase_      ,      UpperCamelCase_      ,      **UpperCamelCase_		)
							else:
										UpperCAmelCase_	:					Dict     	=   processor(UpperCamelCase_      ,      UpperCamelCase_      ,      UpperCamelCase_		)
				return scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					Any      ,      lowerCAmelCase_   :					float		)		->     List[str]:

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

				UpperCAmelCase_	:					Optional[Any]     	=   temperature





	def __call__( self   :					Any      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int		)		->     str:
				UpperCAmelCase_	:					Any     	=   scores / self.temperature
				return scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					List[str]      ,      lowerCAmelCase_   :					float      ,      lowerCAmelCase_   :					float = -float("Inf"		)      ,      lowerCAmelCase_   :					int = 1		)		->     Tuple:

				if not isinstance(UpperCamelCase_      ,      UpperCamelCase_		) 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(UpperCamelCase_      ,      UpperCamelCase_		) 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}"""		)

				UpperCAmelCase_	:					Any     	=   top_p
				UpperCAmelCase_	:					Union[str, Any]     	=   filter_value
				UpperCAmelCase_	:					Dict     	=   min_tokens_to_keep





	def __call__( self   :					Dict      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int		)		->     Optional[Any]:
				UpperCAmelCase_	,      UpperCAmelCase_	:					Dict     	=   lax.top_k(UpperCamelCase_      ,      scores.shape[-1]		)

				UpperCAmelCase_	:					Dict     	=   jnp.full_like(UpperCamelCase_      ,      self.filter_value		)
				UpperCAmelCase_	:					Union[str, Any]     	=   jax.nn.softmax(UpperCamelCase_      ,      axis=-1		).cumsum(axis=-1		)
				UpperCAmelCase_	:					Optional[Any]     	=   cumulative_probs < self.top_p

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

				# min tokens to keep
				UpperCAmelCase_	:					List[Any]     	=   score_mask.at[:, : self.min_tokens_to_keep].set(UpperCamelCase_		)

				UpperCAmelCase_	:					Optional[Any]     	=   jnp.where(UpperCamelCase_      ,      UpperCamelCase_      ,      UpperCamelCase_		)
				UpperCAmelCase_	:					Dict     	=   jax.lax.sort_key_val(UpperCamelCase_      ,      UpperCamelCase_		)[-1]

				return next_scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					int      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					float = -float("Inf"		)      ,      lowerCAmelCase_   :					int = 1		)		->     Any:

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

				UpperCAmelCase_	:					Any     	=   max(UpperCamelCase_      ,      UpperCamelCase_		)
				UpperCAmelCase_	:					List[str]     	=   filter_value





	def __call__( self   :					List[str]      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int		)		->     Union[str, Any]:
				UpperCAmelCase_	,      UpperCAmelCase_	:					Union[str, Any]     	=   scores.shape
				UpperCAmelCase_	:					Union[str, Any]     	=   jnp.full(batch_size * vocab_size      ,      self.filter_value		)

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

				UpperCAmelCase_	:					int     	=   next_scores_flat.at[topk_indices_flat].set(UpperCamelCase_		)
				UpperCAmelCase_	:					int     	=   next_scores_flat.reshape(UpperCamelCase_      ,      UpperCamelCase_		)
				return next_scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					List[str]      ,      lowerCAmelCase_   :					int		)		->     Tuple:
				UpperCAmelCase_	:					Dict     	=   bos_token_id





	def __call__( self   :					Optional[Any]      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int		)		->     Tuple:
				UpperCAmelCase_	:					Optional[int]     	=   jnp.full(scores.shape      ,      -float("inf"		)		)

				UpperCAmelCase_	:					str     	=   1 - jnp.bool_(cur_len - 1		)

				UpperCAmelCase_	:					Dict     	=   jnp.where(UpperCamelCase_      ,      new_scores.at[:, self.bos_token_id].set(0		)      ,      UpperCamelCase_		)

				return scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					int      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					int		)		->     Tuple:
				UpperCAmelCase_	:					Tuple     	=   max_length
				UpperCAmelCase_	:					str     	=   eos_token_id





	def __call__( self   :					Optional[Any]      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int		)		->     Union[str, Any]:
				UpperCAmelCase_	:					Union[str, Any]     	=   jnp.full(scores.shape      ,      -float("inf"		)		)

				UpperCAmelCase_	:					int     	=   1 - jnp.bool_(cur_len - self.max_length + 1		)

				UpperCAmelCase_	:					Optional[Any]     	=   jnp.where(UpperCamelCase_      ,      new_scores.at[:, self.eos_token_id].set(0		)      ,      UpperCamelCase_		)

				return scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					List[str]      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					int		)		->     List[str]:

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

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

				UpperCAmelCase_	:					List[str]     	=   min_length
				UpperCAmelCase_	:					List[str]     	=   eos_token_id





	def __call__( self   :					Tuple      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int		)		->     str:
				UpperCAmelCase_	:					int     	=   1 - jnp.clip(cur_len - self.min_length      ,      0      ,      1		)

				UpperCAmelCase_	:					Optional[int]     	=   jnp.where(UpperCamelCase_      ,      scores.at[:, self.eos_token_id].set(-float("inf"		)		)      ,      UpperCamelCase_		)

				return scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					List[Any]      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					int		)		->     Any:
				UpperCAmelCase_	:					List[Any]     	=   list(UpperCamelCase_		)
				UpperCAmelCase_	:					List[str]     	=   begin_index





	def __call__( self   :					Tuple      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					int		)		->     Optional[Any]:
				UpperCAmelCase_	:					Dict     	=   1 - jnp.bool_(cur_len - self.begin_index		)

				UpperCAmelCase_	:					Optional[int]     	=   jnp.where(UpperCamelCase_      ,      scores.at[:, self.begin_suppress_tokens].set(-float("inf"		)		)      ,      UpperCamelCase_		)

				return scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					Union[str, Any]      ,      lowerCAmelCase_   :					list		)		->     List[str]:
				UpperCAmelCase_	:					str     	=   list(UpperCamelCase_		)





	def __call__( self   :					Tuple      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int		)		->     Union[str, Any]:
				UpperCAmelCase_	:					List[Any]     	=   scores.at[..., self.suppress_tokens].set(-float("inf"		)		)

				return scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					Any      ,      lowerCAmelCase_   :					Tuple		)		->     Optional[Any]:
				UpperCAmelCase_	:					int     	=   dict(UpperCamelCase_		)
				# 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.
				UpperCAmelCase_	:					Optional[Any]     	=   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:
										UpperCAmelCase_	:					Union[str, Any]     	=   force_token_array.at[index].set(UpperCamelCase_		)
				UpperCAmelCase_	:					Any     	=   jnp.intaa(UpperCamelCase_		)





	def __call__( self   :					List[str]      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					jnp.ndarray      ,      lowerCAmelCase_   :					int		)		->     str:

				def _force_token(lowerCAmelCase_   :					Optional[Any]		):
							UpperCAmelCase_	:					Union[str, Any]     	=   scores.shape[0]
							UpperCAmelCase_	:					List[str]     	=   self.force_token_array[generation_idx]

							UpperCAmelCase_	:					Tuple     	=   jnp.ones_like(UpperCamelCase_      ,      dtype=scores.dtype		) * -float("inf"		)
							UpperCAmelCase_	:					List[Any]     	=   jnp.zeros((batch_size, 1)      ,      dtype=scores.dtype		)
							UpperCAmelCase_	:					Any     	=   lax.dynamic_update_slice(UpperCamelCase_      ,      UpperCamelCase_      ,      (0, current_token)		)
							return new_scores

				UpperCAmelCase_	:					Tuple     	=   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(UpperCamelCase_		)      ,      lambda: scores      ,      )      ,      )
				return scores







class      UpperCamelCase_			(a_	):





	def __init__( self   :					List[Any]      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					Dict		)		->     Any:
				UpperCAmelCase_	:					List[str]     	=   generate_config.eos_token_id
				UpperCAmelCase_	:					int     	=   generate_config.no_timestamps_token_id
				UpperCAmelCase_	:					Optional[int]     	=   generate_config.no_timestamps_token_id + 1

				UpperCAmelCase_	:					List[Any]     	=   decoder_input_length + 1

				if generate_config.is_multilingual:
							# room for language token and task token
							self.begin_index += 2
				if hasattr(UpperCamelCase_      ,      "max_initial_timestamp_index"		):
							UpperCAmelCase_	:					List[Any]     	=   generate_config.max_initial_timestamp_index
				else:
							UpperCAmelCase_	:					Optional[Any]     	=   model_config.vocab_size
				if self.max_initial_timestamp_index is None:
							UpperCAmelCase_	:					Tuple     	=   model_config.vocab_size





	def __call__( self   :					str      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					Dict		)		->     Any:
				UpperCAmelCase_	:					Optional[int]     	=   scores.at[:, self.no_timestamps_token_id].set(-float("inf"		)		)

				def handle_pairs(lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					List[Any]		):
							UpperCAmelCase_	:					List[str]     	=   jnp.where((cur_len - self.begin_index) >= 1      ,      UpperCamelCase_      ,      UpperCamelCase_		)
							UpperCAmelCase_	:					List[Any]     	=   jnp.where(
							    input_ids_k[cur_len - 1] >= self.timestamp_begin      ,      True and last_was_timestamp      ,      UpperCamelCase_      ,      )

							UpperCAmelCase_	:					Dict     	=   jnp.where((cur_len - self.begin_index) < 2      ,      UpperCamelCase_      ,      UpperCamelCase_		)
							UpperCAmelCase_	:					Tuple     	=   jnp.where(
							    input_ids_k[cur_len - 2] >= self.timestamp_begin      ,      UpperCamelCase_      ,      UpperCamelCase_      ,      )

							return jnp.where(
							    UpperCamelCase_      ,      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"		)		)      ,      )      ,      UpperCamelCase_      ,      )

				UpperCAmelCase_	:					Optional[int]     	=   jax.vmap(UpperCamelCase_		)(UpperCamelCase_      ,      UpperCamelCase_		)

				UpperCAmelCase_	:					Any     	=   jnp.where(cur_len == self.begin_index      ,      UpperCamelCase_      ,      UpperCamelCase_		)
				UpperCAmelCase_	:					List[str]     	=   jnp.where(
				    self.max_initial_timestamp_index is not None      ,      True and apply_max_initial_timestamp      ,      UpperCamelCase_      ,      )

				UpperCAmelCase_	:					int     	=   self.timestamp_begin + self.max_initial_timestamp_index

				UpperCAmelCase_	:					Optional[int]     	=   jnp.where(
				    UpperCamelCase_      ,      scores.at[:, last_allowed + 1 :].set(-float("inf"		)		)      ,      UpperCamelCase_      ,      )

				# if sum of probability over timestamps is above any other token, sample timestamp
				UpperCAmelCase_	:					Union[str, Any]     	=   jax.nn.log_softmax(UpperCamelCase_      ,      axis=-1		)

				def handle_cumulative_probs(lowerCAmelCase_   :					Dict      ,      lowerCAmelCase_   :					Optional[int]		):
							UpperCAmelCase_	:					List[Any]     	=   jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]      ,      axis=-1		)
							UpperCAmelCase_	:					Any     	=   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"		)		)      ,      UpperCamelCase_      ,      )

				UpperCAmelCase_	:					Any     	=   jax.vmap(UpperCamelCase_		)(UpperCamelCase_      ,      UpperCamelCase_		)

				return scores











"""simple docstring"""
def 							A						(				__snake_case: int  )    ->		int:

       """simple docstring"""





       if divisor % 5 == 0 or divisor % 2 == 0:
              return 0
       __magic_name__      				=    1
       __magic_name__      				=    1
       while repunit:
              __magic_name__      				=    (1_0 * repunit + 1) % divisor
              repunit_index += 1
       return repunit_index


def 							A						(				__snake_case: int = 1_0_0_0_0_0_0  )    ->		int:

       """simple docstring"""





       __magic_name__      				=    limit - 1
       if divisor % 2 == 0:
              divisor += 1
       while least_divisible_repunit(__snake_case  ) <= limit:
              divisor += 2
       return divisor


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