


import unittest

import numpy as np
import requests

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
						import torch

						from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_11
else:
						__lowerCAmelCase      :					int    							=False

if is_vision_available():
						from PIL import Image

						from transformers import PixaStructImageProcessor
class 			_lowercase (     unittest.TestCase			):






				'''simple docstring'''



				def __init__(			self     :Optional[Any]     ,     lowerCAmelCase__     :Dict     ,     lowerCAmelCase__     :List[str]=7     ,     lowerCAmelCase__     :Union[str, Any]=3     ,     lowerCAmelCase__     :int=18     ,     lowerCAmelCase__     :str=30     ,     lowerCAmelCase__     :int=400     ,     lowerCAmelCase__     :Dict=None     ,     lowerCAmelCase__     :Tuple=True     ,     lowerCAmelCase__     :Any=True     ,     lowerCAmelCase__     :Tuple=None     ,     )						->     Union[str, Any]:
										__SCREAMING_SNAKE_CASE		:	str						     =						size if size is not None else {"""height""": 20, """width""": 20}
										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						parent
										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						batch_size
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						num_channels
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						image_size
										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						min_resolution
										__SCREAMING_SNAKE_CASE		:	str						     =						max_resolution
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						size
										__SCREAMING_SNAKE_CASE		:	Any						     =						do_normalize
										__SCREAMING_SNAKE_CASE		:	int						     =						do_convert_rgb
										__SCREAMING_SNAKE_CASE		:	Any						     =						[512, 1_024, 2_048, 4_096]
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						patch_size if patch_size is not None else {"""height""": 16, """width""": 16}


				def 						__magic_name__(			self     :Dict      )						->     Dict:
										return {"do_normalize": self.do_normalize, "do_convert_rgb": self.do_convert_rgb}



				def 						__magic_name__(			self     :Optional[Any]      )						->     List[Any]:
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						"""https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/australia.jpg"""
										__SCREAMING_SNAKE_CASE		:	Any						     =						Image.open(requests.get(__snake_case     ,     stream=__snake_case      ).raw      ).convert('''RGB'''      )
										return raw_image




@unittest.skipIf(
    not is_torch_greater_or_equal_than_1_11		,     reason='''`Pix2StructImageProcessor` requires `torch>=1.11.0`.'''		,     )
