


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


__UpperCAmelCase       =							get_logger(__name__)


__UpperCAmelCase       =							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 		a__    :





							'''simple docstring'''






							@add_start_docstrings(lowerCamelCase_       )
							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														raise NotImplementedError(
														    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""       )




class 		a__    :





							'''simple docstring'''






							@add_start_docstrings(lowerCamelCase_       )
							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														raise NotImplementedError(
														    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""       )




class 		a__    (							a__  ):





							'''simple docstring'''






							@add_start_docstrings(lowerCamelCase_       )
							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   **lowerCamelCase_       )	->							jnp.ndarray:
														for processor in self:
																					lowerCAmelCase__	       =     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."""       )
																												lowerCAmelCase__	       =     processor(lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   **lowerCamelCase_       )
																					else:
																												lowerCAmelCase__	       =     processor(lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )
														return scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_       )	->							Union[str, Any]:
														if not isinstance(lowerCamelCase_ ,   lowerCamelCase_       ) or not (temperature > 0):
																					raise ValueError(F"""`temperature` has to be a strictly positive float, but is {temperature}"""       )

														lowerCAmelCase__	       =     temperature







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														lowerCAmelCase__	       =     scores / self.temperature
														return scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ = -float('''Inf'''       ) ,   lowerCamelCase_ = 1       )	->							Optional[int]:
														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}"""       )

														lowerCAmelCase__	       =     top_p
														lowerCAmelCase__	       =     filter_value
														lowerCAmelCase__	       =     min_tokens_to_keep







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														lowerCAmelCase__       , lowerCAmelCase__	       =     lax.top_k(lowerCamelCase_ ,   scores.shape[-1]       )

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

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

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

														lowerCAmelCase__	       =     jnp.where(lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )
														lowerCAmelCase__	       =     jax.lax.sort_key_val(lowerCamelCase_ ,   lowerCamelCase_       )[-1]

														return next_scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ = -float('''Inf'''       ) ,   lowerCamelCase_ = 1       )	->							str:
														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}"""       )

														lowerCAmelCase__	       =     max(lowerCamelCase_ ,   lowerCamelCase_       )
														lowerCAmelCase__	       =     filter_value







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														lowerCAmelCase__       , lowerCAmelCase__	       =     scores.shape
														lowerCAmelCase__	       =     jnp.full(batch_size * vocab_size ,   self.filter_value       )

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

														lowerCAmelCase__	       =     next_scores_flat.at[topk_indices_flat].set(lowerCamelCase_       )
														lowerCAmelCase__	       =     next_scores_flat.reshape(lowerCamelCase_ ,   lowerCamelCase_       )
														return next_scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_       )	->							List[str]:
														lowerCAmelCase__	       =     bos_token_id







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														lowerCAmelCase__	       =     jnp.full(scores.shape ,   -float('''inf'''       )       )

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

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

														return scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							List[str]:
														lowerCAmelCase__	       =     max_length
														lowerCAmelCase__	       =     eos_token_id







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														lowerCAmelCase__	       =     jnp.full(scores.shape ,   -float('''inf'''       )       )

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

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

														return scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							Optional[Any]:
														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}"""       )

														lowerCAmelCase__	       =     min_length
														lowerCAmelCase__	       =     eos_token_id







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														# create boolean flag to decide if min length penalty should be applied
														lowerCAmelCase__	       =     1 - jnp.clip(cur_len - self.min_length ,   0 ,   1       )

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

														return scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							List[Any]:
														lowerCAmelCase__	       =     list(lowerCamelCase_       )
														lowerCAmelCase__	       =     begin_index







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							Union[str, Any]:
														lowerCAmelCase__	       =     1 - jnp.bool_(cur_len - self.begin_index       )

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

														return scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_       )	->							Optional[Any]:
														lowerCAmelCase__	       =     list(lowerCamelCase_       )







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														lowerCAmelCase__	       =     scores.at[..., self.suppress_tokens].set(-float('''inf'''       )       )

														return scores




