





from __future__ import annotations







class     __lowerCAmelCase	:







      def __init__(   self		,				lowerCAmelCase__							)       ->						str:



        '''simple docstring'''




        a__     : int  =TypeError(
            "Matrices must be formed from a list of zero or more lists containing at "
            "least one and the same number of values, each of which must be of type "
            "int or float."							)
        if len(lowerCAmelCase__							) != 0:
          a__     : List[str]  =len(rows[0]							)
          if cols == 0:
            raise error
          for row in rows:
            if len(lowerCAmelCase__							) != cols:
              raise error
            for value in row:
              if not isinstance(lowerCAmelCase__		,				(int, float)							):
                raise error
          a__     : List[Any]  =rows
        else:
          a__     : str  =[]







      def        _lowercase  (   self							)       ->						list[list[int]]:



        '''simple docstring'''




        return [[row[i] for row in self.rows] for i in range(len(self.rows[0]							)							)]







      @property
      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        return len(self.rows							)







      @property
      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        return len(self.rows[0]							)







      @property
      def        _lowercase  (   self							)       ->						tuple[int, int]:



        '''simple docstring'''




        return (self.num_rows, self.num_columns)







      @property
      def        _lowercase  (   self							)       ->						bool:



        '''simple docstring'''




        return self.order[0] == self.order[1]







      def        _lowercase  (   self							)       ->						Matrix:



        '''simple docstring'''




        a__     : str  =[
            [0 if column_num != row_num else 1 for column_num in range(self.num_rows							)]
            for row_num in range(self.num_rows							)
        ]
        return Matrix(lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        if not self.is_square:
          return 0
        if self.order == (0, 0):
          return 1
        if self.order == (1, 1):
          return int(self.rows[0][0]							)
        if self.order == (2, 2):
          return int(
              (self.rows[0][0] * self.rows[1][1])
              - (self.rows[0][1] * self.rows[1][0])							)
        else:
          return sum(
              self.rows[0][column] * self.cofactors().rows[0][column]
              for column in range(self.num_columns							)							)







      def        _lowercase  (   self							)       ->						bool:



        '''simple docstring'''




        return bool(self.determinant()							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        a__     : List[str]  =[
            [
                self.rows[other_row][other_column]
                for other_column in range(self.num_columns							)
                if other_column != column
            ]
            for other_row in range(self.num_rows							)
            if other_row != row
        ]
        return Matrix(lowerCAmelCase__							).determinant()







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        if (row + column) % 2 == 0:
          return self.get_minor(lowerCAmelCase__		,				lowerCAmelCase__							)
        return -1 * self.get_minor(lowerCAmelCase__		,				lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Matrix:



        '''simple docstring'''




        return Matrix(
            [
                [self.get_minor(lowerCAmelCase__		,				lowerCAmelCase__							) for column in range(self.num_columns							)]
                for row in range(self.num_rows							)
            ]							)







      def        _lowercase  (   self							)       ->						Matrix:



        '''simple docstring'''




        return Matrix(
            [
                [
                    self.minors().rows[row][column]
                    if (row + column) % 2 == 0
                    else self.minors().rows[row][column] * -1
                    for column in range(self.minors().num_columns							)
                ]
                for row in range(self.minors().num_rows							)
            ]							)







      def        _lowercase  (   self							)       ->						Matrix:



        '''simple docstring'''




        a__     : Dict  =[
            [self.cofactors().rows[column][row] for column in range(self.num_columns							)]
            for row in range(self.num_rows							)
        ]
        return Matrix(lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Matrix:



        '''simple docstring'''




        a__     : Union[str, Any]  =self.determinant()
        if not determinant:
          raise TypeError("Only matrices with a non-zero determinant have an inverse"							)
        return self.adjugate() * (1 / determinant)







      def __repr__(   self							)       ->						str:



        '''simple docstring'''




        return str(self.rows							)







      def __str__(   self							)       ->						str:



        '''simple docstring'''




        if self.num_rows == 0:
          return "[]"
        if self.num_rows == 1:
          return "[[" + ". ".join(str(self.rows[0]							)							) + "]]"
        return (
            "["
            + "\n ".join(
                [
                    "[" + ". ".join([str(lowerCAmelCase__							) for value in row]							) + ".]"
                    for row in self.rows
                ]							)
            + "]"
        )







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__ = None							)       ->						None:



        '''simple docstring'''




        a__     : List[str]  =TypeError("Row must be a list containing all ints and/or floats"							)
        if not isinstance(lowerCAmelCase__		,				lowerCAmelCase__							):
          raise type_error
        for value in row:
          if not isinstance(lowerCAmelCase__		,				(int, float)							):
            raise type_error
        if len(lowerCAmelCase__							) != self.num_columns:
          raise ValueError(
              "Row must be equal in length to the other rows in the matrix"							)
        if position is None:
          self.rows.append(lowerCAmelCase__							)
        else:
          a__     : Tuple  =self.rows[0:position] + [row] + self.rows[position:]







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__ = None							)       ->						None:



        '''simple docstring'''




        a__     : str  =TypeError(
            "Column must be a list containing all ints and/or floats"							)
        if not isinstance(lowerCAmelCase__		,				lowerCAmelCase__							):
          raise type_error
        for value in column:
          if not isinstance(lowerCAmelCase__		,				(int, float)							):
            raise type_error
        if len(lowerCAmelCase__							) != self.num_rows:
          raise ValueError(
              "Column must be equal in length to the other columns in the matrix"							)
        if position is None:
          a__     : Optional[Any]  =[self.rows[i] + [column[i]] for i in range(self.num_rows							)]
        else:
          a__     : Any  =[
              self.rows[i][0:position] + [column[i]] + self.rows[i][position:]
              for i in range(self.num_rows							)
          ]







      def __eq__(   self		,				lowerCAmelCase__							)       ->						bool:



        '''simple docstring'''




        if not isinstance(lowerCAmelCase__		,				lowerCAmelCase__							):
          return NotImplemented
        return self.rows == other.rows







      def __ne__(   self		,				lowerCAmelCase__							)       ->						bool:



        '''simple docstring'''




        return not self == other







      def __neg__(   self							)       ->						Matrix:



        '''simple docstring'''




        return self * -1







      def __add__(   self		,				lowerCAmelCase__							)       ->						Matrix:



        '''simple docstring'''




        if self.order != other.order:
          raise ValueError("Addition requires matrices of the same order"							)
        return Matrix(
            [
                [self.rows[i][j] + other.rows[i][j] for j in range(self.num_columns							)]
                for i in range(self.num_rows							)
            ]							)







      def __sub__(   self		,				lowerCAmelCase__							)       ->						Matrix:



        '''simple docstring'''




        if self.order != other.order:
          raise ValueError("Subtraction requires matrices of the same order"							)
        return Matrix(
            [
                [self.rows[i][j] - other.rows[i][j] for j in range(self.num_columns							)]
                for i in range(self.num_rows							)
            ]							)







      def __mul__(   self		,				lowerCAmelCase__							)       ->						Matrix:



        '''simple docstring'''




        if isinstance(lowerCAmelCase__		,				(int, float)							):
          return Matrix(
              [[int(element * other							) for element in row] for row in self.rows]							)
        elif isinstance(lowerCAmelCase__		,				lowerCAmelCase__							):
          if self.num_columns != other.num_rows:
            raise ValueError(
                "The number of columns in the first matrix must "
                "be equal to the number of rows in the second"							)
          return Matrix(
              [
                  [Matrix.dot_product(lowerCAmelCase__		,				lowerCAmelCase__							) for column in other.columns()]
                  for row in self.rows
              ]							)
        else:
          raise TypeError(
              "A Matrix can only be multiplied by an int, float, or another matrix"							)







      def __pow__(   self		,				lowerCAmelCase__							)       ->						Matrix:



        '''simple docstring'''




        if not isinstance(lowerCAmelCase__		,				lowerCAmelCase__							):
          raise TypeError("A Matrix can only be raised to the power of an int"							)
        if not self.is_square:
          raise ValueError("Only square matrices can be raised to a power"							)
        if other == 0:
          return self.identity()
        if other < 0:
          if self.is_invertable():
            return self.inverse() ** (-other)
          raise ValueError(
              "Only invertable matrices can be raised to a negative power"							)
        a__     : Tuple  =self
        for _ in range(other - 1							):
          result *= self
        return result







      @classmethod
      def        _lowercase  (   cls		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        return sum(row[i] * column[i] for i in range(len(lowerCAmelCase__							)							)							)


if __name__ == "__main__":
     import doctest

     doctest.testmod()













import unittest

from transformers import AutoTokenizer, NystromformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         NystromformerForMaskedLM,
         NystromformerForMultipleChoice,
         NystromformerForQuestionAnswering,
         NystromformerForSequenceClassification,
         NystromformerForTokenClassification,
         NystromformerModel,
     )
     from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST







class     __lowerCAmelCase	:







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__=1_3		,				lowerCAmelCase__=7		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=9_9		,				lowerCAmelCase__=3_2		,				lowerCAmelCase__=5		,				lowerCAmelCase__=4		,				lowerCAmelCase__=3_7		,				lowerCAmelCase__="gelu"		,				lowerCAmelCase__=0.1		,				lowerCAmelCase__=0.1		,				lowerCAmelCase__=5_1_2		,				lowerCAmelCase__=1_6		,				lowerCAmelCase__=2		,				lowerCAmelCase__=0.02		,				lowerCAmelCase__=3		,				lowerCAmelCase__=4		,				lowerCAmelCase__=None		,				)       ->						List[Any]:



        '''simple docstring'''




        a__     : Optional[Any]  =parent
        a__     : Tuple  =batch_size
        a__     : List[Any]  =seq_length
        a__     : Dict  =is_training
        a__     : Any  =use_input_mask
        a__     : int  =use_token_type_ids
        a__     : Optional[Any]  =use_labels
        a__     : Optional[Any]  =vocab_size
        a__     : List[str]  =hidden_size
        a__     : int  =num_hidden_layers
        a__     : Tuple  =num_attention_heads
        a__     : Union[str, Any]  =intermediate_size
        a__     : Optional[int]  =hidden_act
        a__     : int  =hidden_dropout_prob
        a__     : Union[str, Any]  =attention_probs_dropout_prob
        a__     : List[Any]  =max_position_embeddings
        a__     : str  =type_vocab_size
        a__     : Optional[Any]  =type_sequence_label_size
        a__     : Union[str, Any]  =initializer_range
        a__     : List[Any]  =num_labels
        a__     : str  =num_choices
        a__     : int  =scope







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : List[Any]  =ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size							)

        a__     : str  =None
        if self.use_input_mask:
          a__     : List[Any]  =random_attention_mask([self.batch_size, self.seq_length]							)

        a__     : str  =None
        if self.use_token_type_ids:
          a__     : List[str]  =ids_tensor([self.batch_size, self.seq_length]		,				self.type_vocab_size							)

        a__     : Dict  =None
        a__     : str  =None
        a__     : str  =None
        if self.use_labels:
          a__     : List[str]  =ids_tensor([self.batch_size]		,				self.type_sequence_label_size							)
          a__     : Tuple  =ids_tensor([self.batch_size, self.seq_length]		,				self.num_labels							)
          a__     : Dict  =ids_tensor([self.batch_size]		,				self.num_choices							)

        a__     : Tuple  =self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




        return NystromformerConfig(
            vocab_size=self.vocab_size		,				hidden_size=self.hidden_size		,				num_hidden_layers=self.num_hidden_layers		,				num_attention_heads=self.num_attention_heads		,				intermediate_size=self.intermediate_size		,				hidden_act=self.hidden_act		,				hidden_dropout_prob=self.hidden_dropout_prob		,				attention_probs_dropout_prob=self.attention_probs_dropout_prob		,				max_position_embeddings=self.max_position_embeddings		,				type_vocab_size=self.type_vocab_size		,				is_decoder=lowerCAmelCase__		,				initializer_range=self.initializer_range		,				)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						List[Any]:



        '''simple docstring'''




        a__     : Tuple  =NystromformerModel(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : Optional[Any]  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__							)
        a__     : str  =model(lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__							)
        a__     : Optional[int]  =model(lowerCAmelCase__							)
        self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.hidden_size)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        a__     : int  =NystromformerForMaskedLM(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : Dict  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Any:



        '''simple docstring'''




        a__     : Optional[int]  =NystromformerForQuestionAnswering(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : str  =model(
            lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				start_positions=lowerCAmelCase__		,				end_positions=lowerCAmelCase__		,				)
        self.parent.assertEqual(result.start_logits.shape		,				(self.batch_size, self.seq_length)							)
        self.parent.assertEqual(result.end_logits.shape		,				(self.batch_size, self.seq_length)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : Optional[Any]  =self.num_labels
        a__     : Dict  =NystromformerForSequenceClassification(lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : List[str]  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_labels)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						List[str]:



        '''simple docstring'''




        a__     : Tuple  =self.num_labels
        a__     : List[str]  =NystromformerForTokenClassification(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : List[Any]  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.num_labels)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Tuple:



        '''simple docstring'''




        a__     : List[Any]  =self.num_choices
        a__     : Optional[Any]  =NystromformerForMultipleChoice(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : List[str]  =input_ids.unsqueeze(1							).expand(-1		,				self.num_choices		,				-1							).contiguous()
        a__     : List[Any]  =token_type_ids.unsqueeze(1							).expand(-1		,				self.num_choices		,				-1							).contiguous()
        a__     : List[Any]  =input_mask.unsqueeze(1							).expand(-1		,				self.num_choices		,				-1							).contiguous()
        a__     : Dict  =model(
            lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__		,				)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_choices)							)







      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




        a__     : Optional[Any]  =self.prepare_config_and_inputs()
        (
            (
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							(
            a__
        ) ,							
        )     : List[str]  =config_and_inputs
        a__     : str  ={"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
        return config, inputs_dict







@require_torch
class     __lowerCAmelCase	(   UpperCamelCase__						,    UpperCamelCase__						,    unittest.TestCase):
      _lowercase    : int            =      (
          (
              NystromformerModel,
              NystromformerForMaskedLM,
              NystromformerForMultipleChoice,
              NystromformerForQuestionAnswering,
              NystromformerForSequenceClassification,
              NystromformerForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _lowercase    : Union[str, Any]            =      (
          {
              """feature-extraction""": NystromformerModel,
              """fill-mask""": NystromformerForMaskedLM,
              """question-answering""": NystromformerForQuestionAnswering,
              """text-classification""": NystromformerForSequenceClassification,
              """token-classification""": NystromformerForTokenClassification,
              """zero-shot""": NystromformerForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _lowercase    : Union[str, Any]            =      False
      _lowercase    : Union[str, Any]            =      False







      def        _lowercase  (   self							)       ->						List[str]:



        '''simple docstring'''




        a__     : Optional[int]  =NystromformerModelTester(self							)
        a__     : Optional[int]  =ConfigTester(self		,				config_class=lowerCAmelCase__		,				hidden_size=3_7							)







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        self.config_tester.run_common_tests()







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        a__     : List[str]  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          a__     : int  =type
          self.model_tester.create_and_check_model(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : Optional[Any]  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        a__     : str  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						List[str]:



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : int  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__							)







      @slow
      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__     : int  =NystromformerModel.from_pretrained(lowerCAmelCase__							)
          self.assertIsNotNone(lowerCAmelCase__							)







@require_torch
class     __lowerCAmelCase	(   unittest.TestCase):







      @slow
      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        a__     : str  =NystromformerModel.from_pretrained("uw-madison/nystromformer-512"							)
        a__     : int  =torch.tensor([[0, 1, 2, 3, 4, 5]]							)

        with torch.no_grad():
          a__     : Tuple  =model(lowerCAmelCase__							)[0]

        a__     : List[str]  =torch.Size((1, 6, 7_6_8)							)
        self.assertEqual(output.shape		,				lowerCAmelCase__							)

        a__     : int  =torch.tensor(
            [[[-0.45_32, -0.09_36, 0.51_37], [-0.26_76, 0.06_28, 0.61_86], [-0.36_29, -0.17_26, 0.47_16]]]							)

        self.assertTrue(torch.allclose(output[:, :3, :3]		,				lowerCAmelCase__		,				atol=1E-4							)							)







      @slow
      def        _lowercase  (   self							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : Union[str, Any]  ="the [MASK] of Belgium is Brussels"

        a__     : str  =AutoTokenizer.from_pretrained("uw-madison/nystromformer-512"							)
        a__     : int  =NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512"							)

        a__     : List[Any]  =tokenizer(lowerCAmelCase__		,				return_tensors="pt"							)

        with torch.no_grad():
          a__     : str  =model(encoding.input_ids							).logits

        a__     : List[str]  =token_logits[:, 2, :].argmax(-1							)[0]

        self.assertEqual(tokenizer.decode(lowerCAmelCase__							)		,				"capital"							)






