






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   :     Any						, __lowercase   :     int						, __lowercase   :     Optional[int]=13						, __lowercase   :     Optional[int]=7						, __lowercase   :     str=True						, __lowercase   :     int=True						, __lowercase   :     Optional[int]=True						, __lowercase   :     Optional[Any]=True						, __lowercase   :     Tuple=99						, __lowercase   :     Any=32						, __lowercase   :     List[Any]=5						, __lowercase   :     Optional[int]=4						, __lowercase   :     Dict=37						, __lowercase   :     List[Any]="gelu"						, __lowercase   :     Tuple=0.1						, __lowercase   :     Union[str, Any]=0.1						, __lowercase   :     str=128						, __lowercase   :     List[str]=32						, __lowercase   :     Union[str, Any]=16						, __lowercase   :     Optional[int]=2						, __lowercase   :     Optional[int]=0.02						, __lowercase   :     Optional[Any]=3						, __lowercase   :     List[str]=4						, __lowercase   :     List[str]=None						, )				-> List[str]:
        __UpperCAmelCase				:      str	 =    parent
        __UpperCAmelCase				:      Dict	 =    batch_size
        __UpperCAmelCase				:      Tuple	 =    seq_length
        __UpperCAmelCase				:      Dict	 =    is_training
        __UpperCAmelCase				:      str	 =    use_input_mask
        __UpperCAmelCase				:      str	 =    use_token_type_ids
        __UpperCAmelCase				:      List[str]	 =    use_labels
        __UpperCAmelCase				:      Optional[Any]	 =    vocab_size
        __UpperCAmelCase				:      List[str]	 =    hidden_size
        __UpperCAmelCase				:      int	 =    num_hidden_layers
        __UpperCAmelCase				:      Union[str, Any]	 =    num_attention_heads
        __UpperCAmelCase				:      List[Any]	 =    intermediate_size
        __UpperCAmelCase				:      List[Any]	 =    hidden_act
        __UpperCAmelCase				:      Tuple	 =    hidden_dropout_prob
        __UpperCAmelCase				:      Any	 =    attention_probs_dropout_prob
        __UpperCAmelCase				:      Any	 =    max_position_embeddings
        __UpperCAmelCase				:      str	 =    type_vocab_size
        __UpperCAmelCase				:      List[Any]	 =    type_sequence_label_size
        __UpperCAmelCase				:      Optional[int]	 =    initializer_range
        __UpperCAmelCase				:      int	 =    num_labels
        __UpperCAmelCase				:      str	 =    num_choices
        __UpperCAmelCase				:      int	 =    scope
      def       UpperCAmelCase ( self   :     Tuple   )				-> Any:
        __UpperCAmelCase				:      Tuple	 =    ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size   )

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

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

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

        __UpperCAmelCase				:      Any	 =    self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
      def       UpperCAmelCase ( self   :     Optional[int]   )				-> Tuple:
        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=__lowercase						, initializer_range=self.initializer_range						, )
      def       UpperCAmelCase ( self   :     Optional[int]   )				-> int:
        (
            (
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					
        )				:      Any	 =    self.prepare_config_and_inputs()

        __UpperCAmelCase				:      Any	 =    True
        __UpperCAmelCase				:      Union[str, Any]	 =    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
        __UpperCAmelCase				:      Tuple	 =    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       UpperCAmelCase ( self   :     Optional[Any]						, __lowercase   :     Dict						, __lowercase   :     Tuple						, __lowercase   :     List[Any]						, __lowercase   :     List[str]						, __lowercase   :     Optional[Any]						, __lowercase   :     Optional[Any]						, __lowercase   :     Optional[Any]   )				-> Any:
        __UpperCAmelCase				:      Dict	 =    NezhaModel(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      List[str]	 =    model(__lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase   )
        __UpperCAmelCase				:      Union[str, Any]	 =    model(__lowercase						, token_type_ids=__lowercase   )
        __UpperCAmelCase				:      List[Any]	 =    model(__lowercase   )
        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       UpperCAmelCase ( self   :     str						, __lowercase   :     str						, __lowercase   :     Any						, __lowercase   :     int						, __lowercase   :     str						, __lowercase   :     int						, __lowercase   :     Dict						, __lowercase   :     Dict						, __lowercase   :     Tuple						, __lowercase   :     Tuple						, )				-> Tuple:
        __UpperCAmelCase				:      Dict	 =    True
        __UpperCAmelCase				:      Union[str, Any]	 =    NezhaModel(__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      Tuple	 =    model(
            __lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, encoder_hidden_states=__lowercase						, encoder_attention_mask=__lowercase						, )
        __UpperCAmelCase				:      Optional[int]	 =    model(
            __lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, encoder_hidden_states=__lowercase						, )
        __UpperCAmelCase				:      Optional[Any]	 =    model(__lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase   )
        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       UpperCAmelCase ( self   :     Optional[int]						, __lowercase   :     Optional[int]						, __lowercase   :     Optional[Any]						, __lowercase   :     Union[str, Any]						, __lowercase   :     List[Any]						, __lowercase   :     Tuple						, __lowercase   :     Tuple						, __lowercase   :     str   )				-> Dict:
        __UpperCAmelCase				:      int	 =    NezhaForMaskedLM(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      Union[str, Any]	 =    model(__lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, labels=__lowercase   )
        self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size)   )
      def       UpperCAmelCase ( self   :     Optional[Any]						, __lowercase   :     Optional[int]						, __lowercase   :     Dict						, __lowercase   :     int						, __lowercase   :     Optional[Any]						, __lowercase   :     Dict						, __lowercase   :     str						, __lowercase   :     Optional[int]   )				-> List[str]:
        __UpperCAmelCase				:      int	 =    NezhaForNextSentencePrediction(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      str	 =    model(
            __lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, labels=__lowercase						, )
        self.parent.assertEqual(result.logits.shape						, (self.batch_size, 2)   )
      def       UpperCAmelCase ( self   :     Tuple						, __lowercase   :     Optional[int]						, __lowercase   :     Union[str, Any]						, __lowercase   :     Dict						, __lowercase   :     str						, __lowercase   :     Dict						, __lowercase   :     Dict						, __lowercase   :     Optional[int]   )				-> Any:
        __UpperCAmelCase				:      Any	 =    NezhaForPreTraining(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      Optional[Any]	 =    model(
            __lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, labels=__lowercase						, next_sentence_label=__lowercase						, )
        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       UpperCAmelCase ( self   :     int						, __lowercase   :     int						, __lowercase   :     Tuple						, __lowercase   :     List[Any]						, __lowercase   :     Any						, __lowercase   :     Union[str, Any]						, __lowercase   :     List[str]						, __lowercase   :     Any   )				-> str:
        __UpperCAmelCase				:      Tuple	 =    NezhaForQuestionAnswering(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      Optional[Any]	 =    model(
            __lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, start_positions=__lowercase						, end_positions=__lowercase						, )
        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       UpperCAmelCase ( self   :     Optional[Any]						, __lowercase   :     Optional[int]						, __lowercase   :     Union[str, Any]						, __lowercase   :     List[Any]						, __lowercase   :     Optional[int]						, __lowercase   :     Optional[Any]						, __lowercase   :     int						, __lowercase   :     List[Any]   )				-> int:
        __UpperCAmelCase				:      Union[str, Any]	 =    self.num_labels
        __UpperCAmelCase				:      List[str]	 =    NezhaForSequenceClassification(__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      Optional[int]	 =    model(__lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, labels=__lowercase   )
        self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_labels)   )
      def       UpperCAmelCase ( self   :     Optional[Any]						, __lowercase   :     Tuple						, __lowercase   :     int						, __lowercase   :     str						, __lowercase   :     int						, __lowercase   :     Optional[int]						, __lowercase   :     Dict						, __lowercase   :     int   )				-> Optional[int]:
        __UpperCAmelCase				:      Any	 =    self.num_labels
        __UpperCAmelCase				:      Tuple	 =    NezhaForTokenClassification(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      List[str]	 =    model(__lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, labels=__lowercase   )
        self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.num_labels)   )
      def       UpperCAmelCase ( self   :     int						, __lowercase   :     Tuple						, __lowercase   :     Dict						, __lowercase   :     str						, __lowercase   :     Union[str, Any]						, __lowercase   :     Tuple						, __lowercase   :     Any						, __lowercase   :     Union[str, Any]   )				-> List[Any]:
        __UpperCAmelCase				:      Any	 =    self.num_choices
        __UpperCAmelCase				:      List[Any]	 =    NezhaForMultipleChoice(config=__lowercase   )
        model.to(__lowercase   )
        model.eval()
        __UpperCAmelCase				:      List[Any]	 =    input_ids.unsqueeze(1   ).expand(-1						, self.num_choices						, -1   ).contiguous()
        __UpperCAmelCase				:      List[Any]	 =    token_type_ids.unsqueeze(1   ).expand(-1						, self.num_choices						, -1   ).contiguous()
        __UpperCAmelCase				:      List[Any]	 =    input_mask.unsqueeze(1   ).expand(-1						, self.num_choices						, -1   ).contiguous()
        __UpperCAmelCase				:      Any	 =    model(
            __lowercase						, attention_mask=__lowercase						, token_type_ids=__lowercase						, labels=__lowercase						, )
        self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_choices)   )




      def       UpperCAmelCase ( self   :     int   )				-> int:
        __UpperCAmelCase				:      Union[str, Any]	 =    self.prepare_config_and_inputs()
        (
            (
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					
        )				:      List[Any]	 =    config_and_inputs
        __UpperCAmelCase				:      Optional[Any]	 =    {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
        return config, inputs_dict


@require_torch
class        a       (					lowercase__				,			lowercase__				,			lowercase__				,			unittest.TestCase  ):



      """simple docstring"""


      a    :					Any							=      (
          (
              NezhaModel,
              NezhaForMaskedLM,
              NezhaForMultipleChoice,
              NezhaForNextSentencePrediction,
              NezhaForPreTraining,
              NezhaForQuestionAnswering,
              NezhaForSequenceClassification,
              NezhaForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      a    :					int							=      (
          {
              'feature-extraction': NezhaModel,
              'fill-mask': NezhaForMaskedLM,
              'question-answering': NezhaForQuestionAnswering,
              'text-classification': NezhaForSequenceClassification,
              'token-classification': NezhaForTokenClassification,
              'zero-shot': NezhaForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      a    :					int							=      True
      def       UpperCAmelCase ( self   :     List[Any]						, __lowercase   :     int						, __lowercase   :     Dict						, __lowercase   :     List[Any]=False   )				-> Dict:
        __UpperCAmelCase				:      List[Any]	 =    super()._prepare_for_class(__lowercase						, __lowercase						, return_labels=__lowercase   )

        if return_labels:
          if model_class in get_values(__lowercase   ):
            __UpperCAmelCase				:      Tuple	 =    torch.zeros(
                (self.model_tester.batch_size, self.model_tester.seq_length)						, dtype=torch.long						, device=__lowercase   )
            __UpperCAmelCase				:      int	 =    torch.zeros(
                self.model_tester.batch_size						, dtype=torch.long						, device=__lowercase   )
        return inputs_dict
      def       UpperCAmelCase ( self   :     List[Any]   )				-> Tuple:
        __UpperCAmelCase				:      Dict	 =    NezhaModelTester(self   )
        __UpperCAmelCase				:      List[Any]	 =    ConfigTester(self						, config_class=__lowercase						, hidden_size=37   )
      def       UpperCAmelCase ( self   :     int   )				-> str:
        self.config_tester.run_common_tests()
      def       UpperCAmelCase ( self   :     List[Any]   )				-> Dict:
        __UpperCAmelCase				:      Any	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase   )
      def       UpperCAmelCase ( self   :     int   )				-> Optional[int]:
        __UpperCAmelCase				:      Any	 =    self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*__lowercase   )
      def       UpperCAmelCase ( self   :     List[str]   )				-> int:
        # This regression test was failing with PyTorch < 1.3
        (
            (
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					(
            __UpperCAmelCase
        ) ,					
        )				:      str	 =    self.model_tester.prepare_config_and_inputs_for_decoder()

        __UpperCAmelCase				:      str	 =    None

        self.model_tester.create_and_check_model_as_decoder(
            __lowercase						, __lowercase						, __lowercase						, __lowercase						, __lowercase						, __lowercase						, __lowercase						, __lowercase						, __lowercase						, )
      def       UpperCAmelCase ( self   :     List[Any]   )				-> Optional[int]:
        __UpperCAmelCase				:      Tuple	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowercase   )
      def       UpperCAmelCase ( self   :     Optional[int]   )				-> Tuple:
        __UpperCAmelCase				:      Dict	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowercase   )
      def       UpperCAmelCase ( self   :     Any   )				-> Union[str, Any]:
        __UpperCAmelCase				:      Optional[Any]	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_next_sequence_prediction(*__lowercase   )
      def       UpperCAmelCase ( self   :     Optional[Any]   )				-> Tuple:
        __UpperCAmelCase				:      List[str]	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*__lowercase   )
      def       UpperCAmelCase ( self   :     List[str]   )				-> Any:
        __UpperCAmelCase				:      int	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowercase   )
      def       UpperCAmelCase ( self   :     int   )				-> Union[str, Any]:
        __UpperCAmelCase				:      str	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowercase   )
      def       UpperCAmelCase ( self   :     Optional[Any]   )				-> str:
        __UpperCAmelCase				:      Union[str, Any]	 =    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowercase   )
      @slow
      def       UpperCAmelCase ( self   :     List[Any]   )				-> Any:
        for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          __UpperCAmelCase				:      Union[str, Any]	 =    NezhaModel.from_pretrained(__lowercase   )
          self.assertIsNotNone(__lowercase   )




      @slow
      @require_torch_gpu
      def       UpperCAmelCase ( self   :     Dict   )				-> Optional[Any]:
        __UpperCAmelCase ,					__UpperCAmelCase				:      Dict	 =    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

          __UpperCAmelCase				:      Tuple	 =    True
          __UpperCAmelCase				:      Dict	 =    model_class(config=__lowercase   )

          __UpperCAmelCase				:      List[str]	 =    self._prepare_for_class(__lowercase						, __lowercase   )
          __UpperCAmelCase				:      Tuple	 =    torch.jit.trace(
              __lowercase						, (inputs_dict["""input_ids"""].to("""cpu"""   ), inputs_dict["""attention_mask"""].to("""cpu"""   ))   )

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