class 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_       )	->							List[Any]:
														lowerCAmelCase__	       =     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.
														lowerCAmelCase__	       =     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:
																												lowerCAmelCase__	       =     force_token_array.at[index].set(lowerCamelCase_       )
														lowerCAmelCase__	       =     jnp.intaa(lowerCamelCase_       )







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							jnp.ndarray:
														def _force_token(lowerCamelCase_       ):
																					lowerCAmelCase__	       =     scores.shape[0]
																					lowerCAmelCase__	       =     self.force_token_array[generation_idx]

																					lowerCAmelCase__	       =     jnp.ones_like(lowerCamelCase_ ,   dtype=scores.dtype       ) * -float('''inf'''       )
																					lowerCAmelCase__	       =     jnp.zeros((batch_size, 1) ,   dtype=scores.dtype       )
																					lowerCAmelCase__	       =     lax.dynamic_update_slice(lowerCamelCase_ ,   lowerCamelCase_ ,   (0, current_token)       )
																					return new_scores

														lowerCAmelCase__	       =     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 		a__    (							a__  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							int:
														lowerCAmelCase__	       =     generate_config.eos_token_id
														lowerCAmelCase__	       =     generate_config.no_timestamps_token_id
														lowerCAmelCase__	       =     generate_config.no_timestamps_token_id + 1

														lowerCAmelCase__	       =     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'''       ):
																					lowerCAmelCase__	       =     generate_config.max_initial_timestamp_index
														else:
																					lowerCAmelCase__	       =     model_config.vocab_size
														if self.max_initial_timestamp_index is None:
																					lowerCAmelCase__	       =     model_config.vocab_size







							def __call__(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							Tuple:
														# suppress <|notimestamps|> which is handled by without_timestamps
														lowerCAmelCase__	       =     scores.at[:, self.no_timestamps_token_id].set(-float('''inf'''       )       )

														def handle_pairs(lowerCamelCase_ ,   lowerCamelCase_       ):
																					lowerCAmelCase__	       =     jnp.where((cur_len - self.begin_index) >= 1 ,   lowerCamelCase_ ,   lowerCamelCase_       )
																					lowerCAmelCase__	       =     jnp.where(
																					    input_ids_k[cur_len - 1] >= self.timestamp_begin ,   True and last_was_timestamp ,   lowerCamelCase_ ,   )

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

														lowerCAmelCase__	       =     jax.vmap(lowerCamelCase_       )(lowerCamelCase_ ,   lowerCamelCase_       )

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

														lowerCAmelCase__	       =     self.timestamp_begin + self.max_initial_timestamp_index

														lowerCAmelCase__	       =     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
														lowerCAmelCase__	       =     jax.nn.log_softmax(lowerCamelCase_ ,   axis=-1       )

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

														lowerCAmelCase__	       =     jax.vmap(lowerCamelCase_       )(lowerCamelCase_ ,   lowerCamelCase_       )

														return scores



'''simple docstring'''






from __future__ import annotations

import numpy as np
from numpy import floataa
from numpy.typing import NDArray


def  _snake_case       (					A			,   A			,   A			,   A			,   )						->   list[float]:
							lowerCAmelCase__       , lowerCAmelCase__	       =     coefficient_matrix.shape
							lowerCAmelCase__       , lowerCAmelCase__	       =     constant_matrix.shape

							if rowsa != colsa:
														lowerCAmelCase__	       =     F"""Coefficient matrix dimensions must be nxn but received {rowsa}x{colsa}"""
														raise ValueError(A	)

							if colsa != 1:
														lowerCAmelCase__	       =     F"""Constant matrix must be nx1 but received {rowsa}x{colsa}"""
														raise ValueError(A	)

							if rowsa != rowsa:
														lowerCAmelCase__	       =     (
														    '''Coefficient and constant matrices dimensions must be nxn and nx1 but '''
														    F"""received {rowsa}x{colsa} and {rowsa}x{colsa}"""
														)
														raise ValueError(A	)

							if len(A	) != rowsa:
														lowerCAmelCase__	       =     (
														    '''Number of initial values must be equal to number of rows in coefficient '''
														    F"""matrix but received {len(A	)} and {rowsa}"""
														)
														raise ValueError(A	)

							if iterations <= 0:
														raise ValueError('''Iterations must be at least 1'''	)

							lowerCAmelCase__	       =     np.concatenate(
							    (coefficient_matrix, constant_matrix)			,   axis=1	)

							lowerCAmelCase__       , lowerCAmelCase__	       =     table.shape

							strictly_diagonally_dominant(A	)

							# Iterates the whole matrix for given number of times
							for _ in range(A	):
														lowerCAmelCase__	       =     []
														for row in range(A	):
																					lowerCAmelCase__	       =     0
																					for col in range(A	):
																												if col == row:
																																			lowerCAmelCase__	       =     table[row][col]
																												elif col == cols - 1:
																																			lowerCAmelCase__	       =     table[row][col]
																												else:
																																			temp += (-1) * table[row][col] * init_val[col]
																					lowerCAmelCase__	       =     (temp + val) / denom
																					new_val.append(A	)
														lowerCAmelCase__	       =     new_val

							return [float(A	) for i in new_val]


def  _snake_case       (					A	)						->   bool:
							lowerCAmelCase__       , lowerCAmelCase__	       =     table.shape

							lowerCAmelCase__	       =     True

							for i in range(0			,   A	):
														lowerCAmelCase__	       =     0
														for j in range(0			,   cols - 1	):
																					if i == j:
																												continue
																					else:
																												total += table[i][j]

														if table[i][i] <= total:
																					raise ValueError('''Coefficient matrix is not strictly diagonally dominant'''	)

							return is_diagonally_dominant


# Test Cases
if __name__ == "__main__":
		import doctest

		doctest.testmod()