





'''simple docstring'''






import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
		import torch
		from torch import nn

		from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
		from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import ViTImageProcessor







class 			lowerCamelCase					:


					def __init__(					self	:  Union[str, Any]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  str=13       ,     __snake_case	:  Union[str, Any]=30       ,     __snake_case	:  Union[str, Any]=2       ,     __snake_case	:  Dict=3       ,     __snake_case	:  Optional[Any]=True       ,     __snake_case	:  Optional[int]=True       ,     __snake_case	:  Optional[int]=32       ,     __snake_case	:  Optional[int]=5       ,     __snake_case	:  Any=4       ,     __snake_case	:  int=37       ,     __snake_case	:  int="gelu"       ,     __snake_case	:  Union[str, Any]=0.1       ,     __snake_case	:  Optional[int]=0.1       ,     __snake_case	:  int=10       ,     __snake_case	:  Any=0.02       ,     __snake_case	:  List[str]=None       ,     __snake_case	:  Tuple=2       ,     ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		parent
											_snake_case:   Tuple  					=		batch_size
											_snake_case:   str  					=		image_size
											_snake_case:   int  					=		patch_size
											_snake_case:   Union[str, Any]  					=		num_channels
											_snake_case:   Dict  					=		is_training
											_snake_case:   Optional[Any]  					=		use_labels
											_snake_case:   Optional[Any]  					=		hidden_size
											_snake_case:   Tuple  					=		num_hidden_layers
											_snake_case:   List[Any]  					=		num_attention_heads
											_snake_case:   Union[str, Any]  					=		intermediate_size
											_snake_case:   List[str]  					=		hidden_act
											_snake_case:   Tuple  					=		hidden_dropout_prob
											_snake_case:   List[Any]  					=		attention_probs_dropout_prob
											_snake_case:   str  					=		type_sequence_label_size
											_snake_case:   Any  					=		initializer_range
											_snake_case:   str  					=		scope
											_snake_case:   Union[str, Any]  					=		encoder_stride

											# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
											_snake_case:   Tuple  					=		(image_size // patch_size) ** 2
											_snake_case:   List[str]  					=		num_patches + 1


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											_snake_case:   List[str]  					=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

											_snake_case:   List[str]  					=		None
											if self.use_labels:
																	_snake_case:   Any  					=		ids_tensor([self.batch_size]       ,     self.type_sequence_label_size       )

											_snake_case:   Union[str, Any]  					=		self.get_config()

											return config, pixel_values, labels


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											return ViTConfig(
											    image_size=self.image_size       ,     patch_size=self.patch_size       ,     num_channels=self.num_channels       ,     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       ,     is_decoder=__snake_case       ,     initializer_range=self.initializer_range       ,     encoder_stride=self.encoder_stride       ,     )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ,     __snake_case	:  int       ,     __snake_case	:  Optional[int]       ,     __snake_case	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		ViTModel(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											_snake_case:   Tuple  					=		model(__snake_case       )
											self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size)       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  int       ):





											'''simple docstring'''





											_snake_case:   int  					=		ViTForMaskedImageModeling(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											_snake_case:   Dict  					=		model(__snake_case       )
											self.parent.assertEqual(
											    result.reconstruction.shape       ,     (self.batch_size, self.num_channels, self.image_size, self.image_size)       )

											# test greyscale images
											_snake_case:   List[str]  					=		1
											_snake_case:   Tuple  					=		ViTForMaskedImageModeling(__snake_case       )
											model.to(__snake_case       )
											model.eval()

											_snake_case:   List[str]  					=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       )
											self.parent.assertEqual(result.reconstruction.shape       ,     (self.batch_size, 1, self.image_size, self.image_size)       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Optional[int]       ,     __snake_case	:  Dict       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		self.type_sequence_label_size
											_snake_case:   Union[str, Any]  					=		ViTForImageClassification(__snake_case       )
											model.to(__snake_case       )
											model.eval()
											_snake_case:   List[Any]  					=		model(__snake_case       ,     labels=__snake_case       )
											self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.type_sequence_label_size)       )

											# test greyscale images
											_snake_case:   Tuple  					=		1
											_snake_case:   Optional[int]  					=		ViTForImageClassification(__snake_case       )
											model.to(__snake_case       )
											model.eval()

											_snake_case:   List[Any]  					=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
											_snake_case:   Any  					=		model(__snake_case       )
											self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.type_sequence_label_size)       )




					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.prepare_config_and_inputs()
											(
											    (
											    _snake_case
											)					,						(
											    _snake_case
											)					,						(
											    _snake_case
											)					,						
											):   int  					=		config_and_inputs
											_snake_case:   Dict  					=		{'pixel_values': pixel_values}
											return config, inputs_dict







