


'''simple docstring'''



from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings





@maybe_allow_in_graph
class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=0.0			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = "geglu"			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = False			,     SCREAMING_SNAKE_CASE_ = False			,     SCREAMING_SNAKE_CASE_ = False			,     SCREAMING_SNAKE_CASE_ = False			,     SCREAMING_SNAKE_CASE_ = True			,     SCREAMING_SNAKE_CASE_ = "layer_norm"			,     SCREAMING_SNAKE_CASE_ = False			,     )				->     List[Any]:
         super().__init__()
         __lowerCamelCase   :	Optional[int]								=		only_cross_attention

         __lowerCamelCase   :	List[str]								=		(num_embeds_ada_norm is not None) and norm_type == 'ada_norm_zero'
         __lowerCamelCase   :	Union[str, Any]								=		(num_embeds_ada_norm is not None) and norm_type == 'ada_norm'

         if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
              raise ValueError(
                  f'`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'
                  f' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'	)

         # Define 3 blocks. Each block has its own normalization layer.
         # 1. Self-Attn
         if self.use_ada_layer_norm:
              __lowerCamelCase   :	Optional[int]								=		AdaLayerNorm(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         elif self.use_ada_layer_norm_zero:
              __lowerCamelCase   :	List[Any]								=		AdaLayerNormZero(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         else:
              __lowerCamelCase   :	Dict								=		nn.LayerNorm(SCREAMING_SNAKE_CASE_			,     elementwise_affine=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		Attention(
             query_dim=SCREAMING_SNAKE_CASE_			,     heads=SCREAMING_SNAKE_CASE_			,     dim_head=SCREAMING_SNAKE_CASE_			,     dropout=SCREAMING_SNAKE_CASE_			,     bias=SCREAMING_SNAKE_CASE_			,     cross_attention_dim=cross_attention_dim if only_cross_attention else None			,     upcast_attention=SCREAMING_SNAKE_CASE_			,     )

         # 2. Cross-Attn
         if cross_attention_dim is not None or double_self_attention:
              # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
              # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
              # the second cross attention block.
              __lowerCamelCase   :	Dict								=		(
                  AdaLayerNorm(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
                  if self.use_ada_layer_norm
                  else nn.LayerNorm(SCREAMING_SNAKE_CASE_			,     elementwise_affine=SCREAMING_SNAKE_CASE_	)
              )
              __lowerCamelCase   :	List[str]								=		Attention(
                  query_dim=SCREAMING_SNAKE_CASE_			,     cross_attention_dim=cross_attention_dim if not double_self_attention else None			,     heads=SCREAMING_SNAKE_CASE_			,     dim_head=SCREAMING_SNAKE_CASE_			,     dropout=SCREAMING_SNAKE_CASE_			,     bias=SCREAMING_SNAKE_CASE_			,     upcast_attention=SCREAMING_SNAKE_CASE_			,     )  # is self-attn if encoder_hidden_states is none
         else:
              __lowerCamelCase   :	Dict								=		None
              __lowerCamelCase   :	Tuple								=		None

         # 3. Feed-forward
         __lowerCamelCase   :	Optional[Any]								=		nn.LayerNorm(SCREAMING_SNAKE_CASE_			,     elementwise_affine=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[Any]								=		FeedForward(SCREAMING_SNAKE_CASE_			,     dropout=SCREAMING_SNAKE_CASE_			,     activation_fn=SCREAMING_SNAKE_CASE_			,     final_dropout=SCREAMING_SNAKE_CASE_	)

         # let chunk size default to None
         __lowerCamelCase   :	str								=		None
         __lowerCamelCase   :	List[str]								=		0







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         # Sets chunk feed-forward
         __lowerCamelCase   :	Optional[Any]								=		chunk_size
         __lowerCamelCase   :	int								=		dim







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = None			,     )				->     List[Any]:
         # Notice that normalization is always applied before the real computation in the following blocks.
         # 1. Self-Attention
         if self.use_ada_layer_norm:
              __lowerCamelCase   :	List[Any]								=		self.norma(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         elif self.use_ada_layer_norm_zero:
              __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase   :	Any								=		self.norma(
                  SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     hidden_dtype=hidden_states.dtype	)
         else:
              __lowerCamelCase   :	Dict								=		self.norma(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Tuple								=		cross_attention_kwargs if cross_attention_kwargs is not None else {}

         __lowerCamelCase   :	str								=		self.attna(
             SCREAMING_SNAKE_CASE_			,     encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None			,     attention_mask=SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_			,     )
         if self.use_ada_layer_norm_zero:
              __lowerCamelCase   :	int								=		gate_msa.unsqueeze(1	) * attn_output
         __lowerCamelCase   :	Any								=		attn_output + hidden_states

         # 2. Cross-Attention
         if self.attna is not None:
              __lowerCamelCase   :	Any								=		(
                  self.norma(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) if self.use_ada_layer_norm else self.norma(SCREAMING_SNAKE_CASE_	)
              )

              __lowerCamelCase   :	int								=		self.attna(
                  SCREAMING_SNAKE_CASE_			,     encoder_hidden_states=SCREAMING_SNAKE_CASE_			,     attention_mask=SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_			,     )
              __lowerCamelCase   :	Optional[int]								=		attn_output + hidden_states

         # 3. Feed-forward
         __lowerCamelCase   :	str								=		self.norma(SCREAMING_SNAKE_CASE_	)

         if self.use_ada_layer_norm_zero:
              __lowerCamelCase   :	Dict								=		norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

         if self._chunk_size is not None:
              # "feed_forward_chunk_size" can be used to save memory
              if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
                   raise ValueError(
                       f'`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'	)

              __lowerCamelCase   :	Any								=		norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
              __lowerCamelCase   :	Optional[int]								=		torch.cat(
                  [self.ff(SCREAMING_SNAKE_CASE_	) for hid_slice in norm_hidden_states.chunk(SCREAMING_SNAKE_CASE_			,     dim=self._chunk_dim	)]			,     dim=self._chunk_dim			,     )
         else:
              __lowerCamelCase   :	Any								=		self.ff(SCREAMING_SNAKE_CASE_	)

         if self.use_ada_layer_norm_zero:
              __lowerCamelCase   :	List[str]								=		gate_mlp.unsqueeze(1	) * ff_output

         __lowerCamelCase   :	int								=		ff_output + hidden_states

         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = 4			,     SCREAMING_SNAKE_CASE_ = 0.0			,     SCREAMING_SNAKE_CASE_ = "geglu"			,     SCREAMING_SNAKE_CASE_ = False			,     )				->     Dict:
         super().__init__()
         __lowerCamelCase   :	Optional[int]								=		int(dim * mult	)
         __lowerCamelCase   :	Any								=		dim_out if dim_out is not None else dim

         if activation_fn == "gelu":
              __lowerCamelCase   :	Optional[Any]								=		GELU(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         if activation_fn == "gelu-approximate":
              __lowerCamelCase   :	Union[str, Any]								=		GELU(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     approximate='tanh'	)
         elif activation_fn == "geglu":
              __lowerCamelCase   :	Dict								=		GEGLU(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         elif activation_fn == "geglu-approximate":
              __lowerCamelCase   :	Union[str, Any]								=		ApproximateGELU(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Tuple								=		nn.ModuleList([]	)
         # project in
         self.net.append(SCREAMING_SNAKE_CASE_	)
         # project dropout
         self.net.append(nn.Dropout(SCREAMING_SNAKE_CASE_	)	)
         # project out
         self.net.append(nn.Linear(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)	)
         # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
         if final_dropout:
              self.net.append(nn.Dropout(SCREAMING_SNAKE_CASE_	)	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Optional[Any]:
         for module in self.net:
              __lowerCamelCase   :	List[str]								=		module(SCREAMING_SNAKE_CASE_	)
         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = "none"	)				->     Tuple:
         super().__init__()
         __lowerCamelCase   :	Optional[Any]								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[Any]								=		approximate







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         if gate.device.type != "mps":
              return F.gelu(SCREAMING_SNAKE_CASE_			,     approximate=self.approximate	)
         # mps: gelu is not implemented for float16
         return F.gelu(gate.to(dtype=torch.floataa	)			,     approximate=self.approximate	).to(dtype=gate.dtype	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[str]:
         __lowerCamelCase   :	str								=		self.proj(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[str]								=		self.gelu(SCREAMING_SNAKE_CASE_	)
         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     str:
         super().__init__()
         __lowerCamelCase   :	str								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     dim_out * 2	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         if gate.device.type != "mps":
              return F.gelu(SCREAMING_SNAKE_CASE_	)
         # mps: gelu is not implemented for float16
         return F.gelu(gate.to(dtype=torch.floataa	)	).to(dtype=gate.dtype	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase				,       __lowerCamelCase   :	Optional[Any]								=		self.proj(SCREAMING_SNAKE_CASE_	).chunk(2			,     dim=-1	)
         return hidden_states * self.gelu(SCREAMING_SNAKE_CASE_	)







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     List[str]:
         super().__init__()
         __lowerCamelCase   :	Optional[int]								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	int								=		self.proj(SCREAMING_SNAKE_CASE_	)
         return x * torch.sigmoid(1.7_0_2 * x	)







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     int:
         super().__init__()
         __lowerCamelCase   :	List[str]								=		nn.Embedding(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[str]								=		nn.SiLU()
         __lowerCamelCase   :	List[str]								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     embedding_dim * 2	)
         __lowerCamelCase   :	str								=		nn.LayerNorm(SCREAMING_SNAKE_CASE_			,     elementwise_affine=SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Dict								=		self.linear(self.silu(self.emb(SCREAMING_SNAKE_CASE_	)	)	)
         __lowerCamelCase				,       __lowerCamelCase   :	List[Any]								=		torch.chunk(SCREAMING_SNAKE_CASE_			,     2	)
         __lowerCamelCase   :	Optional[Any]								=		self.norm(SCREAMING_SNAKE_CASE_	) * (1 + scale) + shift
         return x







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Tuple:
         super().__init__()

         __lowerCamelCase   :	Optional[Any]								=		CombinedTimestepLabelEmbeddings(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Any								=		nn.SiLU()
         __lowerCamelCase   :	Dict								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     6 * embedding_dim			,     bias=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		nn.LayerNorm(SCREAMING_SNAKE_CASE_			,     elementwise_affine=SCREAMING_SNAKE_CASE_			,     eps=1E-6	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=None	)				->     Union[str, Any]:
         __lowerCamelCase   :	Union[str, Any]								=		self.linear(self.silu(self.emb(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     hidden_dtype=SCREAMING_SNAKE_CASE_	)	)	)
         __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase   :	Tuple								=		emb.chunk(6			,     dim=1	)
         __lowerCamelCase   :	Optional[Any]								=		self.norm(SCREAMING_SNAKE_CASE_	) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return x, gate_msa, shift_mlp, scale_mlp, gate_mlp







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = None			,     SCREAMING_SNAKE_CASE_ = 1E-5	)				->     Dict:
         super().__init__()
         __lowerCamelCase   :	Union[str, Any]								=		num_groups
         __lowerCamelCase   :	Any								=		eps

         if act_fn is None:
              __lowerCamelCase   :	Dict								=		None
         else:
              __lowerCamelCase   :	Tuple								=		get_activation(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	List[Any]								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     out_dim * 2	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     int:
         if self.act:
              __lowerCamelCase   :	Tuple								=		self.act(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[int]								=		self.linear(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	str								=		emb[:, :, None, None]
         __lowerCamelCase				,       __lowerCamelCase   :	int								=		emb.chunk(2			,     dim=1	)

         __lowerCamelCase   :	Union[str, Any]								=		F.group_norm(SCREAMING_SNAKE_CASE_			,     self.num_groups			,     eps=self.eps	)
         __lowerCamelCase   :	Dict								=		x * (1 + scale) + shift
         return x










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


A__		:					List[str]												=  get_logger(__name__)


A__		:					str												=  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_  :





    """simple docstring"""






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







class    UpperCAmelCase_  :





    """simple docstring"""






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







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    @add_start_docstrings(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         for processor in self:
              __lowerCamelCase   :	str								=		inspect.signature(processor.__call__	).parameters
              if len(SCREAMING_SNAKE_CASE_	) > 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   :	Tuple								=		processor(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_	)
              else:
                   __lowerCamelCase   :	int								=		processor(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






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

         __lowerCamelCase   :	Optional[int]								=		temperature







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	Dict								=		scores / self.temperature
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = -float('Inf'	)			,     SCREAMING_SNAKE_CASE_ = 1	)				->     Union[str, Any]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) 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(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) 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   :	str								=		top_p
         __lowerCamelCase   :	Tuple								=		filter_value
         __lowerCamelCase   :	Tuple								=		min_tokens_to_keep







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

         __lowerCamelCase   :	int								=		jnp.full_like(SCREAMING_SNAKE_CASE_			,     self.filter_value	)
         __lowerCamelCase   :	Tuple								=		jax.nn.softmax(SCREAMING_SNAKE_CASE_			,     axis=-1	).cumsum(axis=-1	)
         __lowerCamelCase   :	List[str]								=		cumulative_probs < self.top_p

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

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

         __lowerCamelCase   :	Any								=		jnp.where(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		jax.lax.sort_key_val(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)[-1]

         return next_scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = -float('Inf'	)			,     SCREAMING_SNAKE_CASE_ = 1	)				->     str:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or top_k <= 0:
              raise ValueError(f'`top_k` has to be a strictly positive integer, but is {top_k}'	)

         __lowerCamelCase   :	List[str]								=		max(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Dict								=		filter_value







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase				,       __lowerCamelCase   :	List[Any]								=		scores.shape
         __lowerCamelCase   :	Tuple								=		jnp.full(batch_size * vocab_size			,     self.filter_value	)

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

         __lowerCamelCase   :	Tuple								=		next_scores_flat.at[topk_indices_flat].set(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		next_scores_flat.reshape(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         return next_scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         __lowerCamelCase   :	Any								=		bos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	Optional[Any]								=		jnp.full(scores.shape			,     -float('inf'	)	)

         __lowerCamelCase   :	Optional[Any]								=		1 - jnp.bool_(cur_len - 1	)

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

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Tuple								=		max_length
         __lowerCamelCase   :	Any								=		eos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	List[str]								=		jnp.full(scores.shape			,     -float('inf'	)	)

         __lowerCamelCase   :	Any								=		1 - jnp.bool_(cur_len - self.max_length + 1	)

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

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     List[str]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or min_length < 0:
              raise ValueError(f'`min_length` has to be a positive integer, but is {min_length}'	)

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

         __lowerCamelCase   :	str								=		min_length
         __lowerCamelCase   :	Optional[int]								=		eos_token_id







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

         __lowerCamelCase   :	str								=		jnp.where(SCREAMING_SNAKE_CASE_			,     scores.at[:, self.eos_token_id].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Union[str, Any]								=		list(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	str								=		begin_index







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

         __lowerCamelCase   :	str								=		jnp.where(SCREAMING_SNAKE_CASE_			,     scores.at[:, self.begin_suppress_tokens].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[Any]:
         __lowerCamelCase   :	Tuple								=		list(SCREAMING_SNAKE_CASE_	)







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

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






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







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         def _force_token(SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	List[str]								=		scores.shape[0]
              __lowerCamelCase   :	Tuple								=		self.force_token_array[generation_idx]

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

         __lowerCamelCase   :	int								=		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(SCREAMING_SNAKE_CASE_	)			,     lambda: scores			,     )			,     )
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	Any								=		generate_config.eos_token_id
         __lowerCamelCase   :	Dict								=		generate_config.no_timestamps_token_id
         __lowerCamelCase   :	Tuple								=		generate_config.no_timestamps_token_id + 1

         __lowerCamelCase   :	List[str]								=		decoder_input_length + 1

         if generate_config.is_multilingual:
              # room for language token and task token
              self.begin_index += 2
         if hasattr(SCREAMING_SNAKE_CASE_			,     'max_initial_timestamp_index'	):
              __lowerCamelCase   :	str								=		generate_config.max_initial_timestamp_index
         else:
              __lowerCamelCase   :	Optional[int]								=		model_config.vocab_size
         if self.max_initial_timestamp_index is None:
              __lowerCamelCase   :	Tuple								=		model_config.vocab_size







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

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

              __lowerCamelCase   :	Any								=		jnp.where((cur_len - self.begin_index) < 2			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
              __lowerCamelCase   :	str								=		jnp.where(
                  input_ids_k[cur_len - 2] >= self.timestamp_begin			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     )

              return jnp.where(
                  SCREAMING_SNAKE_CASE_			,     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'	)	)			,     )			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	List[Any]								=		jax.vmap(SCREAMING_SNAKE_CASE_	)(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Any								=		jnp.where(cur_len == self.begin_index			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		jnp.where(
             self.max_initial_timestamp_index is not None			,     True and apply_max_initial_timestamp			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	Any								=		self.timestamp_begin + self.max_initial_timestamp_index

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

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

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

         __lowerCamelCase   :	Dict								=		jax.vmap(SCREAMING_SNAKE_CASE_	)(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         return scores






