"""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"\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 			UpperCamelCase__:





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






class 			UpperCamelCase__(	__A       ):





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

													A__		   =							temperature





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






class 			UpperCamelCase__(	__A       ):





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase = -float('Inf'   )						,__UpperCAmelCase = 1   )    -> List[Any]:
													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}'''   )

													A__		   =							top_p
													A__		   =							filter_value
													A__		   =							min_tokens_to_keep





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

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

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

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

													A__		   =							jnp.where(__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )
													A__		   =							jax.lax.sort_key_val(__UpperCAmelCase						,__UpperCAmelCase   )[-1]

													return next_scores






class 			UpperCamelCase__(	__A       ):





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

													A__		   =							max(__UpperCAmelCase						,__UpperCAmelCase   )
													A__		   =							filter_value





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

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

													A__		   =							next_scores_flat.at[topk_indices_flat].set(__UpperCAmelCase   )
													A__		   =							next_scores_flat.reshape(__UpperCAmelCase						,__UpperCAmelCase   )
													return next_scores






class 			UpperCamelCase__(	__A       ):





						def __init__(				self						,__UpperCAmelCase   )    -> Optional[int]:
													A__		   =							bos_token_id





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

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

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

													return scores






class 			UpperCamelCase__(	__A       ):





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase   )    -> Optional[int]:
													A__		   =							max_length
													A__		   =							eos_token_id





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

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

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

													return scores






class 			UpperCamelCase__(	__A       ):





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase   )    -> int:
													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}'''   )

													A__		   =							min_length
													A__		   =							eos_token_id





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

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

													return scores






class 			UpperCamelCase__(	__A       ):





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase   )    -> Union[str, Any]:
													A__		   =							list(__UpperCAmelCase   )
													A__		   =							begin_index





						def __call__(				self						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )    -> str:
													A__		   =							1 - jnp.bool_(cur_len - self.begin_index   )

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

													return scores






class 			UpperCamelCase__(	__A       ):





						def __init__(				self						,__UpperCAmelCase   )    -> Any:
													A__		   =							list(__UpperCAmelCase   )





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

													return scores






class 			UpperCamelCase__(	__A       ):





						def __init__(				self						,__UpperCAmelCase   )    -> Optional[int]:
													A__		   =							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.
													A__		   =							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:
																											A__		   =							force_token_array.at[index].set(__UpperCAmelCase   )
													A__		   =							jnp.intaa(__UpperCAmelCase   )





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

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

													A__		   =							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						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )    -> Optional[Any]:
													A__		   =							generate_config.eos_token_id
													A__		   =							generate_config.no_timestamps_token_id
													A__		   =							generate_config.no_timestamps_token_id + 1

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





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

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

																				A__		   =							jnp.where((cur_len - self.begin_index) < 2						,__UpperCAmelCase						,__UpperCAmelCase   )
																				A__		   =							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						,)

													A__		   =							jax.vmap(__UpperCAmelCase   )(__UpperCAmelCase						,__UpperCAmelCase   )

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

													A__		   =							self.timestamp_begin + self.max_initial_timestamp_index

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

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

													A__		   =							jax.vmap(__UpperCAmelCase   )(__UpperCAmelCase						,__UpperCAmelCase   )

													return scores






"""simple docstring"""




from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase							 =     logging.get_logger(__name__)

__lowerCamelCase							 =     {
    "huggingface/time-series-transformer-tourism-monthly": (
        "https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}
class 			UpperCamelCase__(	__A       ):
						lowerCAmelCase__     :  List[Any]  							=						'time_series_transformer'
						lowerCAmelCase__     :  Dict  							=						{
						    'hidden_size': 'd_model',
						    'num_attention_heads': 'encoder_attention_heads',
						    'num_hidden_layers': 'encoder_layers',
						}





						def __init__(				self						,__UpperCAmelCase = None						,__UpperCAmelCase = None						,__UpperCAmelCase = "student_t"						,__UpperCAmelCase = "nll"						,__UpperCAmelCase = 1						,__UpperCAmelCase = [1, 2, 3, 4, 5, 6, 7]						,__UpperCAmelCase = "mean"						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = None						,__UpperCAmelCase = None						,__UpperCAmelCase = 32						,__UpperCAmelCase = 32						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = True						,__UpperCAmelCase = "gelu"						,__UpperCAmelCase = 64						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 1_00						,__UpperCAmelCase = 0.0_2						,__UpperCAmelCase=True						,**__UpperCAmelCase						,)    -> List[str]:
													# time series specific configuration
													A__		   =							prediction_length
													A__		   =							context_length or prediction_length
													A__		   =							distribution_output
													A__		   =							loss
													A__		   =							input_size
													A__		   =							num_time_features
													A__		   =							lags_sequence
													A__		   =							scaling
													A__		   =							num_dynamic_real_features
													A__		   =							num_static_real_features
													A__		   =							num_static_categorical_features
													if cardinality and num_static_categorical_features > 0:
																				if len(__UpperCAmelCase   ) != num_static_categorical_features:
																											raise ValueError(
																											    'The cardinality should be a list of the same length as `num_static_categorical_features`'   )
																				A__		   =							cardinality
													else:
																				A__		   =							[0]
													if embedding_dimension and num_static_categorical_features > 0:
																				if len(__UpperCAmelCase   ) != num_static_categorical_features:
																											raise ValueError(
																											    'The embedding dimension should be a list of the same length as `num_static_categorical_features`'   )
																				A__		   =							embedding_dimension
													else:
																				A__		   =							[min(50						,(cat + 1) // 2   ) for cat in self.cardinality]
													A__		   =							num_parallel_samples

													# Transformer architecture configuration
													A__		   =							input_size * len(__UpperCAmelCase   ) + self._number_of_features
													A__		   =							d_model
													A__		   =							encoder_attention_heads
													A__		   =							decoder_attention_heads
													A__		   =							encoder_ffn_dim
													A__		   =							decoder_ffn_dim
													A__		   =							encoder_layers
													A__		   =							decoder_layers

													A__		   =							dropout
													A__		   =							attention_dropout
													A__		   =							activation_dropout
													A__		   =							encoder_layerdrop
													A__		   =							decoder_layerdrop

													A__		   =							activation_function
													A__		   =							init_std

													A__		   =							use_cache

													super().__init__(is_encoder_decoder=__UpperCAmelCase						,**__UpperCAmelCase   )





						@property
						def  snake_case__		(				self   )    -> int:
													return (
													    sum(self.embedding_dimension   )
													    + self.num_dynamic_real_features
													    + self.num_time_features
													    + self.num_static_real_features
													    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
													)