@require_torch
class 			lowerCamelCase					(						__UpperCAmelCase     ,							__UpperCAmelCase     ,							unittest.TestCase  ):
					_SCREAMING_SNAKE_CASE							   =    (
					    (
					        ViTModel,
					        ViTForImageClassification,
					        ViTForMaskedImageModeling,
					    )
					    if is_torch_available()
					    else ()
					)
					_SCREAMING_SNAKE_CASE							   =    (
					    {"feature-extraction": ViTModel, "image-classification": ViTForImageClassification}
					    if is_torch_available()
					    else {}
					)
					_SCREAMING_SNAKE_CASE							   =    True

					_SCREAMING_SNAKE_CASE							   =    False
					_SCREAMING_SNAKE_CASE							   =    False
					_SCREAMING_SNAKE_CASE							   =    False


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		ViTModelTester(self       )
											_snake_case:   Union[str, Any]  					=		ConfigTester(self       ,     config_class=__snake_case       ,     has_text_modality=__snake_case       ,     hidden_size=37       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											self.config_tester.run_common_tests()


					@unittest.skip(reason='ViT does not use inputs_embeds'       )
					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											pass


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											_snake_case					,						_snake_case:   Dict  					=		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	_snake_case:   Optional[int]  					=		model_class(__snake_case       )
																	self.assertIsInstance(model.get_input_embeddings()       ,     (nn.Module)       )
																	_snake_case:   Optional[Any]  					=		model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(__snake_case       ,     nn.Linear       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case					,						_snake_case:   List[Any]  					=		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	_snake_case:   int  					=		model_class(__snake_case       )
																	_snake_case:   List[str]  					=		inspect.signature(model.forward       )
																	# signature.parameters is an OrderedDict => so arg_names order is deterministic
																	_snake_case:   List[Any]  					=		[*signature.parameters.keys()]

																	_snake_case:   str  					=		['pixel_values']
																	self.assertListEqual(arg_names[:1]       ,     __snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   List[Any]  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   int  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_image_modeling(*__snake_case       )


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





											'''simple docstring'''





											_snake_case:   str  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*__snake_case       )




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





											'''simple docstring'''





											for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_snake_case:   Any  					=		ViTModel.from_pretrained(__snake_case       )
																	self.assertIsNotNone(__snake_case       )





def       lowercase_      (		)   ->List[Any]:
						_snake_case:   Dict  					=		Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
						return image







@require_torch
@require_vision
class 			lowerCamelCase					(						unittest.TestCase  ):


					@cached_property
					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											return ViTImageProcessor.from_pretrained('google/vit-base-patch16-224'       ) if is_vision_available() else None


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'       ).to(__snake_case       )

											_snake_case:   Dict  					=		self.default_image_processor
											_snake_case:   Optional[Any]  					=		prepare_img()
											_snake_case:   List[str]  					=		image_processor(images=__snake_case       ,     return_tensors='pt'       ).to(__snake_case       )

											# forward pass
											with torch.no_grad():
																	_snake_case:   Optional[int]  					=		model(**__snake_case       )

											# verify the logits
											_snake_case:   Union[str, Any]  					=		torch.Size((1, 10_00)       )
											self.assertEqual(outputs.logits.shape       ,     __snake_case       )

											_snake_case:   Dict  					=		torch.tensor([-0.2_744, 0.8_215, -0.0_836]       ).to(__snake_case       )

											self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,     __snake_case       ,     atol=1e-4       )       )


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											_snake_case:   str  					=		ViTModel.from_pretrained('facebook/dino-vits8'       ).to(__snake_case       )

											_snake_case:   Any  					=		ViTImageProcessor.from_pretrained('facebook/dino-vits8'       ,     size=4_80       )
											_snake_case:   Optional[int]  					=		prepare_img()
											_snake_case:   Dict  					=		image_processor(images=__snake_case       ,     return_tensors='pt'       )
											_snake_case:   Optional[Any]  					=		inputs.pixel_values.to(__snake_case       )

											# forward pass
											with torch.no_grad():
																	_snake_case:   str  					=		model(__snake_case       ,     interpolate_pos_encoding=__snake_case       )

											# verify the logits
											_snake_case:   List[str]  					=		torch.Size((1, 36_01, 3_84)       )
											self.assertEqual(outputs.last_hidden_state.shape       ,     __snake_case       )

											_snake_case:   Any  					=		torch.tensor(
											    [[4.2_340, 4.3_906, -6.6_692], [4.5_463, 1.8_928, -6.7_257], [4.4_429, 0.8_496, -5.8_585]]       ).to(__snake_case       )

											self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]       ,     __snake_case       ,     atol=1e-4       )       )




					@slow
					@require_accelerate
					@require_torch_gpu
					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   List[Any]  					=		ViTModel.from_pretrained('facebook/dino-vits8'       ,     torch_dtype=torch.floataa       ,     device_map='auto'       )
											_snake_case:   Dict  					=		self.default_image_processor

											_snake_case:   Any  					=		prepare_img()
											_snake_case:   str  					=		image_processor(images=__snake_case       ,     return_tensors='pt'       )
											_snake_case:   Any  					=		inputs.pixel_values.to(__snake_case       )

											# forward pass to make sure inference works in fp16
											with torch.no_grad():
																	_snake_case:   int  					=		model(__snake_case       )












