




"""simple docstring"""






import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SPIECE_UNDERLINE, is_sentencepiece_available
from transformers.models.speech_to_text import SpeechaTextTokenizer
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


UpperCamelCase__    =      get_tests_dir('''fixtures/test_sentencepiece.model''')

if is_sentencepiece_available():
				import sentencepiece as sp


UpperCamelCase__    =      5
UpperCamelCase__    =      1_0




@require_sentencepiece
@require_tokenizers
class 	a__      (							__lowercase       ,						unittest.TestCase	):
						snake_case__							   =			SpeechaTextTokenizer
						snake_case__							   =			False
						snake_case__							   =			True





						def    __UpperCamelCase  (				self   :						List[Any])    ->							str:




											"""simple docstring"""







											super().setUp()

											_lowerCAmelCase:Tuple						=	sp.SentencePieceProcessor()
											spm_model.Load(a__)
											_lowerCAmelCase:int						=	["""<s>""", """<pad>""", """</s>""", """<unk>"""]

											vocab += [spm_model.IdToPiece(id_) for id_ in range(len(a__))]
											_lowerCAmelCase:int						=	dict(zip(a__  ,range(len(a__))))

											_lowerCAmelCase:Union[str, Any]						=	Path(self.tmpdirname)
											save_json(a__  ,save_dir / VOCAB_FILES_NAMES['''vocab_file'''])
											if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
																copyfile(a__  ,save_dir / VOCAB_FILES_NAMES['''spm_file'''])

											_lowerCAmelCase:Dict						=	SpeechaTextTokenizer.from_pretrained(self.tmpdirname)
											tokenizer.save_pretrained(self.tmpdirname)





						def    __UpperCamelCase  (				self   :						Dict)    ->							int:




											"""simple docstring"""







											_lowerCAmelCase:Any						=	"""<pad>"""
											_lowerCAmelCase:Optional[Any]						=	1

											self.assertEqual(self.get_tokenizer()._convert_token_to_id(a__)  ,a__)
											self.assertEqual(self.get_tokenizer()._convert_id_to_token(a__)  ,a__)





						def    __UpperCamelCase  (				self   :						List[str])    ->							List[str]:




											"""simple docstring"""







											_lowerCAmelCase:List[Any]						=	list(self.get_tokenizer().get_vocab().keys())

											self.assertEqual(vocab_keys[0]  ,'''<s>''')
											self.assertEqual(vocab_keys[1]  ,'''<pad>''')
											self.assertEqual(vocab_keys[-1]  ,'''j''')
											self.assertEqual(len(a__)  ,1001)





						def    __UpperCamelCase  (				self   :						List[str])    ->							int:




											"""simple docstring"""







											self.assertEqual(self.get_tokenizer().vocab_size  ,1001)





						def    __UpperCamelCase  (				self   :						str)    ->							Tuple:




											"""simple docstring"""







											_lowerCAmelCase:Union[str, Any]						=	SpeechaTextTokenizer.from_pretrained(self.tmpdirname)

											_lowerCAmelCase:List[Any]						=	tokenizer.tokenize('''This is a test''')
											self.assertListEqual(a__  ,['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est'''])

											self.assertListEqual(
											    tokenizer.convert_tokens_to_ids(a__)  ,[289, 50, 14, 174, 386]  ,)

											_lowerCAmelCase:Optional[int]						=	tokenizer.tokenize('''I was born in 92000, and this is falsé.''')
											self.assertListEqual(
											    a__  ,[SPIECE_UNDERLINE + '''I''', SPIECE_UNDERLINE + '''was''', SPIECE_UNDERLINE + '''b''', '''or''', '''n''', SPIECE_UNDERLINE + '''in''', SPIECE_UNDERLINE + '''''', '''9''', '''2''', '''0''', '''0''', '''0''', ''',''', SPIECE_UNDERLINE + '''and''', SPIECE_UNDERLINE + '''this''', SPIECE_UNDERLINE + '''is''', SPIECE_UNDERLINE + '''f''', '''al''', '''s''', '''é''', '''.''']  ,)
											_lowerCAmelCase:Union[str, Any]						=	tokenizer.convert_tokens_to_ids(a__)
											self.assertListEqual(a__  ,[12, 25, 88, 59, 28, 23, 11, 4, 606, 351, 351, 351, 7, 16, 70, 50, 76, 84, 10, 4, 8])

											_lowerCAmelCase:List[Any]						=	tokenizer.convert_ids_to_tokens(a__)
											self.assertListEqual(
											    a__  ,[SPIECE_UNDERLINE + '''I''', SPIECE_UNDERLINE + '''was''', SPIECE_UNDERLINE + '''b''', '''or''', '''n''', SPIECE_UNDERLINE + '''in''', SPIECE_UNDERLINE + '''''', '''<unk>''', '''2''', '''0''', '''0''', '''0''', ''',''', SPIECE_UNDERLINE + '''and''', SPIECE_UNDERLINE + '''this''', SPIECE_UNDERLINE + '''is''', SPIECE_UNDERLINE + '''f''', '''al''', '''s''', '''<unk>''', '''.''']  ,)







						@slow
						def    __UpperCamelCase  (				self   :						int)    ->							Tuple:




											"""simple docstring"""







											_lowerCAmelCase:List[Any]						=	{"""input_ids""": [[3791, 797, 31, 11, 64, 797, 31, 2429, 433, 12, 1176, 12, 20, 786, 915, 142, 2413, 240, 37, 3238, 797, 31, 11, 35, 93, 915, 142, 2413, 240, 37, 5540, 567, 1276, 93, 37, 610, 40, 62, 455, 657, 1042, 123, 780, 177, 37, 309, 241, 1298, 514, 20, 292, 2737, 114, 2469, 241, 85, 64, 302, 548, 528, 423, 4, 509, 406, 423, 37, 601, 4, 777, 302, 548, 528, 423, 284, 4, 3388, 511, 459, 4, 3555, 40, 321, 302, 705, 4, 3388, 511, 583, 326, 5, 5, 5, 62, 3310, 560, 177, 2680, 217, 1508, 32, 31, 853, 418, 64, 583, 511, 1605, 62, 35, 93, 560, 177, 2680, 217, 1508, 1521, 64, 583, 511, 519, 62, 20, 1515, 764, 20, 149, 261, 5625, 7972, 20, 5540, 567, 1276, 93, 3925, 1675, 11, 15, 802, 7972, 576, 217, 1508, 11, 35, 93, 1253, 2441, 15, 289, 652, 31, 416, 321, 3842, 115, 40, 911, 8, 476, 619, 4, 380, 142, 423, 335, 240, 35, 93, 264, 8, 11, 335, 569, 420, 163, 5, 2], [260, 548, 528, 423, 20, 451, 20, 2681, 1153, 3434, 20, 5540, 37, 567, 126, 1253, 2441, 3376, 449, 210, 431, 1563, 177, 767, 5540, 11, 1203, 472, 11, 2953, 685, 285, 364, 706, 1153, 20, 6799, 20, 2869, 20, 4464, 126, 40, 2429, 20, 1040, 866, 2664, 418, 20, 318, 20, 1726, 186, 20, 265, 522, 35, 93, 2191, 4634, 20, 1040, 12, 6799, 15, 228, 2356, 142, 31, 11, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [2575, 2666, 684, 1582, 1176, 12, 627, 149, 619, 20, 4902, 563, 11, 20, 149, 261, 3420, 2356, 174, 142, 4714, 131, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], """attention_mask""": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
											# fmt: on

											self.tokenizer_integration_test_util(
											    expected_encoding=a__  ,model_name='''facebook/s2t-small-mustc-en-de-st'''  ,revision='''a14f04cf0776c02f62a8cb800cf7909e15ea23ad'''  ,)




@require_sentencepiece
class 	a__      (							unittest.TestCase	):
						snake_case__							   =			'''valhalla/s2t_mustc_multilinguial_medium'''

						snake_case__							   =			'''C\'est trop cool'''
						snake_case__							   =			'''Esto es genial'''





						@classmethod
						def    __UpperCamelCase  (				cls   :						Optional[Any])    ->							int:




											"""simple docstring"""







											_lowerCAmelCase:SpeechaTextTokenizer						=	SpeechaTextTokenizer.from_pretrained(cls.checkpoint_name)
											return cls





						def    __UpperCamelCase  (				self   :						List[Any])    ->							Optional[int]:




											"""simple docstring"""







											self.assertEqual(self.tokenizer.lang_code_to_id['''pt''']  ,4)
											self.assertEqual(self.tokenizer.lang_code_to_id['''ru''']  ,6)
											self.assertEqual(self.tokenizer.lang_code_to_id['''it''']  ,9)
											self.assertEqual(self.tokenizer.lang_code_to_id['''de''']  ,11)





						def    __UpperCamelCase  (				self   :						Optional[int])    ->							Dict:




											"""simple docstring"""







											self.assertEqual(self.tokenizer.vocab_size  ,1_0000)





						def    __UpperCamelCase  (				self   :						List[Any])    ->							Optional[Any]:




											"""simple docstring"""







											self.assertIn(a__  ,self.tokenizer.all_special_ids)
											_lowerCAmelCase:Tuple						=	[ES_CODE, 4, 1601, 47, 7647, 2]
											_lowerCAmelCase:Any						=	self.tokenizer.decode(a__  ,skip_special_tokens=a__)
											_lowerCAmelCase:int						=	self.tokenizer.decode(generated_ids[1:]  ,skip_special_tokens=a__)
											self.assertEqual(a__  ,a__)
											self.assertNotIn(self.tokenizer.eos_token  ,a__)





						def    __UpperCamelCase  (				self   :						Union[str, Any])    ->							Any:




											"""simple docstring"""







											_lowerCAmelCase:Optional[Any]						=	"""fr"""
											_lowerCAmelCase:int						=	self.tokenizer(self.french_text).input_ids
											self.assertEqual(encoded[0]  ,a__)
											self.assertEqual(encoded[-1]  ,self.tokenizer.eos_token_id)







						def    __UpperCamelCase  (				self   :						Dict)    ->							Union[str, Any]:




											"""simple docstring"""







											_lowerCAmelCase:Dict						=	"""fr"""
											self.assertListEqual(self.tokenizer.prefix_tokens  ,[FR_CODE])

											_lowerCAmelCase:Any						=	"""es"""
											self.assertListEqual(self.tokenizer.prefix_tokens  ,[ES_CODE])













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

		def __init__(   self					:   List[str]			)					->    Optional[Any]:


					"""simple docstring"""

					super().__init__()
					__snake_case					:      List[Any]							      =     nn.Linear(3  ,      4			)
					__snake_case					:      str							      =     nn.BatchNormad(4			)
					__snake_case					:      Optional[Any]							      =     nn.Linear(4  ,      5			)

		def 	lowercase__     (   self					:   str  ,      __magic_name__					:   Dict			)					->    List[str]:


					"""simple docstring"""

					return self.lineara(self.batchnorm(self.lineara(__magic_name__			)			)			)





class  _A (					__lowercase		):

		def 	lowercase__     (   self					:   List[str]  ,      __magic_name__					:   Tuple  ,      *__magic_name__					:   Dict  ,      **__magic_name__					:   Optional[Any]			)					->    Tuple:


					"""simple docstring"""

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





class  _A (					__lowercase		):

		def 	lowercase__     (   self					:   str  ,      __magic_name__					:   Union[str, Any]  ,      __magic_name__					:   Tuple			)					->    Union[str, Any]:


					"""simple docstring"""

					return output + 1





class  _A (					unittest.TestCase		):

		def 	lowercase__     (   self					:   Dict			)					->    Any:


					"""simple docstring"""

					__snake_case					:      int							      =     ModelForTest()
					__snake_case					:      Tuple							      =     ModelHook()

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

		def 	lowercase__     (   self					:   Tuple			)					->    List[str]:


					"""simple docstring"""

					__snake_case					:      List[Any]							      =     ModelForTest()
					__snake_case					:      Optional[int]							      =     ModelHook()

					add_hook_to_module(__magic_name__  ,      __magic_name__			)
					add_hook_to_module(__magic_name__  ,      __magic_name__  ,      append=__magic_name__			)

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

		def 	lowercase__     (   self					:   str			)					->    Union[str, Any]:


					"""simple docstring"""

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

					__snake_case					:      Union[str, Any]							      =     PreForwardHook()
					add_hook_to_module(__magic_name__  ,      __magic_name__			)
					__snake_case					:      int							      =     test_model(__magic_name__			)
					self.assertTrue(torch.allclose(__magic_name__  ,      __magic_name__  ,      atol=1E-5			)			)

					# Attaching a hook to a model when it already has one replaces, does not chain
					__snake_case					:      Optional[int]							      =     PreForwardHook()
					add_hook_to_module(__magic_name__  ,      __magic_name__			)
					__snake_case					:      List[Any]							      =     test_model(__magic_name__			)
					self.assertTrue(torch.allclose(__magic_name__  ,      __magic_name__  ,      atol=1E-5			)			)

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

					__snake_case					:      List[str]							      =     test_model(__magic_name__			)
					assert torch.allclose(__magic_name__  ,      __magic_name__  ,      atol=1E-5			)

		def 	lowercase__     (   self					:   Union[str, Any]			)					->    List[str]:


					"""simple docstring"""

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

					__snake_case					:      Any							      =     PostForwardHook()
					add_hook_to_module(__magic_name__  ,      __magic_name__			)
					__snake_case					:      Any							      =     test_model(__magic_name__			)
					self.assertTrue(torch.allclose(__magic_name__  ,      output + 1  ,      atol=1E-5			)			)

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

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

					__snake_case					:      int							      =     test_model(__magic_name__			)
					assert torch.allclose(__magic_name__  ,      output + 2  ,      atol=1E-5			)

		def 	lowercase__     (   self					:   str			)					->    int:


					"""simple docstring"""

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

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

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

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

		def 	lowercase__     (   self					:   Union[str, Any]			)					->    str:


					"""simple docstring"""

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

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

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

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

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

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

					# 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					:      str							      =     torch.randn(2  ,      3			)
					__snake_case					:      str							      =     model(__magic_name__			)
					self.assertEqual(output.device  ,      __magic_name__			)

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

		def 	lowercase__     (   self					:   Dict			)					->    str:


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

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

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

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

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

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

		def 	lowercase__     (   self					:   Any			)					->    Union[str, Any]:


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

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

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

					# 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[str]							      =     torch.randn(2  ,      3			)
					__snake_case					:      Dict							      =     model(__magic_name__			)
					self.assertEqual(output.device  ,      __magic_name__			)

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