@require_torch
class        a       (					unittest.TestCase  ):



      """simple docstring"""


      @slow
      def       UpperCAmelCase ( self   :     Any   )				-> List[Any]:
        __UpperCAmelCase				:      Optional[Any]	 =    NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""   )
        __UpperCAmelCase				:      Tuple	 =    torch.tensor([[0, 1, 2, 3, 4, 5]]   )
        __UpperCAmelCase				:      str	 =    torch.tensor([[0, 1, 1, 1, 1, 1]]   )
        with torch.no_grad():
          __UpperCAmelCase				:      Optional[Any]	 =    model(__lowercase						, attention_mask=__lowercase   )[0]
        __UpperCAmelCase				:      Union[str, Any]	 =    torch.Size((1, 6, 768)   )
        self.assertEqual(output.shape						, __lowercase   )
        __UpperCAmelCase				:      List[Any]	 =    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]						, __lowercase						, atol=1e-4   )   )




      @slow
      def       UpperCAmelCase ( self   :     List[Any]   )				-> Optional[Any]:
        __UpperCAmelCase				:      str	 =    NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""   )
        __UpperCAmelCase				:      Dict	 =    torch.tensor([[0, 1, 2, 3, 4, 5]]   )
        __UpperCAmelCase				:      Union[str, Any]	 =    torch.tensor([[1, 1, 1, 1, 1, 1]]   )
        with torch.no_grad():
          __UpperCAmelCase				:      int	 =    model(__lowercase						, attention_mask=__lowercase   )[0]
        __UpperCAmelCase				:      Optional[int]	 =    torch.Size((1, 6, 21128)   )
        self.assertEqual(output.shape						, __lowercase   )
        __UpperCAmelCase				:      List[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]						, __lowercase						, atol=1e-4   )   )



