

"""simple docstring"""


import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      NezhaForMaskedLM,
      NezhaForMultipleChoice,
      NezhaForNextSentencePrediction,
      NezhaForPreTraining,
      NezhaForQuestionAnswering,
      NezhaForSequenceClassification,
      NezhaForTokenClassification,
      NezhaModel,
  )
  from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST

class      __a							:






     '''simple docstring'''

     def __init__(						self	,    _a	,    _a=13	,    _a=7	,    _a=True	,    _a=True	,    _a=True	,    _a=True	,    _a=99	,    _a=32	,    _a=5	,    _a=4	,    _a=37	,    _a="gelu"	,    _a=0.1	,    _a=0.1	,    _a=128	,    _a=32	,    _a=16	,    _a=2	,    _a=0.02	,    _a=3	,    _a=4	,    _a=None	,    )  ->      int:


           """simple docstring"""


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

     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size  )

           SCREAMING_SNAKE_CASE__   :      str      		=	None
           if self.use_input_mask:
                 SCREAMING_SNAKE_CASE__   :      int      		=	random_attention_mask([self.batch_size, self.seq_length]  )

           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	None
           if self.use_token_type_ids:
                 SCREAMING_SNAKE_CASE__   :      Tuple      		=	ids_tensor([self.batch_size, self.seq_length]	,    self.type_vocab_size  )

           SCREAMING_SNAKE_CASE__   :      str      		=	None
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	None
           SCREAMING_SNAKE_CASE__   :      Any      		=	None
           if self.use_labels:
                 SCREAMING_SNAKE_CASE__   :      Any      		=	ids_tensor([self.batch_size]	,    self.type_sequence_label_size  )
                 SCREAMING_SNAKE_CASE__   :      Dict      		=	ids_tensor([self.batch_size, self.seq_length]	,    self.num_labels  )
                 SCREAMING_SNAKE_CASE__   :      str      		=	ids_tensor([self.batch_size]	,    self.num_choices  )

           SCREAMING_SNAKE_CASE__   :      int      		=	self.get_config()

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

     def 					_a			(						self  )  ->      Dict:


           """simple docstring"""


           return NezhaConfig(
               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=_a	,    initializer_range=self.initializer_range	,    )

     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           (
               (
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		
           )   :      List[Any]      		=	self.prepare_config_and_inputs()

           SCREAMING_SNAKE_CASE__   :      Tuple      		=	True
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	ids_tensor([self.batch_size, self.seq_length]	,    vocab_size=2  )

           return (
               config,
               input_ids,
               token_type_ids,
               input_mask,
               sequence_labels,
               token_labels,
               choice_labels,
               encoder_hidden_states,
               encoder_attention_mask,
           )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	NezhaModel(config=_a  )
           model.to(_a  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      int      		=	model(_a	,    attention_mask=_a	,    token_type_ids=_a  )
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	model(_a	,    token_type_ids=_a  )
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	model(_a  )
           self.parent.assertEqual(result.last_hidden_state.shape	,    (self.batch_size, self.seq_length, self.hidden_size)  )
           self.parent.assertEqual(result.pooler_output.shape	,    (self.batch_size, self.hidden_size)  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	True
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	NezhaModel(_a  )
           model.to(_a  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Tuple      		=	model(
               _a	,    attention_mask=_a	,    token_type_ids=_a	,    encoder_hidden_states=_a	,    encoder_attention_mask=_a	,    )
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	model(
               _a	,    attention_mask=_a	,    token_type_ids=_a	,    encoder_hidden_states=_a	,    )
           SCREAMING_SNAKE_CASE__   :      str      		=	model(_a	,    attention_mask=_a	,    token_type_ids=_a  )
           self.parent.assertEqual(result.last_hidden_state.shape	,    (self.batch_size, self.seq_length, self.hidden_size)  )
           self.parent.assertEqual(result.pooler_output.shape	,    (self.batch_size, self.hidden_size)  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	NezhaForMaskedLM(config=_a  )
           model.to(_a  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Any      		=	model(_a	,    attention_mask=_a	,    token_type_ids=_a	,    labels=_a  )
           self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.vocab_size)  )

     def 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Dict:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	NezhaForNextSentencePrediction(config=_a  )
           model.to(_a  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Tuple      		=	model(
               _a	,    attention_mask=_a	,    token_type_ids=_a	,    labels=_a	,    )
           self.parent.assertEqual(result.logits.shape	,    (self.batch_size, 2)  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	NezhaForPreTraining(config=_a  )
           model.to(_a  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	model(
               _a	,    attention_mask=_a	,    token_type_ids=_a	,    labels=_a	,    next_sentence_label=_a	,    )
           self.parent.assertEqual(result.prediction_logits.shape	,    (self.batch_size, self.seq_length, self.vocab_size)  )
           self.parent.assertEqual(result.seq_relationship_logits.shape	,    (self.batch_size, 2)  )

     def 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Dict:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Any      		=	NezhaForQuestionAnswering(config=_a  )
           model.to(_a  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	model(
               _a	,    attention_mask=_a	,    token_type_ids=_a	,    start_positions=_a	,    end_positions=_a	,    )
           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 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Optional[int]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	self.num_labels
           SCREAMING_SNAKE_CASE__   :      int      		=	NezhaForSequenceClassification(_a  )
           model.to(_a  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	model(_a	,    attention_mask=_a	,    token_type_ids=_a	,    labels=_a  )
           self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.num_labels)  )

     def 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	self.num_labels
           SCREAMING_SNAKE_CASE__   :      Dict      		=	NezhaForTokenClassification(config=_a  )
           model.to(_a  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Dict      		=	model(_a	,    attention_mask=_a	,    token_type_ids=_a	,    labels=_a  )
           self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.num_labels)  )

     def 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Tuple:


           """simple docstring"""


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





     def 					_a			(						self  )  ->      Tuple:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      int      		=	self.prepare_config_and_inputs()
           (
               (
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		
           )   :      str      		=	config_and_inputs
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
           return config, inputs_dict





@require_torch
class      __a							(UpperCamelCase_   ,	UpperCamelCase_   ,	UpperCamelCase_   ,	unittest.TestCase):






     '''simple docstring'''
     _SCREAMING_SNAKE_CASE						:Union[str, Any]       = (
         (
             NezhaModel,
             NezhaForMaskedLM,
             NezhaForMultipleChoice,
             NezhaForNextSentencePrediction,
             NezhaForPreTraining,
             NezhaForQuestionAnswering,
             NezhaForSequenceClassification,
             NezhaForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     _SCREAMING_SNAKE_CASE						:List[Any]       = (
         {
             """feature-extraction""": NezhaModel,
             """fill-mask""": NezhaForMaskedLM,
             """question-answering""": NezhaForQuestionAnswering,
             """text-classification""": NezhaForSequenceClassification,
             """token-classification""": NezhaForTokenClassification,
             """zero-shot""": NezhaForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     _SCREAMING_SNAKE_CASE						:Optional[Any]       = True

     def 					_a			(						self	,    _a	,    _a	,    _a=False  )  ->      Tuple:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	super()._prepare_for_class(_a	,    _a	,    return_labels=_a  )

           if return_labels:
                 if model_class in get_values(_a  ):
                       SCREAMING_SNAKE_CASE__   :      Tuple      		=	torch.zeros(
                           (self.model_tester.batch_size, self.model_tester.seq_length)	,    dtype=torch.long	,    device=_a  )
                       SCREAMING_SNAKE_CASE__   :      Tuple      		=	torch.zeros(
                           self.model_tester.batch_size	,    dtype=torch.long	,    device=_a  )
           return inputs_dict

     def 					_a			(						self  )  ->      Any:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Tuple      		=	NezhaModelTester(self  )
           SCREAMING_SNAKE_CASE__   :      int      		=	ConfigTester(self	,    config_class=_a	,    hidden_size=37  )

     def 					_a			(						self  )  ->      Optional[int]:


           """simple docstring"""


           self.config_tester.run_common_tests()

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*_a  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Dict      		=	self.model_tester.prepare_config_and_inputs_for_decoder()
           self.model_tester.create_and_check_model_as_decoder(*_a  )

     def 					_a			(						self  )  ->      List[Any]:


           """simple docstring"""


           (
               (
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		(
               SCREAMING_SNAKE_CASE__
           )		,		
           )   :      Optional[int]      		=	self.model_tester.prepare_config_and_inputs_for_decoder()

           SCREAMING_SNAKE_CASE__   :      Any      		=	None

           self.model_tester.create_and_check_model_as_decoder(
               _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    )

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*_a  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*_a  )

     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      int      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_next_sequence_prediction(*_a  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*_a  )

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*_a  )

     def 					_a			(						self  )  ->      List[Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Any      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*_a  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      int      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*_a  )

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


           """simple docstring"""


           for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	NezhaModel.from_pretrained(_a  )
                 self.assertIsNotNone(_a  )





     @slow
     @require_torch_gpu
     def 					_a			(						self  )  ->      Any:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__		,		SCREAMING_SNAKE_CASE__   :      str      		=	self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes:
                 # NezhaForMultipleChoice behaves incorrectly in JIT environments.
                 if model_class == NezhaForMultipleChoice:
                       return

                 SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	True
                 SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	model_class(config=_a  )

                 SCREAMING_SNAKE_CASE__   :      Any      		=	self._prepare_for_class(_a	,    _a  )
                 SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	torch.jit.trace(
                     _a	,    (inputs_dict["""input_ids"""].to("""cpu"""  ), inputs_dict["""attention_mask"""].to("""cpu"""  ))  )

                 with tempfile.TemporaryDirectory() as tmp:
                       torch.jit.save(_a	,    os.path.join(_a	,    """bert.pt"""  )  )
                       SCREAMING_SNAKE_CASE__   :      Any      		=	torch.jit.load(os.path.join(_a	,    """bert.pt"""  )	,    map_location=_a  )
                       loaded(inputs_dict["""input_ids"""].to(_a  )	,    inputs_dict["""attention_mask"""].to(_a  )  )





@require_torch
class      __a							(unittest.TestCase):






     '''simple docstring'''

     @slow
     def 					_a			(						self  )  ->      Any:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""  )
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	torch.tensor([[0, 1, 2, 3, 4, 5]]  )
           SCREAMING_SNAKE_CASE__   :      str      		=	torch.tensor([[0, 1, 1, 1, 1, 1]]  )
           with torch.no_grad():
                 SCREAMING_SNAKE_CASE__   :      Dict      		=	model(_a	,    attention_mask=_a  )[0]
           SCREAMING_SNAKE_CASE__   :      str      		=	torch.Size((1, 6, 768)  )
           self.assertEqual(output.shape	,    _a  )
           SCREAMING_SNAKE_CASE__   :      int      		=	torch.tensor([[[0.0_685, 0.2_441, 0.1_102], [0.0_600, 0.1_906, 0.1_349], [0.0_221, 0.0_819, 0.0_586]]]  )

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





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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""  )
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	torch.tensor([[0, 1, 2, 3, 4, 5]]  )
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	torch.tensor([[1, 1, 1, 1, 1, 1]]  )
           with torch.no_grad():
                 SCREAMING_SNAKE_CASE__   :      Dict      		=	model(_a	,    attention_mask=_a  )[0]
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	torch.Size((1, 6, 21_128)  )
           self.assertEqual(output.shape	,    _a  )
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	torch.tensor(
               [[-2.7_939, -1.7_902, -2.2_189], [-2.8_585, -1.8_908, -2.3_723], [-2.6_499, -1.7_750, -2.2_558]]  )

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







"""simple docstring"""


from math import sqrt




def  _lowercase					(    __lowerCAmelCase							)       ->				bool:

      if 1 < number < 4:
            # 2 and 3 are primes
            return True
      elif number < 2 or number % 2 == 0 or number % 3 == 0:
            # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
            return False

      # All primes number are in format of 6k +/- 1
      for i in range(5  ,					int(sqrt(__lowerCAmelCase							) + 1							)  ,					6							):
            if number % i == 0 or number % (i + 2) == 0:
                  return False
      return True




def  _lowercase					(    __lowerCAmelCase = 1_0001							)       ->				int:
      SCREAMING_SNAKE_CASE__   :      Dict      		=	0
      SCREAMING_SNAKE_CASE__   :      Tuple      		=	1
      while count != nth and number < 3:
            number += 1
            if is_prime(__lowerCAmelCase							):
                  count += 1
      while count != nth:
            number += 2
            if is_prime(__lowerCAmelCase							):
                  count += 1
      return number


if __name__ == "__main__":
  print(f'{solution() = }')




