

from __future__ import annotations

from random import random
from typing import Generic, TypeVar

lowerCAmelCase   :      List[Any]     =     TypeVar('KT')
lowerCAmelCase   :      Tuple     =     TypeVar('VT')
class        _A						(     Generic[KT, VT]):

  def __init__( self       ,    _SCREAMING_SNAKE_CASE = "root"       ,    _SCREAMING_SNAKE_CASE = None    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     Optional[Any]							=   key
        SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   value
        SCREAMING_SNAKE_CASE_ :     list[Node[KT, VT]]							=   []

  def __repr__( self    ):





        """simple docstring"""







        return f"Node({self.key}: {self.value})"



  @property
  def        UpperCAmelCase			( self    ):





        """simple docstring"""







        return len(self.forward    )

class        _A						(     Generic[KT, VT]):

  def __init__( self       ,    _SCREAMING_SNAKE_CASE = 0.5       ,    _SCREAMING_SNAKE_CASE = 16    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     Node[KT, VT]							=   Node[KT, VT]()
        SCREAMING_SNAKE_CASE_ :     Tuple							=   0
        SCREAMING_SNAKE_CASE_ :     Any							=   p
        SCREAMING_SNAKE_CASE_ :     Optional[Any]							=   max_level

  def __str__( self    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     str							=   list(self    )

        if len(snake_case_    ) == 0:
              return f"SkipList(level={self.level})"

        SCREAMING_SNAKE_CASE_ :     str							=   max((len(str(snake_case_    )    ) for item in items)       ,    default=4    )
        SCREAMING_SNAKE_CASE_ :     Any							=   max(snake_case_       ,    4    ) + 4

        SCREAMING_SNAKE_CASE_ :     Any							=   self.head
        SCREAMING_SNAKE_CASE_ :     Any							=   []

        SCREAMING_SNAKE_CASE_ :     List[Any]							=   node.forward.copy()
        lines.append(f"[{node.key}]".ljust(snake_case_       ,    '-'    ) + '* ' * len(snake_case_    )    )
        lines.append(' ' * label_size + '| ' * len(snake_case_    )    )

        while len(node.forward    ) != 0:
              SCREAMING_SNAKE_CASE_ :     Optional[Any]							=   node.forward[0]

              lines.append(
                  f"[{node.key}]".ljust(snake_case_       ,    '-'    )
                  + ' '.join(str(n.key    ) if n.key == node.key else '|' for n in forwards    )    )
              lines.append(' ' * label_size + '| ' * len(snake_case_    )    )
              SCREAMING_SNAKE_CASE_ :     Dict							=   node.forward

        lines.append('None'.ljust(snake_case_    ) + '* ' * len(snake_case_    )    )
        return f"SkipList(level={self.level})\n" + "\n".join(snake_case_    )

  def __iter__( self    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     List[str]							=   self.head

        while len(node.forward    ) != 0:
              yield node.forward[0].key
              SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   node.forward[0]

  def        UpperCAmelCase			( self    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     List[str]							=   1
        while random() < self.p and level < self.max_level:
              level += 1

        return level

  def        UpperCAmelCase			( self       ,    _SCREAMING_SNAKE_CASE    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     Dict							=   []

        SCREAMING_SNAKE_CASE_ :     Optional[Any]							=   self.head

        for i in reversed(range(self.level    )    ):
              # i < node.level - When node level is lesser than `i` decrement `i`.
              # node.forward[i].key < key - Jumping to node with key value higher
              #                             or equal to searched key would result
              #                             in skipping searched key.
              while i < node.level and node.forward[i].key < key:
                    SCREAMING_SNAKE_CASE_ :     str							=   node.forward[i]
              # Each leftmost node (relative to searched node) will potentially have to
              # be updated.
              update_vector.append(snake_case_    )

        update_vector.reverse()  # Note that we were inserting values in reverse order.

        # len(node.forward) != 0 - If current node doesn't contain any further
        #                          references then searched key is not present.
        # node.forward[0].key == key - Next node key should be equal to search key
        #                              if key is present.
        if len(node.forward    ) != 0 and node.forward[0].key == key:
              return node.forward[0], update_vector
        else:
              return None, update_vector

  def        UpperCAmelCase			( self       ,    _SCREAMING_SNAKE_CASE    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     Tuple							=   self._locate_node(snake_case_    )

        if node is not None:
              for i, update_node in enumerate(snake_case_    ):
                    # Remove or replace all references to removed node.
                    if update_node.level > i and update_node.forward[i].key == key:
                          if node.level > i:
                                SCREAMING_SNAKE_CASE_ :     Optional[int]							=   node.forward[i]
                          else:
                                SCREAMING_SNAKE_CASE_ :     int							=   update_node.forward[:i]

  def        UpperCAmelCase			( self       ,    _SCREAMING_SNAKE_CASE       ,    _SCREAMING_SNAKE_CASE    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     Optional[int]							=   self._locate_node(snake_case_    )
        if node is not None:
              SCREAMING_SNAKE_CASE_ :     List[Any]							=   value
        else:
              SCREAMING_SNAKE_CASE_ :     Any							=   self.random_level()

              if level > self.level:
                    # After level increase we have to add additional nodes to head.
                    for _ in range(self.level - 1       ,    snake_case_    ):
                          update_vector.append(self.head    )
                    SCREAMING_SNAKE_CASE_ :     List[str]							=   level

              SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   Node(snake_case_       ,    snake_case_    )

              for i, update_node in enumerate(update_vector[:level]    ):
                    # Change references to pass through new node.
                    if update_node.level > i:
                          new_node.forward.append(update_node.forward[i]    )

                    if update_node.level < i + 1:
                          update_node.forward.append(snake_case_    )
                    else:
                          SCREAMING_SNAKE_CASE_ :     Dict							=   new_node



  def        UpperCAmelCase			( self       ,    _SCREAMING_SNAKE_CASE    ):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_ :     Dict							=   self._locate_node(snake_case_    )

        if node is not None:
              return node.value

        return None

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     Optional[int]							=   SkipList()
      skip_list.insert('Key1'    ,      3     )
      skip_list.insert('Key2'    ,      1_2     )
      skip_list.insert('Key3'    ,      4_1     )
      skip_list.insert('Key4'    ,      -1_9     )

      SCREAMING_SNAKE_CASE_ :     Tuple							=   skip_list.head
      SCREAMING_SNAKE_CASE_ :     List[str]							=   {}
      while node.level != 0:
            SCREAMING_SNAKE_CASE_ :     List[Any]							=   node.forward[0]
            SCREAMING_SNAKE_CASE_ :     List[str]							=   node.value

      assert len(_UpperCAmelCase     ) == 4
      assert all_values["Key1"] == 3
      assert all_values["Key2"] == 1_2
      assert all_values["Key3"] == 4_1
      assert all_values["Key4"] == -1_9

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   SkipList()
      skip_list.insert('Key1'    ,      1_0     )
      skip_list.insert('Key1'    ,      1_2     )

      skip_list.insert('Key5'    ,      7     )
      skip_list.insert('Key7'    ,      1_0     )
      skip_list.insert('Key10'    ,      5     )

      skip_list.insert('Key7'    ,      7     )
      skip_list.insert('Key5'    ,      5     )
      skip_list.insert('Key10'    ,      1_0     )

      SCREAMING_SNAKE_CASE_ :     Dict							=   skip_list.head
      SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   {}
      while node.level != 0:
            SCREAMING_SNAKE_CASE_ :     str							=   node.forward[0]
            SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   node.value

      if len(_UpperCAmelCase     ) != 4:
            print()
      assert len(_UpperCAmelCase     ) == 4
      assert all_values["Key1"] == 1_2
      assert all_values["Key7"] == 7
      assert all_values["Key5"] == 5
      assert all_values["Key10"] == 1_0

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     Dict							=   SkipList()
      assert skip_list.find('Some key'     ) is None

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     Tuple							=   SkipList()

      skip_list.insert('Key2'    ,      2_0     )
      assert skip_list.find('Key2'     ) == 2_0

      skip_list.insert('Some Key'    ,      1_0     )
      skip_list.insert('Key2'    ,      8     )
      skip_list.insert('V'    ,      1_3     )

      assert skip_list.find('Y'     ) is None
      assert skip_list.find('Key2'     ) == 8
      assert skip_list.find('Some Key'     ) == 1_0
      assert skip_list.find('V'     ) == 1_3

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     int							=   SkipList()
      skip_list.delete('Some key'     )

      assert len(skip_list.head.forward     ) == 0

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     int							=   SkipList()

      skip_list.insert('Key1'    ,      1_2     )
      skip_list.insert('V'    ,      1_3     )
      skip_list.insert('X'    ,      1_4     )
      skip_list.insert('Key2'    ,      1_5     )

      skip_list.delete('V'     )
      skip_list.delete('Key2'     )

      assert skip_list.find('V'     ) is None
      assert skip_list.find('Key2'     ) is None

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     Union[str, Any]							=   SkipList()

      skip_list.insert('Key1'    ,      1_2     )
      skip_list.insert('V'    ,      1_3     )
      skip_list.insert('X'    ,      1_4     )
      skip_list.insert('Key2'    ,      1_5     )

      skip_list.delete('V'     )
      assert skip_list.find('V'     ) is None
      assert skip_list.find('X'     ) == 1_4
      assert skip_list.find('Key1'     ) == 1_2
      assert skip_list.find('Key2'     ) == 1_5

      skip_list.delete('X'     )
      assert skip_list.find('V'     ) is None
      assert skip_list.find('X'     ) is None
      assert skip_list.find('Key1'     ) == 1_2
      assert skip_list.find('Key2'     ) == 1_5

      skip_list.delete('Key1'     )
      assert skip_list.find('V'     ) is None
      assert skip_list.find('X'     ) is None
      assert skip_list.find('Key1'     ) is None
      assert skip_list.find('Key2'     ) == 1_5

      skip_list.delete('Key2'     )
      assert skip_list.find('V'     ) is None
      assert skip_list.find('X'     ) is None
      assert skip_list.find('Key1'     ) is None
      assert skip_list.find('Key2'     ) is None

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     Optional[Any]							=   SkipList()

      skip_list.insert('Key1'    ,      1_2     )
      skip_list.insert('V'    ,      1_3     )
      skip_list.insert('X'    ,      1_4_2     )
      skip_list.insert('Key2'    ,      1_5     )

      skip_list.delete('X'     )

      def traverse_keys(a     ):
            yield node.key
            for forward_node in node.forward:
                  yield from traverse_keys(_UpperCAmelCase     )

      assert len(set(traverse_keys(skip_list.head     )     )     ) == 4

def       A_     (							):

      """simple docstring"""



      def is_sorted(a     ):
            return all(next_item >= item for item, next_item in zip(_UpperCAmelCase    ,      lst[1:]     )     )

      SCREAMING_SNAKE_CASE_ :     List[str]							=   SkipList()
      for i in range(1_0     ):
            skip_list.insert(_UpperCAmelCase    ,      _UpperCAmelCase     )
      assert is_sorted(list(_UpperCAmelCase     )     )
      skip_list.delete(5     )
      skip_list.delete(8     )
      skip_list.delete(2     )
      assert is_sorted(list(_UpperCAmelCase     )     )
      skip_list.insert(-1_2    ,      -1_2     )
      skip_list.insert(7_7    ,      7_7     )
      assert is_sorted(list(_UpperCAmelCase     )     )

def       A_     (							):

      """simple docstring"""



      for _ in range(1_0_0     ):
            # Repeat test 100 times due to the probabilistic nature of skip list
            # random values == random bugs
            test_insert()
            test_insert_overrides_existing_value()

            test_searching_empty_list_returns_none()
            test_search()

            test_deleting_item_from_empty_list_do_nothing()
            test_deleted_items_are_not_founded_by_find_method()
            test_delete_removes_only_given_key()
            test_delete_doesnt_leave_dead_nodes()

            test_iter_always_yields_sorted_values()

def       A_     (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_ :     Optional[Any]							=   SkipList()
      skip_list.insert(2    ,      '2'     )
      skip_list.insert(4    ,      '4'     )
      skip_list.insert(6    ,      '4'     )
      skip_list.insert(4    ,      '5'     )
      skip_list.insert(8    ,      '4'     )
      skip_list.insert(9    ,      '4'     )

      skip_list.delete(4     )

      print(_UpperCAmelCase     )


if __name__ == "__main__":
    import doctest

    doctest.testmod()
    main()













"""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_     :						Dict						=					get_logger(__name__)


lowerCamelCase_     :						List[str]						=					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  :



      '''simple docstring'''

      @add_start_docstrings(snake_case_  )
      def __call__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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


class  _UpperCAmelCase  :



      '''simple docstring'''

      @add_start_docstrings(snake_case_  )
      def __call__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        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(snake_case_  )
      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      **snake_case_  ):

        """simple docstring"""

        for processor in self:
          A_  :       Tuple										=  inspect.signature(processor.__call__  ).parameters
          if len(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."""  )
            A_  :       Tuple										=  processor(snake_case_					,      snake_case_					,      snake_case_					,      **snake_case_  )
          else:
            A_  :       Optional[Any]										=  processor(snake_case_					,      snake_case_					,      snake_case_  )
        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

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

        A_  :       Optional[int]										=  temperature

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       int										=  scores / self.temperature
        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_ = -float('Inf'  )					,      snake_case_ = 1  ):

        """simple docstring"""

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

        A_  :       str										=  top_p
        A_  :       Union[str, Any]										=  filter_value
        A_  :       int										=  min_tokens_to_keep

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_						,       A_  :       Tuple										=  lax.top_k(snake_case_					,      scores.shape[-1]  )

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

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

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

        A_  :       Optional[Any]										=  jnp.where(snake_case_					,      snake_case_					,      snake_case_  )
        A_  :       List[Any]										=  jax.lax.sort_key_val(snake_case_					,      snake_case_  )[-1]

        return next_scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_ = -float('Inf'  )					,      snake_case_ = 1  ):

        """simple docstring"""

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

        A_  :       str										=  max(snake_case_					,      snake_case_  )
        A_  :       Union[str, Any]										=  filter_value

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_						,       A_  :       int										=  scores.shape
        A_  :       Tuple										=  jnp.full(batch_size * vocab_size					,      self.filter_value  )

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

        A_  :       List[str]										=  next_scores_flat.at[topk_indices_flat].set(snake_case_  )
        A_  :       Union[str, Any]										=  next_scores_flat.reshape(snake_case_					,      snake_case_  )
        return next_scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  bos_token_id

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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

        A_  :       Union[str, Any]										=  1 - jnp.bool_(cur_len - 1  )

        A_  :       str										=  jnp.where(snake_case_					,      new_scores.at[:, self.bos_token_id].set(0  )					,      snake_case_  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Dict										=  max_length
        A_  :       Optional[int]										=  eos_token_id

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  jnp.full(scores.shape					,      -float('inf'  )  )

        A_  :       Dict										=  1 - jnp.bool_(cur_len - self.max_length + 1  )

        A_  :       Dict										=  jnp.where(snake_case_					,      new_scores.at[:, self.eos_token_id].set(0  )					,      snake_case_  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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

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

        A_  :       Any										=  min_length
        A_  :       List[Any]										=  eos_token_id

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       int										=  1 - jnp.clip(cur_len - self.min_length					,      0					,      1  )

        A_  :       Optional[Any]										=  jnp.where(snake_case_					,      scores.at[:, self.eos_token_id].set(-float('inf'  )  )					,      snake_case_  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       List[Any]										=  list(snake_case_  )
        A_  :       Tuple										=  begin_index

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Dict										=  1 - jnp.bool_(cur_len - self.begin_index  )

        A_  :       int										=  jnp.where(snake_case_					,      scores.at[:, self.begin_suppress_tokens].set(-float('inf'  )  )					,      snake_case_  )

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       List[Any]										=  list(snake_case_  )

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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

        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Any										=  dict(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.
        A_  :       Tuple										=  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_  :       Tuple										=  force_token_array.at[index].set(snake_case_  )
        A_  :       Any										=  jnp.intaa(snake_case_  )

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        def _force_token(snake_case_  ):
          A_  :       List[Any]										=  scores.shape[0]
          A_  :       Any										=  self.force_token_array[generation_idx]

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

        A_  :       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(snake_case_  )					,      lambda: scores					,      )					,      )
        return scores


class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Tuple										=  generate_config.eos_token_id
        A_  :       Optional[int]										=  generate_config.no_timestamps_token_id
        A_  :       List[str]										=  generate_config.no_timestamps_token_id + 1

        A_  :       Any										=  decoder_input_length + 1

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

      def __call__(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

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

        def handle_pairs(snake_case_					,      snake_case_  ):
          A_  :       Any										=  jnp.where((cur_len - self.begin_index) >= 1					,      snake_case_					,      snake_case_  )
          A_  :       Tuple										=  jnp.where(
              input_ids_k[cur_len - 1] >= self.timestamp_begin					,      True and last_was_timestamp					,      snake_case_					,      )

          A_  :       Tuple										=  jnp.where((cur_len - self.begin_index) < 2					,      snake_case_					,      snake_case_  )
          A_  :       Any										=  jnp.where(
              input_ids_k[cur_len - 2] >= self.timestamp_begin					,      snake_case_					,      snake_case_					,      )

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

        A_  :       Tuple										=  jax.vmap(snake_case_  )(snake_case_					,      snake_case_  )

        A_  :       Optional[Any]										=  jnp.where(cur_len == self.begin_index					,      snake_case_					,      snake_case_  )
        A_  :       Tuple										=  jnp.where(
            self.max_initial_timestamp_index is not None					,      True and apply_max_initial_timestamp					,      snake_case_					,      )

        A_  :       int										=  self.timestamp_begin + self.max_initial_timestamp_index

        A_  :       List[Any]										=  jnp.where(
            snake_case_					,      scores.at[:, last_allowed + 1 :].set(-float('inf'  )  )					,      snake_case_					,      )

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

        def handle_cumulative_probs(snake_case_					,      snake_case_  ):
          A_  :       Dict										=  jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]					,      axis=-1  )
          A_  :       Optional[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'  )  )					,      snake_case_					,      )

        A_  :       Union[str, Any]										=  jax.vmap(snake_case_  )(snake_case_					,      snake_case_  )

        return scores