






"""simple docstring"""




import unittest

from transformers import DebertaVaConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
 import torch

 from transformers import (
     DebertaVaForMaskedLM,
     DebertaVaForMultipleChoice,
     DebertaVaForQuestionAnswering,
     DebertaVaForSequenceClassification,
     DebertaVaForTokenClassification,
     DebertaVaModel,
 )
 from transformers.models.deberta_va.modeling_deberta_va import DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST






class        UpperCamelCase			(	snake_case_      ):


 def __init__(     self		:						List[str]							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						Tuple=13							,						UpperCAmelCase__		:						Any=7							,						UpperCAmelCase__		:						List[Any]=True							,						UpperCAmelCase__		:						Optional[Any]=True							,						UpperCAmelCase__		:						Dict=True							,						UpperCAmelCase__		:						Dict=True							,						UpperCAmelCase__		:						int=99							,						UpperCAmelCase__		:						Optional[Any]=32							,						UpperCAmelCase__		:						int=5							,						UpperCAmelCase__		:						int=4							,						UpperCAmelCase__		:						Dict=37							,						UpperCAmelCase__		:						Optional[Any]="gelu"							,						UpperCAmelCase__		:						int=0.1							,						UpperCAmelCase__		:						Union[str, Any]=0.1							,						UpperCAmelCase__		:						List[str]=512							,						UpperCAmelCase__		:						Tuple=16							,						UpperCAmelCase__		:						Optional[Any]=2							,						UpperCAmelCase__		:						Optional[Any]=0.0_2							,						UpperCAmelCase__		:						Any=False							,						UpperCAmelCase__		:						Tuple=True							,						UpperCAmelCase__		:						Any="None"							,						UpperCAmelCase__		:						Optional[int]=3							,						UpperCAmelCase__		:						List[Any]=4							,						UpperCAmelCase__		:						Optional[int]=None							,						)					->       List[str]:
   _a  :       Any             =	parent
   _a  :       Union[str, Any]             =	batch_size
   _a  :       List[str]             =	seq_length
   _a  :       List[str]             =	is_training
   _a  :       Optional[Any]             =	use_input_mask
   _a  :       Any             =	use_token_type_ids
   _a  :       str             =	use_labels
   _a  :       Tuple             =	vocab_size
   _a  :       List[str]             =	hidden_size
   _a  :       List[str]             =	num_hidden_layers
   _a  :       List[Any]             =	num_attention_heads
   _a  :       Tuple             =	intermediate_size
   _a  :       Dict             =	hidden_act
   _a  :       Tuple             =	hidden_dropout_prob
   _a  :       Tuple             =	attention_probs_dropout_prob
   _a  :       Dict             =	max_position_embeddings
   _a  :       Tuple             =	type_vocab_size
   _a  :       Optional[Any]             =	type_sequence_label_size
   _a  :       int             =	initializer_range
   _a  :       Optional[int]             =	num_labels
   _a  :       List[str]             =	num_choices
   _a  :       Any             =	relative_attention
   _a  :       Optional[Any]             =	position_biased_input
   _a  :       List[Any]             =	pos_att_type
   _a  :       str             =	scope


 def _lowercase	(     self		:						int							)					->       Any:
   _a  :       List[Any]             =	ids_tensor([self.batch_size, self.seq_length]							,						self.vocab_size							)

   _a  :       Optional[int]             =	None
   if self.use_input_mask:
     _a  :       str             =	ids_tensor([self.batch_size, self.seq_length]							,						vocab_size=2							)

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

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

   _a  :       Optional[int]             =	self.get_config()

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


 def _lowercase	(     self		:						List[str]							)					->       Optional[Any]:
   return DebertaVaConfig(
       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							,						initializer_range=self.initializer_range							,						relative_attention=self.relative_attention							,						position_biased_input=self.position_biased_input							,						pos_att_type=self.pos_att_type							,						)


 def _lowercase	(     self		:						List[Any]							,						UpperCAmelCase__		:						Union[str, Any]							)					->       Union[str, Any]:
   self.parent.assertListEqual(list(result.loss.size()							)							,						[]							)


 def _lowercase	(     self		:						Optional[Any]							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						Dict							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						Union[str, Any]							)					->       Tuple:
   _a  :       Any             =	DebertaVaModel(config=UpperCAmelCase__							)
   model.to(UpperCAmelCase__							)
   model.eval()
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							,						attention_mask=UpperCAmelCase__							,						token_type_ids=UpperCAmelCase__							)[0]
   _a  :       List[str]             =	model(UpperCAmelCase__							,						token_type_ids=UpperCAmelCase__							)[0]
   _a  :       List[str]             =	model(UpperCAmelCase__							)[0]

   self.parent.assertListEqual(list(sequence_output.size()							)							,						[self.batch_size, self.seq_length, self.hidden_size]							)


 def _lowercase	(     self		:						Dict							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Union[str, Any]							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						Optional[int]							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Any							)					->       int:
   _a  :       Optional[Any]             =	DebertaVaForMaskedLM(config=UpperCAmelCase__							)
   model.to(UpperCAmelCase__							)
   model.eval()
   _a  :       List[str]             =	model(UpperCAmelCase__							,						attention_mask=UpperCAmelCase__							,						token_type_ids=UpperCAmelCase__							,						labels=UpperCAmelCase__							)

   self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.seq_length, self.vocab_size)							)


 def _lowercase	(     self		:						List[Any]							,						UpperCAmelCase__		:						Dict							,						UpperCAmelCase__		:						Union[str, Any]							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						Optional[int]							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						str							)					->       str:
   _a  :       List[str]             =	self.num_labels
   _a  :       List[Any]             =	DebertaVaForSequenceClassification(UpperCAmelCase__							)
   model.to(UpperCAmelCase__							)
   model.eval()
   _a  :       Dict             =	model(UpperCAmelCase__							,						attention_mask=UpperCAmelCase__							,						token_type_ids=UpperCAmelCase__							,						labels=UpperCAmelCase__							)
   self.parent.assertListEqual(list(result.logits.size()							)							,						[self.batch_size, self.num_labels]							)
   self.check_loss_output(UpperCAmelCase__							)


 def _lowercase	(     self		:						List[str]							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						int							)					->       Any:
   _a  :       Optional[Any]             =	self.num_labels
   _a  :       Any             =	DebertaVaForTokenClassification(config=UpperCAmelCase__							)
   model.to(UpperCAmelCase__							)
   model.eval()
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							,						attention_mask=UpperCAmelCase__							,						token_type_ids=UpperCAmelCase__							,						labels=UpperCAmelCase__							)
   self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.seq_length, self.num_labels)							)


 def _lowercase	(     self		:						str							,						UpperCAmelCase__		:						Dict							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Optional[int]							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						Dict							)					->       Optional[Any]:
   _a  :       int             =	DebertaVaForQuestionAnswering(config=UpperCAmelCase__							)
   model.to(UpperCAmelCase__							)
   model.eval()
   _a  :       str             =	model(
       UpperCAmelCase__							,						attention_mask=UpperCAmelCase__							,						token_type_ids=UpperCAmelCase__							,						start_positions=UpperCAmelCase__							,						end_positions=UpperCAmelCase__							,						)
   self.parent.assertEqual(result.start_logits.shape							,						(self.batch_size, self.seq_length)							)
   self.parent.assertEqual(result.end_logits.shape							,						(self.batch_size, self.seq_length)							)


 def _lowercase	(     self		:						List[str]							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						int							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						List[str]							)					->       Union[str, Any]:
   _a  :       List[Any]             =	DebertaVaForMultipleChoice(config=UpperCAmelCase__							)
   model.to(UpperCAmelCase__							)
   model.eval()
   _a  :       Optional[int]             =	input_ids.unsqueeze(1							).expand(-1							,						self.num_choices							,						-1							).contiguous()
   _a  :       Optional[Any]             =	token_type_ids.unsqueeze(1							).expand(-1							,						self.num_choices							,						-1							).contiguous()
   _a  :       Any             =	input_mask.unsqueeze(1							).expand(-1							,						self.num_choices							,						-1							).contiguous()
   _a  :       Any             =	model(
       UpperCAmelCase__							,						attention_mask=UpperCAmelCase__							,						token_type_ids=UpperCAmelCase__							,						labels=UpperCAmelCase__							,						)
   self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.num_choices)							)







 def _lowercase	(     self		:						int							)					->       Dict:
   _a  :       int             =	self.prepare_config_and_inputs()
   (
       (
       _a
   )      ,  (
       _a
   )      ,  (
       _a
   )      ,  (
       _a
   )      ,  (
       _a
   )      ,  (
       _a
   )      ,  (
       _a
   )      ,  
   )  :       Any             =	config_and_inputs
   _a  :       Tuple             =	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
   return config, inputs_dict