@require_torch
@require_vision
class 			_lowercase (     UpperCamelCase__		,     unittest.TestCase			):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__							:     Optional[Any]         =						PixaStructImageProcessor if is_vision_available() else None


				def 						__magic_name__(			self     :Union[str, Any]      )						->     Dict:
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						PixaStructImageProcessingTester(self      )


				@property
				def 						__magic_name__(			self     :str      )						->     str:
										return self.image_processor_tester.prepare_image_processor_dict()


				def 						__magic_name__(			self     :List[Any]      )						->     str:
										__SCREAMING_SNAKE_CASE		:	Any						     =						self.image_processing_class(**self.image_processor_dict      )
										self.assertTrue(hasattr(__snake_case     ,     '''do_normalize'''      )      )
										self.assertTrue(hasattr(__snake_case     ,     '''do_convert_rgb'''      )      )


				def 						__magic_name__(			self     :int      )						->     Optional[Any]:
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						self.image_processor_tester.prepare_dummy_image()

										__SCREAMING_SNAKE_CASE		:	Dict						     =						self.image_processing_class(**self.image_processor_dict      )
										__SCREAMING_SNAKE_CASE		:	str						     =						2_048

										__SCREAMING_SNAKE_CASE		:	Tuple						     =						image_processor(__snake_case     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      )
										self.assertTrue(torch.allclose(inputs.flattened_patches.mean()     ,     torch.tensor(0.0606      )     ,     atol=1E-3     ,     rtol=1E-3      )      )


				def 						__magic_name__(			self     :Optional[Any]      )						->     int:
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						self.image_processing_class(**self.image_processor_dict      )
										# create random PIL images
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						prepare_image_inputs(self.image_processor_tester     ,     equal_resolution=__snake_case      )
										for image in image_inputs:
																self.assertIsInstance(__snake_case     ,     Image.Image      )

										# Test not batched input
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						(
										    (self.image_processor_tester.patch_size["""height"""] * self.image_processor_tester.patch_size["""width"""])
										    * self.image_processor_tester.num_channels
										) + 2

										for max_patch in self.image_processor_tester.max_patches:
																# Test not batched input
																__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						image_processor(
																    image_inputs[0]     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (1, max_patch, expected_hidden_dim)     ,     )

																# Test batched
																__SCREAMING_SNAKE_CASE		:	List[str]						     =						image_processor(
																    __snake_case     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim)     ,     )


				def 						__magic_name__(			self     :Union[str, Any]      )						->     List[str]:
										__SCREAMING_SNAKE_CASE		:	Any						     =						self.image_processing_class(**self.image_processor_dict      )
										# create random PIL images
										__SCREAMING_SNAKE_CASE		:	str						     =						prepare_image_inputs(self.image_processor_tester     ,     equal_resolution=__snake_case      )
										for image in image_inputs:
																self.assertIsInstance(__snake_case     ,     Image.Image      )

										# Test not batched input
										__SCREAMING_SNAKE_CASE		:	str						     =						(
										    (self.image_processor_tester.patch_size["""height"""] * self.image_processor_tester.patch_size["""width"""])
										    * self.image_processor_tester.num_channels
										) + 2

										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						True

										for max_patch in self.image_processor_tester.max_patches:
																# Test not batched input
																with self.assertRaises(__snake_case      ):
																						__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						image_processor(
																						    image_inputs[0]     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches

																__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						"""Hello"""

																__SCREAMING_SNAKE_CASE		:	List[Any]						     =						image_processor(
																    image_inputs[0]     ,     return_tensors='''pt'''     ,     max_patches=__snake_case     ,     header_text=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (1, max_patch, expected_hidden_dim)     ,     )

																# Test batched
																__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						image_processor(
																    __snake_case     ,     return_tensors='''pt'''     ,     max_patches=__snake_case     ,     header_text=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim)     ,     )


				def 						__magic_name__(			self     :Dict      )						->     Tuple:
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						self.image_processing_class(**self.image_processor_dict      )
										# create random numpy tensors
										__SCREAMING_SNAKE_CASE		:	str						     =						prepare_image_inputs(self.image_processor_tester     ,     equal_resolution=__snake_case     ,     numpify=__snake_case      )
										for image in image_inputs:
																self.assertIsInstance(__snake_case     ,     np.ndarray      )

										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						(
										    (self.image_processor_tester.patch_size["""height"""] * self.image_processor_tester.patch_size["""width"""])
										    * self.image_processor_tester.num_channels
										) + 2

										for max_patch in self.image_processor_tester.max_patches:
																# Test not batched input
																__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						image_processor(
																    image_inputs[0]     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (1, max_patch, expected_hidden_dim)     ,     )

																# Test batched
																__SCREAMING_SNAKE_CASE		:	str						     =						image_processor(
																    __snake_case     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim)     ,     )



				def 						__magic_name__(			self     :Union[str, Any]      )						->     Dict:
										__SCREAMING_SNAKE_CASE		:	Dict						     =						self.image_processing_class(**self.image_processor_dict      )
										# create random PyTorch tensors
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						prepare_image_inputs(self.image_processor_tester     ,     equal_resolution=__snake_case     ,     torchify=__snake_case      )
										for image in image_inputs:
																self.assertIsInstance(__snake_case     ,     torch.Tensor      )

										# Test not batched input
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						(
										    (self.image_processor_tester.patch_size["""height"""] * self.image_processor_tester.patch_size["""width"""])
										    * self.image_processor_tester.num_channels
										) + 2

										for max_patch in self.image_processor_tester.max_patches:
																# Test not batched input
																__SCREAMING_SNAKE_CASE		:	Tuple						     =						image_processor(
																    image_inputs[0]     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (1, max_patch, expected_hidden_dim)     ,     )

																# Test batched
																__SCREAMING_SNAKE_CASE		:	str						     =						image_processor(
																    __snake_case     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim)     ,     )




@unittest.skipIf(
    not is_torch_greater_or_equal_than_1_11		,     reason='''`Pix2StructImageProcessor` requires `torch>=1.11.0`.'''		,     )
