





"""simple docstring"""





import collections
import os
import re
from pathlib import Path


lowercase__				:int				   =       '''src/transformers'''


# Matches is_xxx_available()
lowercase__				:int				   =       re.compile(R'is\_([a-z_]*)_available()')
# Catches a one-line _import_struct = {xxx}
lowercase__				:int				   =       re.compile(R'^_import_structure\s+=\s+\{([^\}]+)\}')
# Catches a line with a key-values pattern: "bla": ["foo", "bar"]
lowercase__				:Optional[int]				   =       re.compile(R'\s+"\S*":\s+\[([^\]]*)\]')
# Catches a line if not is_foo_available
lowercase__				:Dict				   =       re.compile(R'^\s*if\s+not\s+is\_[a-z_]*\_available\(\)')
# Catches a line _import_struct["bla"].append("foo")
lowercase__				:Any				   =       re.compile(R'^\s*_import_structure\["\S*"\]\.append\("(\S*)"\)')
# Catches a line _import_struct["bla"].extend(["foo", "bar"]) or _import_struct["bla"] = ["foo", "bar"]
lowercase__				:str				   =       re.compile(R'^\s*_import_structure\[\S*\](?:\.extend\(|\s*=\s+)\[([^\]]*)\]')
# Catches a line with an object between quotes and a comma:     "MyModel",
lowercase__				:Tuple				   =       re.compile(R'^\s+"([^"]+)",')
# Catches a line with objects between brackets only:    ["foo", "bar"],
lowercase__				:Any				   =       re.compile(R'^\s+\[([^\]]+)\]')
# Catches a line with from foo import bar, bla, boo
lowercase__				:str				   =       re.compile(R'\s+from\s+\S*\s+import\s+([^\(\s].*)\n')
# Catches a line with try:
lowercase__				:Optional[int]				   =       re.compile(R'^\s*try:')
# Catches a line with else:
lowercase__				:Dict				   =       re.compile(R'^\s*else:')


def       lowerCamelCase_       (   UpperCAmelCase_  )						->Dict:





		"""simple docstring"""


		if _re_test_backend.search(UpperCAmelCase_  ) is None:
				return None
		__UpperCAmelCase    :						Tuple       =						[b[0] for b in _re_backend.findall(UpperCAmelCase_  )]
		backends.sort()
		return "_and_".join(UpperCAmelCase_  )