'''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 			lowerCamelCase					(						nn.Module  ):


					def __init__(					self	:  Optional[int]       ):





											'''simple docstring'''





											super().__init__()
											_snake_case:   int  					=		nn.Linear(3       ,     4       )
											_snake_case:   Any  					=		nn.BatchNormad(4       )
											_snake_case:   Union[str, Any]  					=		nn.Linear(4       ,     5       )




					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  List[Any]       ):





											'''simple docstring'''





											return self.lineara(self.batchnorm(self.lineara(__snake_case       )       )       )







class 			lowerCamelCase					(						__UpperCAmelCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  Optional[Any]       ,     *__snake_case	:  Optional[Any]       ,     **__snake_case	:  List[Any]       ):





											'''simple docstring'''





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







class 			lowerCamelCase					(						__UpperCAmelCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Tuple       ):





											'''simple docstring'''





											return output + 1







class 			lowerCamelCase					(						unittest.TestCase  ):


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





											'''simple docstring'''





											_snake_case:   str  					=		ModelForTest()
											_snake_case:   List[str]  					=		ModelHook()

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


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		ModelForTest()
											_snake_case:   Union[str, Any]  					=		ModelHook()

											add_hook_to_module(__snake_case       ,     __snake_case       )
											add_hook_to_module(__snake_case       ,     __snake_case       ,     append=__snake_case       )

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


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





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

											_snake_case:   List[str]  					=		PreForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   List[Any]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )       )

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

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

											_snake_case:   int  					=		test_model(__snake_case       )
											assert torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   int  					=		ModelForTest()
											_snake_case:   str  					=		torch.randn(2       ,     3       )
											_snake_case:   int  					=		test_model(__snake_case       )

											_snake_case:   Union[str, Any]  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   int  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       ,     atol=1e-5       )       )

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

											# You need to use the sequential hook to chain two or more hooks
											_snake_case:   List[str]  					=		SequentialHook(PostForwardHook()       ,     PostForwardHook()       )
											add_hook_to_module(__snake_case       ,     __snake_case       )

											_snake_case:   List[Any]  					=		test_model(__snake_case       )
											assert torch.allclose(__snake_case       ,     output + 2       ,     atol=1e-5       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		ModelForTest()
											_snake_case:   Dict  					=		torch.randn(2       ,     3       )
											_snake_case:   Tuple  					=		test_model(__snake_case       )

											_snake_case:   Tuple  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   List[str]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       )       )
											self.assertTrue(outputa.requires_grad       )

											_snake_case:   Any  					=		True
											_snake_case:   str  					=		test_model(__snake_case       )
											self.assertFalse(outputa.requires_grad       )


					@require_multi_gpu
					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   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
											_snake_case:   int  					=		torch.randn(2       ,     3       )
											_snake_case:   Tuple  					=		model(__snake_case       )
											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(__snake_case       ,     AlignDevicesHook(io_same_device=__snake_case       )       )
											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       ).to(0       )
											_snake_case:   Any  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     torch.device(0       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Union[str, 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
											_snake_case:   Optional[int]  					=		{'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

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

											# 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
											_snake_case:   Optional[int]  					=		torch.device(hook_kwargs['execution_device']       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

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

											# 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
											_snake_case:   Optional[Any]  					=		{
											    'execution_device': 0 if torch.cuda.is_available() else 'cpu',
											    'offload': True,
											    'offload_buffers': True,
											}

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

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

											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# 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    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Union[str, 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
											_snake_case:   Optional[Any]  					=		0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(__snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       )

											# 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
											_snake_case:   List[str]  					=		torch.device(__snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

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

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

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

											_snake_case:   Any  					=		torch.randn(2       ,     3       )
											_snake_case:   List[Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__snake_case       )
											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    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Tuple  					=		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
											_snake_case:   Tuple  					=		0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(
											    __snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       ,     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
											_snake_case:   str  					=		torch.device(__snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

											_snake_case:   Dict  					=		torch.randn(2       ,     3       )
											_snake_case:   Dict  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

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

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

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

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