@require_torch
@require_vision
class 			_lowercase (     UpperCamelCase__		,     unittest.TestCase			):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__							:     List[Any]         =						PixaStructImageProcessor if is_vision_available() else None


				def 						__magic_name__(			self     :Optional[Any]      )						->     Dict:
										__SCREAMING_SNAKE_CASE		:	Tuple						     =						PixaStructImageProcessingTester(self     ,     num_channels=4      )
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						3


				@property
				def 						__magic_name__(			self     :Any      )						->     Union[str, Any]:
										return self.image_processor_tester.prepare_image_processor_dict()


				def 						__magic_name__(			self     :Optional[int]      )						->     Any:
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						self.image_processing_class(**self.image_processor_dict      )
										self.assertTrue(hasattr(__snake_case     ,     '''do_normalize'''      )      )
										self.assertTrue(hasattr(__snake_case     ,     '''do_convert_rgb'''      )      )



				def 						__magic_name__(			self     :Union[str, Any]      )						->     Union[str, Any]:
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						self.image_processing_class(**self.image_processor_dict      )
										# create random PIL images
										__SCREAMING_SNAKE_CASE		:	Tuple						     =						prepare_image_inputs(self.image_processor_tester     ,     equal_resolution=__snake_case      )
										for image in image_inputs:
																self.assertIsInstance(__snake_case     ,     Image.Image      )

										# Test not batched input
										__SCREAMING_SNAKE_CASE		:	str						     =						(
										    (self.image_processor_tester.patch_size["""height"""] * self.image_processor_tester.patch_size["""width"""])
										    * (self.image_processor_tester.num_channels - 1)
										) + 2

										for max_patch in self.image_processor_tester.max_patches:
																# Test not batched input
																__SCREAMING_SNAKE_CASE		:	Any						     =						image_processor(
																    image_inputs[0]     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (1, max_patch, expected_hidden_dim)     ,     )

																# Test batched
																__SCREAMING_SNAKE_CASE		:	List[Any]						     =						image_processor(
																    __snake_case     ,     return_tensors='''pt'''     ,     max_patches=__snake_case      ).flattened_patches
																self.assertEqual(
																    encoded_images.shape     ,     (self.image_processor_tester.batch_size, max_patch, expected_hidden_dim)     ,     )






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

						def __init__(		self :List[Any]			):
										'''simple docstring'''

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


						def      A__    (		self :Dict     ,			__snake_case :Tuple			):
										'''simple docstring'''

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




class 		__A   (      UpperCamelCase__      ):

						def      A__    (		self :Any     ,			__snake_case :Optional[Any]     ,			*__snake_case :List[Any]     ,			**__snake_case :Any			):
										'''simple docstring'''

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




class 		__A   (      UpperCamelCase__      ):

						def      A__    (		self :List[str]     ,			__snake_case :Tuple     ,			__snake_case :Union[str, Any]			):
										'''simple docstring'''

										return output + 1




class 		__A   (      unittest.TestCase      ):

						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						int										=ModelForTest()
										__magic_name__						:						Tuple										=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      A__    (		self :Tuple			):
										'''simple docstring'''

										__magic_name__						:						int										=ModelForTest()
										__magic_name__						:						List[str]										=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      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Any										=ModelForTest()
										__magic_name__						:						Any										=torch.randn(2     ,			3			)
										__magic_name__						:						Any										=test_model(x + 1			)
										__magic_name__						:						Optional[Any]										=test_model(x + 2			)

										__magic_name__						:						int										=PreForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						int										=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
										__magic_name__						:						str										=PreForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						List[str]										=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
										__magic_name__						:						Optional[Any]										=SequentialHook(PreForwardHook()     ,			PreForwardHook()			)
										add_hook_to_module(__snake_case     ,			__snake_case			)

										__magic_name__						:						Any										=test_model(__snake_case			)
										assert torch.allclose(__snake_case     ,			__snake_case     ,			atol=1E-5			)

						def      A__    (		self :Any			):
										'''simple docstring'''

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

										__magic_name__						:						Dict										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						Any										=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
										__magic_name__						:						Union[str, Any]										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						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
										__magic_name__						:						Union[str, Any]										=SequentialHook(PostForwardHook()     ,			PostForwardHook()			)
										add_hook_to_module(__snake_case     ,			__snake_case			)

										__magic_name__						:						Union[str, Any]										=test_model(__snake_case			)
										assert torch.allclose(__snake_case     ,			output + 2     ,			atol=1E-5			)

						def      A__    (		self :Tuple			):
										'''simple docstring'''

										__magic_name__						:						Tuple										=ModelForTest()
										__magic_name__						:						int										=torch.randn(2     ,			3			)
										__magic_name__						:						Union[str, Any]										=test_model(__snake_case			)

										__magic_name__						:						Union[str, Any]										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						Dict										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			output + 1			)			)
										self.assertTrue(outputa.requires_grad			)

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

						@require_multi_gpu
						def      A__    (		self :List[str]			):
										'''simple docstring'''

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

						def      A__    (		self :List[Any]			):
										'''simple docstring'''

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

										__magic_name__						:						Union[str, Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						Optional[int]										=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
										__magic_name__						:						Tuple										={
										    """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"""			)			)

										__magic_name__						:						Tuple										=torch.randn(2     ,			3			)
										__magic_name__						:						int										=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      A__    (		self :List[Any]			):
										'''simple docstring'''

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

										__magic_name__						:						List[Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						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"""			)			)

										__magic_name__						:						Optional[int]										=torch.randn(2     ,			3			)
										__magic_name__						:						Union[str, 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      A__    (		self :List[str]			):
										'''simple docstring'''

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

										__magic_name__						:						int										=torch.randn(2     ,			3			)
										__magic_name__						:						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"""			)			)

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

										__magic_name__						:						List[Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						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"""			)			)