'''simple docstring'''




from manim import *




class        a__ (  UpperCAmelCase__					):






     def   SCREAMING_SNAKE_CASE__  (					self						:						List[Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    Rectangle(height=0.5    ,							width=0.5       )
         __lowerCamelCase            =    Rectangle(height=0.46    ,							width=0.46       ).set_stroke(width=0       )

         __lowerCamelCase            =    Rectangle(height=0.25    ,							width=0.25       )

         __lowerCamelCase            =    [mem.copy() for i in range(6       )]
         __lowerCamelCase            =    [mem.copy() for i in range(6       )]
         __lowerCamelCase            =    VGroup(*a       ).arrange(a    ,							buff=0       )
         __lowerCamelCase            =    VGroup(*a       ).arrange(a    ,							buff=0       )
         __lowerCamelCase            =    VGroup(a    ,							a       ).arrange(a    ,							buff=0       )
         __lowerCamelCase            =    Text('''CPU'''    ,							font_size=24       )
         __lowerCamelCase            =    Group(a    ,							a       ).arrange(a    ,							buff=0.5    ,							aligned_edge=a       )
         cpu.move_to([-2.5, -0.5, 0]       )
         self.add(a       )

         __lowerCamelCase            =    [mem.copy() for i in range(4       )]
         __lowerCamelCase            =    VGroup(*a       ).arrange(a    ,							buff=0       )
         __lowerCamelCase            =    Text('''GPU'''    ,							font_size=24       )
         __lowerCamelCase            =    Group(a    ,							a       ).arrange(a    ,							buff=0.5    ,							aligned_edge=a       )
         gpu.move_to([-1, -1, 0]       )
         self.add(a       )

         __lowerCamelCase            =    [mem.copy() for i in range(6       )]
         __lowerCamelCase            =    VGroup(*a       ).arrange(a    ,							buff=0       )

         __lowerCamelCase            =    Text('''Model'''    ,							font_size=24       )
         __lowerCamelCase            =    Group(a    ,							a       ).arrange(a    ,							buff=0.5    ,							aligned_edge=a       )
         model.move_to([3, -1.0, 0]       )
         self.add(a       )

         __lowerCamelCase            =    []
         __lowerCamelCase            =    []

         for i, rect in enumerate(a       ):
             __lowerCamelCase            =    fill.copy().set_fill(a    ,							opacity=0.8       )
             target.move_to(a       )
             model_arr.append(a       )

             __lowerCamelCase            =    Rectangle(height=0.46    ,							width=0.46       ).set_stroke(width=0.0       ).set_fill(a    ,							opacity=0.8       )
             cpu_target.move_to(cpu_left_col_base[i]       )
             model_cpu_arr.append(a       )

         self.add(*a    ,							*a       )

         __lowerCamelCase            =    [meta_mem.copy() for i in range(6       )]
         __lowerCamelCase            =    [meta_mem.copy() for i in range(6       )]
         __lowerCamelCase            =    VGroup(*a       ).arrange(a    ,							buff=0       )
         __lowerCamelCase            =    VGroup(*a       ).arrange(a    ,							buff=0       )
         __lowerCamelCase            =    VGroup(a    ,							a       ).arrange(a    ,							buff=0       )
         __lowerCamelCase            =    Text('''Disk'''    ,							font_size=24       )
         __lowerCamelCase            =    Group(a    ,							a       ).arrange(a    ,							buff=0.5    ,							aligned_edge=a       )
         disk.move_to([-4, -1.25, 0]       )
         self.add(a    ,							a       )

         __lowerCamelCase            =    Square(side_length=2.2       )
         key.move_to([-5, 2, 0]       )

         __lowerCamelCase            =    MarkupText(
             f"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""    ,							font_size=18    ,							)

         key_text.move_to([-5, 2.4, 0]       )

         self.add(a    ,							a       )

         __lowerCamelCase            =    MarkupText(
             f"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""    ,							font_size=18    ,							)

         blue_text.next_to(a    ,							DOWN * 2.4    ,							aligned_edge=key_text.get_left()       )
         self.add(a       )

         __lowerCamelCase            =    MarkupText(
             f"""Now watch as an input is passed through the model\nand how the memory is utilized and handled."""    ,							font_size=24    ,							)
         step_a.move_to([2, 2, 0]       )

         self.play(Write(a       )       )

         __lowerCamelCase            =    Square(0.3       )
         input.set_fill(a    ,							opacity=1.0       )
         input.set_stroke(width=0.0       )
         input.next_to(model_base[0]    ,							a    ,							buff=0.5       )

         self.play(Write(a       )       )

         input.generate_target()
         input.target.next_to(model_arr[0]    ,							direction=a    ,							buff=0.02       )
         self.play(MoveToTarget(a       )       )

         self.play(FadeOut(a       )       )

         __lowerCamelCase            =    Arrow(start=a    ,							end=a    ,							color=a    ,							buff=0.5       )
         a.next_to(model_arr[0].get_left()    ,							a    ,							buff=0.2       )

         model_cpu_arr[0].generate_target()
         model_cpu_arr[0].target.move_to(gpu_rect[0]       )

         __lowerCamelCase            =    MarkupText(
             f"""As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back."""    ,							font_size=24    ,							)
         step_a.move_to([2, 2, 0]       )

         self.play(Write(a    ,							run_time=3       )       )

         __lowerCamelCase            =    {'''run_time''': 1, '''fade_in''': True, '''fade_out''': True, '''buff''': 0.02}

         self.play(
             Write(a       )    ,							Circumscribe(model_arr[0]    ,							color=a    ,							**a       )    ,							Circumscribe(model_cpu_arr[0]    ,							color=a    ,							**a       )    ,							Circumscribe(gpu_rect[0]    ,							color=a    ,							**a       )    ,							)
         self.play(MoveToTarget(model_cpu_arr[0]       )       )

         __lowerCamelCase            =    a.copy()
         for i in range(6       ):
             a_c.next_to(model_arr[i].get_right() + 0.02    ,							a    ,							buff=0.2       )

             input.generate_target()
             input.target.move_to(model_arr[i].get_right() + 0.02       )

             __lowerCamelCase            =    AnimationGroup(
                 FadeOut(a    ,							run_time=0.5       )    ,							MoveToTarget(a    ,							run_time=0.5       )    ,							FadeIn(a    ,							run_time=0.5       )    ,							lag_ratio=0.2       )

             self.play(a       )

             model_cpu_arr[i].generate_target()
             model_cpu_arr[i].target.move_to(cpu_left_col_base[i]       )

             if i < 5:
                 model_cpu_arr[i + 1].generate_target()
                 model_cpu_arr[i + 1].target.move_to(gpu_rect[0]       )
                 if i >= 1:
                     __lowerCamelCase            =    0.7

                 self.play(
                     Circumscribe(model_arr[i]    ,							**a       )    ,							Circumscribe(cpu_left_col_base[i]    ,							**a       )    ,							Circumscribe(cpu_left_col_base[i + 1]    ,							color=a    ,							**a       )    ,							Circumscribe(gpu_rect[0]    ,							color=a    ,							**a       )    ,							Circumscribe(model_arr[i + 1]    ,							color=a    ,							**a       )    ,							)
                 if i < 1:
                     self.play(
                         MoveToTarget(model_cpu_arr[i]       )    ,							MoveToTarget(model_cpu_arr[i + 1]       )    ,							)
                 else:
                     self.play(
                         MoveToTarget(model_cpu_arr[i]    ,							run_time=0.7       )    ,							MoveToTarget(model_cpu_arr[i + 1]    ,							run_time=0.7       )    ,							)
             else:
                 model_cpu_arr[i].generate_target()
                 model_cpu_arr[i].target.move_to(cpu_left_col_base[-1]       )
                 input.generate_target()
                 input.target.next_to(model_arr[-1].get_right()    ,							RIGHT + 0.02    ,							buff=0.2       )

                 self.play(
                     Circumscribe(model_arr[-1]    ,							color=a    ,							**a       )    ,							Circumscribe(cpu_left_col_base[-1]    ,							color=a    ,							**a       )    ,							Circumscribe(gpu_rect[0]    ,							color=a    ,							**a       )    ,							)

                 self.play(MoveToTarget(model_cpu_arr[i]       )       )

             __lowerCamelCase            =    a_c
             __lowerCamelCase            =    a_c.copy()

         input.generate_target()
         input.target.next_to(model_base[-1]    ,							RIGHT + 0.02    ,							buff=0.5       )
         self.play(
             FadeOut(a       )    ,							FadeOut(a    ,							run_time=0.5       )    ,							)

         __lowerCamelCase            =    MarkupText(f"""Inference on a model too large for GPU memory\nis successfully completed."""    ,							font_size=24       )
         step_a.move_to([2, 2, 0]       )

         self.play(Write(a    ,							run_time=3       )    ,							MoveToTarget(a       )       )

         self.wait()
