





"""simple docstring"""





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 							snake_case							:

      '''simple docstring'''



      def __init__(							self       :      Dict						,				__lowercase       :      Union[str, Any]						,				__lowercase       :      Dict=13						,				__lowercase       :      Union[str, Any]=7						,				__lowercase       :      Dict=True						,				__lowercase       :      int=True						,				__lowercase       :      int=True						,				__lowercase       :      List[Any]=True						,				__lowercase       :      str=99						,				__lowercase       :      Optional[Any]=32						,				__lowercase       :      Tuple=5						,				__lowercase       :      Any=4						,				__lowercase       :      List[str]=37						,				__lowercase       :      Any="gelu"						,				__lowercase       :      Optional[int]=0.1						,				__lowercase       :      Tuple=0.1						,				__lowercase       :      Any=512						,				__lowercase       :      Tuple=16						,				__lowercase       :      List[Any]=2						,				__lowercase       :      str=0.0_2						,				__lowercase       :      Optional[int]=3						,				__lowercase       :      Dict=4						,				__lowercase       :      int=None						,				):


        '''simple docstring'''




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

      def       A_					(							self       :      Optional[Any] ):


        '''simple docstring'''




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

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

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

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

        __UpperCAmelCase    :						int       =						self.get_config()

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

      def       A_					(							self       :      int ):


        '''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=__lowercase						,				initializer_range=self.initializer_range						,				)

      def       A_					(							self       :      Union[str, Any]						,				__lowercase       :      List[str]						,				__lowercase       :      Tuple						,				__lowercase       :      int						,				__lowercase       :      List[Any]						,				__lowercase       :      List[Any]						,				__lowercase       :      int						,				__lowercase       :      List[str] ):


        '''simple docstring'''




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

      def       A_					(							self       :      Union[str, Any]						,				__lowercase       :      Optional[int]						,				__lowercase       :      Optional[Any]						,				__lowercase       :      List[str]						,				__lowercase       :      Tuple						,				__lowercase       :      str						,				__lowercase       :      Tuple						,				__lowercase       :      Union[str, Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						NystromformerForMaskedLM(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, self.seq_length, self.vocab_size) )

      def       A_					(							self       :      Dict						,				__lowercase       :      List[Any]						,				__lowercase       :      Tuple						,				__lowercase       :      Dict						,				__lowercase       :      Dict						,				__lowercase       :      Dict						,				__lowercase       :      Union[str, Any]						,				__lowercase       :      List[Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Tuple       =						NystromformerForQuestionAnswering(config=__lowercase )
        model.to(__lowercase )
        model.eval()
        __UpperCAmelCase    :						Optional[int]       =						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       A_					(							self       :      str						,				__lowercase       :      Union[str, Any]						,				__lowercase       :      List[Any]						,				__lowercase       :      Optional[Any]						,				__lowercase       :      str						,				__lowercase       :      Any						,				__lowercase       :      str						,				__lowercase       :      Dict ):


        '''simple docstring'''




        __UpperCAmelCase    :						Any       =						self.num_labels
        __UpperCAmelCase    :						List[str]       =						NystromformerForSequenceClassification(__lowercase )
        model.to(__lowercase )
        model.eval()
        __UpperCAmelCase    :						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       A_					(							self       :      List[Any]						,				__lowercase       :      Dict						,				__lowercase       :      Union[str, Any]						,				__lowercase       :      List[Any]						,				__lowercase       :      Union[str, Any]						,				__lowercase       :      List[Any]						,				__lowercase       :      Tuple						,				__lowercase       :      List[Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						self.num_labels
        __UpperCAmelCase    :						Dict       =						NystromformerForTokenClassification(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, self.seq_length, self.num_labels) )

      def       A_					(							self       :      Union[str, Any]						,				__lowercase       :      Tuple						,				__lowercase       :      Any						,				__lowercase       :      Tuple						,				__lowercase       :      List[str]						,				__lowercase       :      Optional[Any]						,				__lowercase       :      Union[str, Any]						,				__lowercase       :      Union[str, Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						str       =						self.num_choices
        __UpperCAmelCase    :						Union[str, Any]       =						NystromformerForMultipleChoice(config=__lowercase )
        model.to(__lowercase )
        model.eval()
        __UpperCAmelCase    :						Any       =						input_ids.unsqueeze(1 ).expand(-1						,				self.num_choices						,				-1 ).contiguous()
        __UpperCAmelCase    :						Tuple       =						token_type_ids.unsqueeze(1 ).expand(-1						,				self.num_choices						,				-1 ).contiguous()
        __UpperCAmelCase    :						Optional[Any]       =						input_mask.unsqueeze(1 ).expand(-1						,				self.num_choices						,				-1 ).contiguous()
        __UpperCAmelCase    :						Optional[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       A_					(							self       :      str ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[str]       =						self.prepare_config_and_inputs()
        (
            (
            __UpperCAmelCase
        )       ,     (
            __UpperCAmelCase
        )       ,     (
            __UpperCAmelCase
        )       ,     (
            __UpperCAmelCase
        )       ,     (
            __UpperCAmelCase
        )       ,     (
            __UpperCAmelCase
        )       ,     (
            __UpperCAmelCase
        )       ,     
        )    :						Union[str, 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 							snake_case							(			__UpperCAmelCase      ,  __UpperCAmelCase      ,  unittest.TestCase ):

      '''simple docstring'''



      _A     :       Any         =       (
          (
              NystromformerModel,
              NystromformerForMaskedLM,
              NystromformerForMultipleChoice,
              NystromformerForQuestionAnswering,
              NystromformerForSequenceClassification,
              NystromformerForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _A     :       List[str]         =       (
          {
              'feature-extraction': NystromformerModel,
              'fill-mask': NystromformerForMaskedLM,
              'question-answering': NystromformerForQuestionAnswering,
              'text-classification': NystromformerForSequenceClassification,
              'token-classification': NystromformerForTokenClassification,
              'zero-shot': NystromformerForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _A     :       str         =       False
      _A     :       List[Any]         =       False

      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[Any]       =						NystromformerModelTester(self )
        __UpperCAmelCase    :						Tuple       =						ConfigTester(self						,				config_class=__lowercase						,				hidden_size=37 )

      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




        self.config_tester.run_common_tests()

      def       A_					(							self       :      Union[str, Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Dict       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase )

      def       A_					(							self       :      Any ):


        '''simple docstring'''




        __UpperCAmelCase    :						Tuple       =						self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          __UpperCAmelCase    :						Dict       =						type
          self.model_tester.create_and_check_model(*__lowercase )

      def       A_					(							self       :      str ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowercase )

      def       A_					(							self       :      str ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowercase )

      def       A_					(							self       :      int ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowercase )

      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        __UpperCAmelCase    :						str       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowercase )

      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        __UpperCAmelCase    :						Tuple       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowercase )







      @slow
      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          __UpperCAmelCase    :						str       =						NystromformerModel.from_pretrained(__lowercase )
          self.assertIsNotNone(__lowercase )




@require_torch
class 							snake_case							(			unittest.TestCase ):

      '''simple docstring'''



      @slow
      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Union[str, Any]       =						NystromformerModel.from_pretrained('''uw-madison/nystromformer-512''' )
        __UpperCAmelCase    :						Optional[Any]       =						torch.tensor([[0, 1, 2, 3, 4, 5]] )

        with torch.no_grad():
          __UpperCAmelCase    :						str       =						model(__lowercase )[0]

        __UpperCAmelCase    :						Union[str, Any]       =						torch.Size((1, 6, 768) )
        self.assertEqual(output.shape						,				__lowercase )

        __UpperCAmelCase    :						Optional[Any]       =						torch.tensor(
            [[[-0.4_5_3_2, -0.0_9_3_6, 0.5_1_3_7], [-0.2_6_7_6, 0.0_6_2_8, 0.6_1_8_6], [-0.3_6_2_9, -0.1_7_2_6, 0.4_7_1_6]]] )

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







      @slow
      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						'''the [MASK] of Belgium is Brussels'''

        __UpperCAmelCase    :						List[Any]       =						AutoTokenizer.from_pretrained('''uw-madison/nystromformer-512''' )
        __UpperCAmelCase    :						Dict       =						NystromformerForMaskedLM.from_pretrained('''uw-madison/nystromformer-512''' )

        __UpperCAmelCase    :						str       =						tokenizer(__lowercase						,				return_tensors='''pt''' )

        with torch.no_grad():
          __UpperCAmelCase    :						Optional[int]       =						model(encoding.input_ids ).logits

        __UpperCAmelCase    :						Any       =						token_logits[:, 2, :].argmax(-1 )[0]

        self.assertEqual(tokenizer.decode(__lowercase )						,				'''capital''' )






"""simple docstring"""





import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu




class 							snake_case							(			nn.Module ):

      '''simple docstring'''



      def __init__(							self       :      List[Any] ):


        '''simple docstring'''




        super().__init__()
        __UpperCAmelCase    :						Dict       =						nn.Linear(3						,				4 )
        __UpperCAmelCase    :						Union[str, Any]       =						nn.BatchNormad(4 )
        __UpperCAmelCase    :						List[str]       =						nn.Linear(4						,				5 )







      def       A_					(							self       :      Any						,				__lowercase       :      Any ):


        '''simple docstring'''




        return self.lineara(self.batchnorm(self.lineara(__lowercase ) ) )




class 							snake_case							(			__UpperCAmelCase ):

      '''simple docstring'''



      def       A_					(							self       :      Union[str, Any]						,				__lowercase       :      Optional[int]						,				*__lowercase       :      str						,				**__lowercase       :      Optional[int] ):


        '''simple docstring'''




        return (args[0] + 1,) + args[1:], kwargs




class 							snake_case							(			__UpperCAmelCase ):

      '''simple docstring'''



      def       A_					(							self       :      Any						,				__lowercase       :      Tuple						,				__lowercase       :      Any ):


        '''simple docstring'''




        return output + 1




class 							snake_case							(			unittest.TestCase ):

      '''simple docstring'''



      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Dict       =						ModelForTest()
        __UpperCAmelCase    :						Optional[int]       =						ModelHook()

        add_hook_to_module(__lowercase						,				__lowercase )
        self.assertEqual(test_model._hf_hook						,				__lowercase )
        self.assertTrue(hasattr(__lowercase						,				'''_old_forward''' ) )

        # Check adding the hook did not change the name or the signature
        self.assertEqual(test_model.forward.__name__						,				'''forward''' )
        self.assertListEqual(list(inspect.signature(test_model.forward ).parameters )						,				['''x'''] )

        remove_hook_from_module(__lowercase )
        self.assertFalse(hasattr(__lowercase						,				'''_hf_hook''' ) )
        self.assertFalse(hasattr(__lowercase						,				'''_old_forward''' ) )

      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        __UpperCAmelCase    :						str       =						ModelForTest()
        __UpperCAmelCase    :						Tuple       =						ModelHook()

        add_hook_to_module(__lowercase						,				__lowercase )
        add_hook_to_module(__lowercase						,				__lowercase						,				append=__lowercase )

        self.assertEqual(isinstance(test_model._hf_hook						,				__lowercase )						,				__lowercase )
        self.assertEqual(len(test_model._hf_hook.hooks )						,				2 )
        self.assertTrue(hasattr(__lowercase						,				'''_old_forward''' ) )

        # Check adding the hook did not change the name or the signature
        self.assertEqual(test_model.forward.__name__						,				'''forward''' )
        self.assertListEqual(list(inspect.signature(test_model.forward ).parameters )						,				['''x'''] )

        remove_hook_from_module(__lowercase )
        self.assertFalse(hasattr(__lowercase						,				'''_hf_hook''' ) )
        self.assertFalse(hasattr(__lowercase						,				'''_old_forward''' ) )

      def       A_					(							self       :      Any ):


        '''simple docstring'''




        __UpperCAmelCase    :						Dict       =						ModelForTest()
        __UpperCAmelCase    :						Tuple       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						Optional[int]       =						test_model(x + 1 )
        __UpperCAmelCase    :						Optional[Any]       =						test_model(x + 2 )

        __UpperCAmelCase    :						Optional[int]       =						PreForwardHook()
        add_hook_to_module(__lowercase						,				__lowercase )
        __UpperCAmelCase    :						Union[str, Any]       =						test_model(__lowercase )
        self.assertTrue(torch.allclose(__lowercase						,				__lowercase						,				atol=1e-5 ) )

        # Attaching a hook to a model when it already has one replaces, does not chain
        __UpperCAmelCase    :						int       =						PreForwardHook()
        add_hook_to_module(__lowercase						,				__lowercase )
        __UpperCAmelCase    :						Any       =						test_model(__lowercase )
        self.assertTrue(torch.allclose(__lowercase						,				__lowercase						,				atol=1e-5 ) )

        # You need to use the sequential hook to chain two or more hooks
        __UpperCAmelCase    :						Any       =						SequentialHook(PreForwardHook()						,				PreForwardHook() )
        add_hook_to_module(__lowercase						,				__lowercase )

        __UpperCAmelCase    :						Tuple       =						test_model(__lowercase )
        assert torch.allclose(__lowercase						,				__lowercase						,				atol=1e-5 )

      def       A_					(							self       :      Any ):


        '''simple docstring'''




        __UpperCAmelCase    :						int       =						ModelForTest()
        __UpperCAmelCase    :						List[Any]       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						Tuple       =						test_model(__lowercase )

        __UpperCAmelCase    :						int       =						PostForwardHook()
        add_hook_to_module(__lowercase						,				__lowercase )
        __UpperCAmelCase    :						Optional[int]       =						test_model(__lowercase )
        self.assertTrue(torch.allclose(__lowercase						,				output + 1						,				atol=1e-5 ) )

        # Attaching a hook to a model when it already has one replaces, does not chain
        __UpperCAmelCase    :						str       =						PostForwardHook()
        add_hook_to_module(__lowercase						,				__lowercase )
        __UpperCAmelCase    :						List[str]       =						test_model(__lowercase )
        self.assertTrue(torch.allclose(__lowercase						,				output + 1						,				atol=1e-5 ) )

        # You need to use the sequential hook to chain two or more hooks
        __UpperCAmelCase    :						Optional[int]       =						SequentialHook(PostForwardHook()						,				PostForwardHook() )
        add_hook_to_module(__lowercase						,				__lowercase )

        __UpperCAmelCase    :						Dict       =						test_model(__lowercase )
        assert torch.allclose(__lowercase						,				output + 2						,				atol=1e-5 )

      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[str]       =						ModelForTest()
        __UpperCAmelCase    :						Union[str, Any]       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						str       =						test_model(__lowercase )

        __UpperCAmelCase    :						Union[str, Any]       =						PostForwardHook()
        add_hook_to_module(__lowercase						,				__lowercase )
        __UpperCAmelCase    :						Optional[int]       =						test_model(__lowercase )
        self.assertTrue(torch.allclose(__lowercase						,				output + 1 ) )
        self.assertTrue(outputa.requires_grad )

        __UpperCAmelCase    :						Optional[Any]       =						True
        __UpperCAmelCase    :						int       =						test_model(__lowercase )
        self.assertFalse(outputa.requires_grad )

      @require_multi_gpu
      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						ModelForTest()
        # Everything is on CPU
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )

        # This will move each submodule on different devices
        add_hook_to_module(model.lineara						,				AlignDevicesHook(execution_device=0 ) )
        add_hook_to_module(model.batchnorm						,				AlignDevicesHook(execution_device=0 ) )
        add_hook_to_module(model.lineara						,				AlignDevicesHook(execution_device=1 ) )

        self.assertEqual(model.lineara.weight.device						,				torch.device(0 ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device(0 ) )
        self.assertEqual(model.batchnorm.running_mean.device						,				torch.device(0 ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device(1 ) )

        # We can still make a forward pass. The input does not need to be on any particular device
        __UpperCAmelCase    :						Dict       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						Any       =						model(__lowercase )
        self.assertEqual(output.device						,				torch.device(1 ) )

        # We can add a general hook to put back output on same device as input.
        add_hook_to_module(__lowercase						,				AlignDevicesHook(io_same_device=__lowercase ) )
        __UpperCAmelCase    :						List[Any]       =						torch.randn(2						,				3 ).to(0 )
        __UpperCAmelCase    :						int       =						model(__lowercase )
        self.assertEqual(output.device						,				torch.device(0 ) )

      def       A_					(							self       :      Tuple ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						ModelForTest()

        # Everything is on CPU
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )

        # This will move each submodule on different devices
        __UpperCAmelCase    :						Tuple       =						{'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

        add_hook_to_module(model.lineara						,				AlignDevicesHook(**__lowercase ) )
        add_hook_to_module(model.batchnorm						,				AlignDevicesHook(**__lowercase ) )
        add_hook_to_module(model.lineara						,				AlignDevicesHook(**__lowercase ) )

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        # Buffers are not included in the offload by default, so are on the execution device
        __UpperCAmelCase    :						Optional[int]       =						torch.device(hook_kwargs['''execution_device'''] )
        self.assertEqual(model.batchnorm.running_mean.device						,				__lowercase )

        __UpperCAmelCase    :						int       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						Optional[int]       =						model(__lowercase )
        self.assertEqual(output.device						,				__lowercase )

        # Removing hooks loads back the weights in the model.
        remove_hook_from_module(model.lineara )
        remove_hook_from_module(model.batchnorm )
        remove_hook_from_module(model.lineara )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )

        # Now test with buffers included in the offload
        __UpperCAmelCase    :						str       =						{
            '''execution_device''': 0 if torch.cuda.is_available() else '''cpu''',
            '''offload''': True,
            '''offload_buffers''': True,
        }

        add_hook_to_module(model.lineara						,				AlignDevicesHook(**__lowercase ) )
        add_hook_to_module(model.batchnorm						,				AlignDevicesHook(**__lowercase ) )
        add_hook_to_module(model.lineara						,				AlignDevicesHook(**__lowercase ) )

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.running_mean.device						,				torch.device('''meta''' ) )

        __UpperCAmelCase    :						Optional[Any]       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						List[Any]       =						model(__lowercase )
        self.assertEqual(output.device						,				__lowercase )

        # Removing hooks loads back the weights in the model.
        remove_hook_from_module(model.lineara )
        remove_hook_from_module(model.batchnorm )
        remove_hook_from_module(model.lineara )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )

      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[str]       =						ModelForTest()

        # Everything is on CPU
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )

        # This will move each submodule on different devices
        __UpperCAmelCase    :						Optional[Any]       =						0 if torch.cuda.is_available() else '''cpu'''
        attach_align_device_hook(__lowercase						,				execution_device=__lowercase						,				offload=__lowercase )

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        # Buffers are not included in the offload by default, so are on the execution device
        __UpperCAmelCase    :						Dict       =						torch.device(__lowercase )
        self.assertEqual(model.batchnorm.running_mean.device						,				__lowercase )

        __UpperCAmelCase    :						Optional[int]       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						Dict       =						model(__lowercase )
        self.assertEqual(output.device						,				__lowercase )

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(__lowercase )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )

        # Now test with buffers included in the offload
        attach_align_device_hook(__lowercase						,				execution_device=__lowercase						,				offload=__lowercase						,				offload_buffers=__lowercase )

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.running_mean.device						,				torch.device('''meta''' ) )

        __UpperCAmelCase    :						Dict       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						str       =						model(__lowercase )
        self.assertEqual(output.device						,				__lowercase )

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(__lowercase )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )







      def       A_					(							self       :      Any ):


        '''simple docstring'''




        __UpperCAmelCase    :						Dict       =						ModelForTest()

        # Everything is on CPU
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )

        # This will move each submodule on different devices
        __UpperCAmelCase    :						str       =						0 if torch.cuda.is_available() else '''cpu'''
        attach_align_device_hook(
            __lowercase						,				execution_device=__lowercase						,				offload=__lowercase						,				weights_map=model.state_dict() )

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        # Buffers are not included in the offload by default, so are on the execution device
        __UpperCAmelCase    :						Optional[Any]       =						torch.device(__lowercase )
        self.assertEqual(model.batchnorm.running_mean.device						,				__lowercase )

        __UpperCAmelCase    :						Any       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						Dict       =						model(__lowercase )
        self.assertEqual(output.device						,				__lowercase )

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(__lowercase )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )

        # Now test with buffers included in the offload
        attach_align_device_hook(
            __lowercase						,				execution_device=__lowercase						,				offload=__lowercase						,				weights_map=model.state_dict()						,				offload_buffers=__lowercase						,				)

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''meta''' ) )
        self.assertEqual(model.batchnorm.running_mean.device						,				torch.device('''meta''' ) )

        __UpperCAmelCase    :						List[str]       =						torch.randn(2						,				3 )
        __UpperCAmelCase    :						Optional[int]       =						model(__lowercase )
        self.assertEqual(output.device						,				__lowercase )

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(__lowercase )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.batchnorm.weight.device						,				torch.device('''cpu''' ) )
        self.assertEqual(model.lineara.weight.device						,				torch.device('''cpu''' ) )