def       lowerCamelCase_       (   UpperCAmelCase_  )						->Any:





		"""simple docstring"""


		with open(UpperCAmelCase_  ,			'''r'''  ,			encoding='''utf-8'''  ,			newline='''\n'''  ) as f:
				__UpperCAmelCase    :						List[Any]       =						f.readlines()

		__UpperCAmelCase    :						Union[str, Any]       =						0
		while line_index < len(UpperCAmelCase_  ) and not lines[line_index].startswith('''_import_structure = {'''  ):
				line_index += 1

		# If this is a traditional init, just return.
		if line_index >= len(UpperCAmelCase_  ):
				return None

		# First grab the objects without a specific backend in _import_structure
		__UpperCAmelCase    :						List[Any]       =						[]
		while not lines[line_index].startswith('''if TYPE_CHECKING'''  ) and find_backend(lines[line_index]  ) is None:
				__UpperCAmelCase    :						str       =						lines[line_index]
				# If we have everything on a single line, let's deal with it.
				if _re_one_line_import_struct.search(UpperCAmelCase_  ):
						__UpperCAmelCase    :						List[str]       =						_re_one_line_import_struct.search(UpperCAmelCase_  ).groups()[0]
						__UpperCAmelCase    :						Dict       =						re.findall(R'''\[([^\]]+)\]'''  ,			UpperCAmelCase_  )
						for imp in imports:
								objects.extend([obj[1:-1] for obj in imp.split(''', '''  )]  )
						line_index += 1
						continue
				__UpperCAmelCase    :						Tuple       =						_re_import_struct_key_value.search(UpperCAmelCase_  )
				if single_line_import_search is not None:
						__UpperCAmelCase    :						List[str]       =						[obj[1:-1] for obj in single_line_import_search.groups()[0].split(''', '''  ) if len(UpperCAmelCase_  ) > 0]
						objects.extend(UpperCAmelCase_  )
				elif line.startswith(''' ''' * 8 + '''\"'''  ):
						objects.append(line[9:-3]  )
				line_index += 1

		__UpperCAmelCase    :						Union[str, Any]       =						{'''none''': objects}
		# Let's continue with backend-specific objects in _import_structure
		while not lines[line_index].startswith('''if TYPE_CHECKING'''  ):
				# If the line is an if not is_backend_available, we grab all objects associated.
				__UpperCAmelCase    :						List[str]       =						find_backend(lines[line_index]  )
				# Check if the backend declaration is inside a try block:
				if _re_try.search(lines[line_index - 1]  ) is None:
						__UpperCAmelCase    :						Tuple       =						None

				if backend is not None:
						line_index += 1

						# Scroll until we hit the else block of try-except-else
						while _re_else.search(lines[line_index]  ) is None:
								line_index += 1

						line_index += 1

						__UpperCAmelCase    :						int       =						[]
						# Until we unindent, add backend objects to the list
						while len(lines[line_index]  ) <= 1 or lines[line_index].startswith(''' ''' * 4  ):
								__UpperCAmelCase    :						List[str]       =						lines[line_index]
								if _re_import_struct_add_one.search(UpperCAmelCase_  ) is not None:
										objects.append(_re_import_struct_add_one.search(UpperCAmelCase_  ).groups()[0]  )
								elif _re_import_struct_add_many.search(UpperCAmelCase_  ) is not None:
										__UpperCAmelCase    :						str       =						_re_import_struct_add_many.search(UpperCAmelCase_  ).groups()[0].split(''', '''  )
										__UpperCAmelCase    :						List[Any]       =						[obj[1:-1] for obj in imports if len(UpperCAmelCase_  ) > 0]
										objects.extend(UpperCAmelCase_  )
								elif _re_between_brackets.search(UpperCAmelCase_  ) is not None:
										__UpperCAmelCase    :						Dict       =						_re_between_brackets.search(UpperCAmelCase_  ).groups()[0].split(''', '''  )
										__UpperCAmelCase    :						Dict       =						[obj[1:-1] for obj in imports if len(UpperCAmelCase_  ) > 0]
										objects.extend(UpperCAmelCase_  )
								elif _re_quote_object.search(UpperCAmelCase_  ) is not None:
										objects.append(_re_quote_object.search(UpperCAmelCase_  ).groups()[0]  )
								elif line.startswith(''' ''' * 8 + '''\"'''  ):
										objects.append(line[9:-3]  )
								elif line.startswith(''' ''' * 12 + '''\"'''  ):
										objects.append(line[13:-3]  )
								line_index += 1

						__UpperCAmelCase    :						str       =						objects
				else:
						line_index += 1

    # At this stage we are in the TYPE_CHECKING part, first grab the objects without a specific backend
		__UpperCAmelCase    :						Dict       =						[]
		while (
		    line_index < len(UpperCAmelCase_  )
		    and find_backend(lines[line_index]  ) is None
		    and not lines[line_index].startswith('''else'''  )
		):
				__UpperCAmelCase    :						List[Any]       =						lines[line_index]
				__UpperCAmelCase    :						List[Any]       =						_re_import.search(UpperCAmelCase_  )
				if single_line_import_search is not None:
						objects.extend(single_line_import_search.groups()[0].split(''', '''  )  )
				elif line.startswith(''' ''' * 8  ):
						objects.append(line[8:-2]  )
				line_index += 1

		__UpperCAmelCase    :						Optional[Any]       =						{'''none''': objects}
		# Let's continue with backend-specific objects
		while line_index < len(UpperCAmelCase_  ):
				# If the line is an if is_backend_available, we grab all objects associated.
				__UpperCAmelCase    :						Optional[int]       =						find_backend(lines[line_index]  )
				# Check if the backend declaration is inside a try block:
				if _re_try.search(lines[line_index - 1]  ) is None:
						__UpperCAmelCase    :						str       =						None

				if backend is not None:
						line_index += 1

						# Scroll until we hit the else block of try-except-else
						while _re_else.search(lines[line_index]  ) is None:
								line_index += 1

						line_index += 1

						__UpperCAmelCase    :						Union[str, Any]       =						[]
						# Until we unindent, add backend objects to the list
						while len(lines[line_index]  ) <= 1 or lines[line_index].startswith(''' ''' * 8  ):
								__UpperCAmelCase    :						List[Any]       =						lines[line_index]
								__UpperCAmelCase    :						Optional[Any]       =						_re_import.search(UpperCAmelCase_  )
								if single_line_import_search is not None:
										objects.extend(single_line_import_search.groups()[0].split(''', '''  )  )
								elif line.startswith(''' ''' * 12  ):
										objects.append(line[12:-2]  )
								line_index += 1

						__UpperCAmelCase    :						Tuple       =						objects
				else:
						line_index += 1

		return import_dict_objects, type_hint_objects


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  )						->Any:





		"""simple docstring"""


		def find_duplicates(UpperCAmelCase_  ):
				return [k for k, v in collections.Counter(UpperCAmelCase_  ).items() if v > 1]

		if list(import_dict_objects.keys()  ) != list(type_hint_objects.keys()  ):
				return ["Both sides of the init do not have the same backends!"]

		__UpperCAmelCase    :						Tuple       =						[]
		for key in import_dict_objects.keys():
				__UpperCAmelCase    :						Union[str, Any]       =						find_duplicates(import_dict_objects[key]  )
				if duplicate_imports:
						errors.append(f'''Duplicate _import_structure definitions for: {duplicate_imports}'''  )
				__UpperCAmelCase    :						Dict       =						find_duplicates(type_hint_objects[key]  )
				if duplicate_type_hints:
						errors.append(f'''Duplicate TYPE_CHECKING objects for: {duplicate_type_hints}'''  )

				if sorted(set(import_dict_objects[key]  )  ) != sorted(set(type_hint_objects[key]  )  ):
						__UpperCAmelCase    :						Optional[int]       =						'''base imports''' if key == '''none''' else f'''{key} backend'''
						errors.append(f'''Differences for {name}:'''  )
						for a in type_hint_objects[key]:
								if a not in import_dict_objects[key]:
										errors.append(f'''  {a} in TYPE_HINT but not in _import_structure.'''  )
						for a in import_dict_objects[key]:
								if a not in type_hint_objects[key]:
										errors.append(f'''  {a} in _import_structure but not in TYPE_HINT.'''  )
		return errors