@require_torch
class        UpperCamelCase			(	snake_case_		,     snake_case_		,     unittest.TestCase      ):
 UpperCamelCase					:       Optional[int]			=					(
     (
         DebertaVaModel,
         DebertaVaForMaskedLM,
         DebertaVaForSequenceClassification,
         DebertaVaForTokenClassification,
         DebertaVaForQuestionAnswering,
         DebertaVaForMultipleChoice,
     )
     if is_torch_available()
     else ()
 )
 UpperCamelCase					:       Optional[int]			=					(
     {
         '''feature-extraction''': DebertaVaModel,
         '''fill-mask''': DebertaVaForMaskedLM,
         '''question-answering''': DebertaVaForQuestionAnswering,
         '''text-classification''': DebertaVaForSequenceClassification,
         '''token-classification''': DebertaVaForTokenClassification,
         '''zero-shot''': DebertaVaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )

 UpperCamelCase					:       Tuple			=					True
 UpperCamelCase					:       Any			=					False
 UpperCamelCase					:       Tuple			=					False
 UpperCamelCase					:       int			=					False
 UpperCamelCase					:       str			=					False


 def _lowercase	(     self		:						str							)					->       List[Any]:
   _a  :       List[str]             =	DebertaVaModelTester(self							)
   _a  :       int             =	ConfigTester(self							,						config_class=UpperCAmelCase__							,						hidden_size=37							)


 def _lowercase	(     self		:						str							)					->       Dict:
   self.config_tester.run_common_tests()


 def _lowercase	(     self		:						Dict							)					->       List[str]:
   _a  :       Any             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_deberta_model(*UpperCAmelCase__							)


 def _lowercase	(     self		:						Tuple							)					->       str:
   _a  :       List[str]             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_deberta_for_sequence_classification(*UpperCAmelCase__							)


 def _lowercase	(     self		:						Union[str, Any]							)					->       List[Any]:
   _a  :       str             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_deberta_for_masked_lm(*UpperCAmelCase__							)


 def _lowercase	(     self		:						Optional[int]							)					->       int:
   _a  :       List[str]             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_deberta_for_question_answering(*UpperCAmelCase__							)


 def _lowercase	(     self		:						str							)					->       Union[str, Any]:
   _a  :       Any             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_deberta_for_token_classification(*UpperCAmelCase__							)


 def _lowercase	(     self		:						Any							)					->       Tuple:
   _a  :       Tuple             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_deberta_for_multiple_choice(*UpperCAmelCase__							)







 @slow
 def _lowercase	(     self		:						Tuple							)					->       Optional[int]:
   for model_name in DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _a  :       Dict             =	DebertaVaModel.from_pretrained(UpperCAmelCase__							)
     self.assertIsNotNone(UpperCAmelCase__							)







@require_torch
@require_sentencepiece
@require_tokenizers
class        UpperCamelCase			(	unittest.TestCase      ):


 @unittest.skip(reason="""Model not available yet"""							)
 def _lowercase	(     self		:						Union[str, Any]							)					->       Optional[Any]:
   pass







 @slow
 def _lowercase	(     self		:						List[Any]							)					->       Union[str, Any]:
   _a  :       str             =	DebertaVaModel.from_pretrained("""microsoft/deberta-v2-xlarge"""							)

   _a  :       str             =	torch.tensor([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]							)
   _a  :       str             =	torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
   with torch.no_grad():
     _a  :       List[Any]             =	model(UpperCAmelCase__							,						attention_mask=UpperCAmelCase__							)[0]
   # compare the actual values for a slice.
   _a  :       List[Any]             =	torch.tensor(
       [[[0.2_3_5_6, 0.1_9_4_8, 0.0_3_6_9], [-0.1_0_6_3, 0.3_5_8_6, -0.5_1_5_2], [-0.6_3_9_9, -0.0_2_5_9, -0.2_5_2_5]]]							)
   self.assertTrue(torch.allclose(output[:, 1:4, 1:4]							,						UpperCAmelCase__							,						atol=1E-4							)							,						f"""{output[:, 1:4, 1:4]}"""							)









"""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        UpperCamelCase			(	nn.Module      ):


 def __init__(     self		:						Union[str, Any]							)					->       int:
   super().__init__()
   _a  :       Optional[Any]             =	nn.Linear(3							,						4							)
   _a  :       Tuple             =	nn.BatchNormad(4							)
   _a  :       Dict             =	nn.Linear(4							,						5							)







 def _lowercase	(     self		:						Optional[int]							,						UpperCAmelCase__		:						List[str]							)					->       int:
   return self.lineara(self.batchnorm(self.lineara(UpperCAmelCase__							)							)							)







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Any							,						UpperCAmelCase__		:						Any							,						*UpperCAmelCase__		:						List[str]							,						**UpperCAmelCase__		:						Optional[int]							)					->       Optional[int]:
   return (args[0] + 1,) + args[1:], kwargs







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Optional[Any]							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Optional[Any]							)					->       List[str]:
   return output + 1







class        UpperCamelCase			(	unittest.TestCase      ):


 def _lowercase	(     self		:						Dict							)					->       str:
   _a  :       List[Any]             =	ModelForTest()
   _a  :       str             =	ModelHook()

   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   self.assertEqual(test_model._hf_hook							,						UpperCAmelCase__							)
   self.assertTrue(hasattr(UpperCAmelCase__							,						"""_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(UpperCAmelCase__							)
   self.assertFalse(hasattr(UpperCAmelCase__							,						"""_hf_hook"""							)							)
   self.assertFalse(hasattr(UpperCAmelCase__							,						"""_old_forward"""							)							)


 def _lowercase	(     self		:						Optional[int]							)					->       Optional[int]:
   _a  :       Dict             =	ModelForTest()
   _a  :       Dict             =	ModelHook()

   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							,						append=UpperCAmelCase__							)

   self.assertEqual(isinstance(test_model._hf_hook							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(len(test_model._hf_hook.hooks							)							,						2							)
   self.assertTrue(hasattr(UpperCAmelCase__							,						"""_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(UpperCAmelCase__							)
   self.assertFalse(hasattr(UpperCAmelCase__							,						"""_hf_hook"""							)							)
   self.assertFalse(hasattr(UpperCAmelCase__							,						"""_old_forward"""							)							)


 def _lowercase	(     self		:						Dict							)					->       int:
   _a  :       str             =	ModelForTest()
   _a  :       List[Any]             =	torch.randn(2							,						3							)
   _a  :       Optional[Any]             =	test_model(x + 1							)
   _a  :       str             =	test_model(x + 2							)

   _a  :       Union[str, Any]             =	PreForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       Tuple             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						UpperCAmelCase__							,						atol=1E-5							)							)

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

   # You need to use the sequential hook to chain two or more hooks
   _a  :       int             =	SequentialHook(PreForwardHook()							,						PreForwardHook()							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Tuple             =	test_model(UpperCAmelCase__							)
   assert torch.allclose(UpperCAmelCase__							,						UpperCAmelCase__							,						atol=1E-5							)


 def _lowercase	(     self		:						Tuple							)					->       int:
   _a  :       Tuple             =	ModelForTest()
   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       Optional[int]             =	test_model(UpperCAmelCase__							)

   _a  :       int             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       List[str]             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						output + 1							,						atol=1E-5							)							)

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

   # You need to use the sequential hook to chain two or more hooks
   _a  :       Any             =	SequentialHook(PostForwardHook()							,						PostForwardHook()							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Optional[int]             =	test_model(UpperCAmelCase__							)
   assert torch.allclose(UpperCAmelCase__							,						output + 2							,						atol=1E-5							)


 def _lowercase	(     self		:						Dict							)					->       Optional[Any]:
   _a  :       Any             =	ModelForTest()
   _a  :       List[Any]             =	torch.randn(2							,						3							)
   _a  :       Dict             =	test_model(UpperCAmelCase__							)

   _a  :       Any             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       List[str]             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						output + 1							)							)
   self.assertTrue(outputa.requires_grad							)

   _a  :       Any             =	True
   _a  :       Union[str, Any]             =	test_model(UpperCAmelCase__							)
   self.assertFalse(outputa.requires_grad							)


 @require_multi_gpu
 def _lowercase	(     self		:						Optional[Any]							)					->       str:
   _a  :       List[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
   _a  :       Optional[int]             =	torch.randn(2							,						3							)
   _a  :       Any             =	model(UpperCAmelCase__							)
   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(UpperCAmelCase__							,						AlignDevicesHook(io_same_device=UpperCAmelCase__							)							)
   _a  :       str             =	torch.randn(2							,						3							).to(0							)
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						torch.device(0							)							)


 def _lowercase	(     self		:						str							)					->       Union[str, Any]:
   _a  :       int             =	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
   _a  :       List[Any]             =	{"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

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

   # 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
   _a  :       Dict             =	torch.device(hook_kwargs["""execution_device"""]							)
   self.assertEqual(model.batchnorm.running_mean.device							,						UpperCAmelCase__							)

   _a  :       int             =	torch.randn(2							,						3							)
   _a  :       str             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # 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
   _a  :       List[str]             =	{
       """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
       """offload""": True,
       """offload_buffers""": True,
   }

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

   # 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"""							)							)

   _a  :       Tuple             =	torch.randn(2							,						3							)
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # 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 _lowercase	(     self		:						Tuple							)					->       List[str]:
   _a  :       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
   _a  :       Union[str, Any]             =	0 if torch.cuda.is_available() else """cpu"""
   attach_align_device_hook(UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							)

   # 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
   _a  :       Dict             =	torch.device(UpperCAmelCase__							)
   self.assertEqual(model.batchnorm.running_mean.device							,						UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       List[Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   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(UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							,						offload_buffers=UpperCAmelCase__							)

   # 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"""							)							)

   _a  :       List[str]             =	torch.randn(2							,						3							)
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   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 _lowercase	(     self		:						Dict							)					->       str:
   _a  :       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
   _a  :       str             =	0 if torch.cuda.is_available() else """cpu"""
   attach_align_device_hook(
       UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							,						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
   _a  :       Union[str, Any]             =	torch.device(UpperCAmelCase__							)
   self.assertEqual(model.batchnorm.running_mean.device							,						UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       int             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   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(
       UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							,						weights_map=model.state_dict()							,						offload_buffers=UpperCAmelCase__							,						)

   # 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"""							)							)

   _a  :       Any             =	torch.randn(2							,						3							)
   _a  :       int             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   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"""							)							)