def       lowerCamelCase_       (   )						->Dict:





		"""simple docstring"""


		__UpperCAmelCase    :						List[Any]       =						[]
		for root, _, files in os.walk(UpperCAmelCase_  ):
				if "__init__.py" in files:
						__UpperCAmelCase    :						List[Any]       =						os.path.join(UpperCAmelCase_  ,			'''__init__.py'''  )
						__UpperCAmelCase    :						Any       =						parse_init(UpperCAmelCase_  )
						if objects is not None:
								__UpperCAmelCase    :						List[str]       =						analyze_results(*UpperCAmelCase_  )
								if len(UpperCAmelCase_  ) > 0:
										__UpperCAmelCase    :						List[str]       =						f'''Problem in {fname}, both halves do not define the same objects.\n{errors[0]}'''
										failures.append('''\n'''.join(UpperCAmelCase_  )  )
		if len(UpperCAmelCase_  ) > 0:
				raise ValueError('''\n\n'''.join(UpperCAmelCase_  )  )


def       lowerCamelCase_       (   )						->Tuple:





		"""simple docstring"""


		__UpperCAmelCase    :						Any       =						[]
		for path, directories, files in os.walk(UpperCAmelCase_  ):
				for folder in directories:
						# Ignore private modules
						if folder.startswith('''_'''  ):
								directories.remove(UpperCAmelCase_  )
								continue
						# Ignore leftovers from branches (empty folders apart from pycache)
						if len(list((Path(UpperCAmelCase_  ) / folder).glob('''*.py'''  )  )  ) == 0:
								continue
						__UpperCAmelCase    :						Optional[int]       =						str((Path(UpperCAmelCase_  ) / folder).relative_to(UpperCAmelCase_  )  )
						__UpperCAmelCase    :						List[str]       =						short_path.replace(os.path.sep  ,			'''.'''  )
						submodules.append(UpperCAmelCase_  )
				for fname in files:
						if fname == "__init__.py":
								continue
						__UpperCAmelCase    :						int       =						str((Path(UpperCAmelCase_  ) / fname).relative_to(UpperCAmelCase_  )  )
						__UpperCAmelCase    :						Dict       =						short_path.replace('''.py'''  ,			''''''  ).replace(os.path.sep  ,			'''.'''  )
						if len(submodule.split('''.'''  )  ) == 1:
								submodules.append(UpperCAmelCase_  )
		return submodules


lowercase__				:List[str]				   =       [
    '''convert_pytorch_checkpoint_to_tf2''',
    '''modeling_flax_pytorch_utils''',
    '''models.esm.openfold_utils''',
]


def       lowerCamelCase_       (   )						->Dict:





		"""simple docstring"""


		from transformers.utils import direct_transformers_import

		__UpperCAmelCase    :						Tuple       =						direct_transformers_import(UpperCAmelCase_  )

		__UpperCAmelCase    :						List[Any]       =						set(transformers._import_structure.keys()  )
		# This contains all the base keys of the _import_structure object defined in the init, but if the user is missing
		# some optional dependencies, they may not have all of them. Thus we read the init to read all additions and
		# (potentiall re-) add them.
		with open(os.path.join(UpperCAmelCase_  ,			'''__init__.py'''  )  ,			'''r'''  ) as f:
				__UpperCAmelCase    :						int       =						f.read()
		import_structure_keys.update(set(re.findall(R'''import_structure\[\"([^\"]*)\"\]'''  ,			UpperCAmelCase_  )  )  )

		__UpperCAmelCase    :						Any       =						[
		    module
		    for module in get_transformers_submodules()
		    if module not in IGNORE_SUBMODULES and module not in import_structure_keys
		]

		if len(UpperCAmelCase_  ) > 0:
				__UpperCAmelCase    :						Optional[int]       =						'''\n'''.join(f'''- {module}''' for module in module_not_registered  )
				raise ValueError(
				    '''The following submodules are not properly registed in the main init of Transformers:\n'''
				    f'''{list_of_modules}\n'''
				    '''Make sure they appear somewhere in the keys of `_import_structure` with an empty list as value.'''  )


if __name__ == "__main__":
		check_all_inits()
		check_submodules()






"""simple docstring"""





import inspect
import unittest

import torch
import torch.nn as nn

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




class 							snake_case							(			nn.Module ):

      '''simple docstring'''



      def __init__(							self       :      List[Any] ):


        '''simple docstring'''




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







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


        '''simple docstring'''




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




class 							snake_case							(			__UpperCAmelCase ):

      '''simple docstring'''



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


        '''simple docstring'''




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




class 							snake_case							(			__UpperCAmelCase ):

      '''simple docstring'''



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


        '''simple docstring'''




        return output + 1




class 							snake_case							(			unittest.TestCase ):

      '''simple docstring'''



      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




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

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

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

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

      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




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

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

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

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

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

      def       A_					(							self       :      Any ):


        '''simple docstring'''




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

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

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

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

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

      def       A_					(							self       :      Any ):


        '''simple docstring'''




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

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

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

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

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

      def       A_					(							self       :      Dict ):


        '''simple docstring'''




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

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

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

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


        '''simple docstring'''




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

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

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

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

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

      def       A_					(							self       :      Tuple ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						ModelForTest()

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

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

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

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

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

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

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

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

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

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

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

      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[str]       =						ModelForTest()

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

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

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

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

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

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

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

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

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







      def       A_					(							self       :      Any ):


        '''simple docstring'''




        __UpperCAmelCase    :						Dict       =						ModelForTest()

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

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

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

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

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

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

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

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

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