


import json
import os
import unittest

from transformers.models.gptsan_japanese.tokenization_gptsan_japanese import (
    VOCAB_FILES_NAMES,
    GPTSanJapaneseTokenizer,
)
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin

@require_tokenizers
class       __magic_name__	(     __lowerCAmelCase					,				unittest.TestCase):

     '''simple docstring'''





     A:		Union[str, Any]			=      GPTSanJapaneseTokenizer
     A:		Tuple			=      False
     A:		Any			=      {"do_clean_text": False, "add_prefix_space": False}


     def      UpperCAmelCase__	(      self							:     Any			)  ->      str:





      '''simple docstring'''






      super().setUp()

      # fmt: off
      UpperCamelCase__					:    int          =   ['''こん''', '''こんに''', '''にちは''', '''ばんは''', '''世界,㔺界''', '''、''', '''。''', '''<BR>''', '''<SP>''', '''<TAB>''', '''<URL>''', '''<EMAIL>''', '''<TEL>''', '''<DATE>''', '''<PRICE>''', '''<BLOCK>''', '''<KIGOU>''', '''<U2000U2BFF>''', '''<|emoji1|>''', '''<unk>''', '''<|bagoftoken|>''', '''<|endoftext|>''']
      # fmt: on
      UpperCamelCase__					:    Tuple          =   {'''emoji''': {'''\ud83d\ude00''': '''<|emoji1|>'''}, '''emoji_inv''': {'''<|emoji1|>''': '''\ud83d\ude00'''}}  # 😀
      UpperCamelCase__					:    Optional[int]          =   {'''unk_token''': '''<unk>'''}

      UpperCamelCase__					:    str          =   os.path.join(self.tmpdirname							,	VOCAB_FILES_NAMES['''vocab_file''']			)
      UpperCamelCase__					:    Optional[Any]          =   os.path.join(self.tmpdirname							,	VOCAB_FILES_NAMES['''emoji_file''']			)
      with open(self.vocab_file							,	'''w'''							,	encoding='''utf-8'''			) as vocab_writer:
       vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]			)			)
      with open(self.emoji_file							,	'''w'''			) as emoji_writer:
       emoji_writer.write(json.dumps(lowerCamelCase__			)			)


     def      UpperCAmelCase__	(      self							:     Optional[int]							,	**lowerCamelCase__							:     Optional[int]			)  ->      Any:





      '''simple docstring'''






      kwargs.update(self.special_tokens_map			)
      return GPTSanJapaneseTokenizer.from_pretrained(self.tmpdirname							,	**lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     Optional[int]							,	lowerCamelCase__							:     Dict			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   '''こんにちは、世界。 \nこんばんは、㔺界。😀'''
      UpperCamelCase__					:    List[Any]          =   '''こんにちは、世界。 \nこんばんは、世界。😀'''
      return input_text, output_text


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





      '''simple docstring'''






      UpperCamelCase__					:    Union[str, Any]          =   self.get_input_output_texts(lowerCamelCase__			)
      UpperCamelCase__					:    List[str]          =   tokenizer.encode(lowerCamelCase__							,	add_special_tokens=lowerCamelCase__			)
      UpperCamelCase__					:    str          =   tokenizer.decode(lowerCamelCase__							,	clean_up_tokenization_spaces=lowerCamelCase__			)
      return text, ids


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





      '''simple docstring'''






      pass  # TODO add if relevant


     def      UpperCAmelCase__	(      self							:     Optional[Any]			)  ->      Any:





      '''simple docstring'''






      pass  # TODO add if relevant


     def      UpperCAmelCase__	(      self							:     Any			)  ->      List[Any]:





      '''simple docstring'''






      pass  # TODO add if relevant


     def      UpperCAmelCase__	(      self							:     Any			)  ->      Any:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   self.get_tokenizer()

      # Testing tokenization
      UpperCamelCase__					:    List[Any]          =   '''こんにちは、世界。　こんばんは、㔺界。'''
      UpperCamelCase__					:    Optional[Any]          =   ['''こん''', '''にちは''', '''、''', '''世界''', '''。''', '''<SP>''', '''こん''', '''ばんは''', '''、''', '''㔺界''', '''。''']
      UpperCamelCase__					:    int          =   tokenizer.tokenize(lowerCamelCase__			)
      self.assertListEqual(lowerCamelCase__							,	lowerCamelCase__			)

      # Testing conversion to ids without special tokens
      UpperCamelCase__					:    List[str]          =   [0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6]
      UpperCamelCase__					:    Any          =   tokenizer.convert_tokens_to_ids(lowerCamelCase__			)
      self.assertListEqual(lowerCamelCase__							,	lowerCamelCase__			)

      # Testing conversion to ids with special tokens
      UpperCamelCase__					:    str          =   tokens + [tokenizer.unk_token]
      UpperCamelCase__					:    Tuple          =   [0, 2, 5, 4, 6, 8, 0, 3, 5, 4, 6, 19]
      UpperCamelCase__					:    Union[str, Any]          =   tokenizer.convert_tokens_to_ids(lowerCamelCase__			)
      self.assertListEqual(lowerCamelCase__							,	lowerCamelCase__			)


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





      '''simple docstring'''






      UpperCamelCase__					:    Optional[int]          =   self.get_tokenizer()

      # Testing tokenization
      UpperCamelCase__					:    Optional[int]          =   '''こんにちは、<|bagoftoken|>世界。こんばんは、<|bagoftoken|>㔺界。'''
      UpperCamelCase__					:    Optional[Any]          =   '''こんにちは、、、、世界。こんばんは、、、、世界。'''
      UpperCamelCase__					:    Tuple          =   tokenizer.encode(lowerCamelCase__			)
      UpperCamelCase__					:    Optional[Any]          =   tokenizer.decode(lowerCamelCase__			)
      self.assertEqual(lowerCamelCase__							,	lowerCamelCase__			)


     @slow
     def      UpperCAmelCase__	(      self							:     int			)  ->      Optional[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   self.tokenizer_class.from_pretrained('''Tanrei/GPTSAN-japanese'''			)

      # Testing tokenization
      UpperCamelCase__					:    str          =   '''こんにちは、世界。'''
      UpperCamelCase__					:    Optional[int]          =   '''こんばんは、㔺界。😀'''
      UpperCamelCase__					:    Optional[int]          =   '''こんにちは、世界。こんばんは、世界。😀'''
      UpperCamelCase__					:    List[str]          =   tokenizer.encode(prefix_text + input_text			)
      UpperCamelCase__					:    List[str]          =   tokenizer.encode(''''''							,	prefix_text=prefix_text + input_text			)
      UpperCamelCase__					:    str          =   tokenizer.encode(lowerCamelCase__							,	prefix_text=lowerCamelCase__			)
      UpperCamelCase__					:    Tuple          =   tokenizer.decode(lowerCamelCase__			)
      UpperCamelCase__					:    List[str]          =   tokenizer.decode(lowerCamelCase__			)
      UpperCamelCase__					:    Optional[int]          =   tokenizer.decode(lowerCamelCase__			)
      self.assertEqual(lowerCamelCase__							,	lowerCamelCase__			)
      self.assertEqual(lowerCamelCase__							,	lowerCamelCase__			)
      self.assertEqual(lowerCamelCase__							,	lowerCamelCase__			)


     @slow
     def      UpperCAmelCase__	(      self							:     Any			)  ->      List[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    List[Any]          =   self.tokenizer_class.from_pretrained('''Tanrei/GPTSAN-japanese'''			)

      # Testing tokenization
      UpperCamelCase__					:    Optional[int]          =   '''こんにちは、世界。'''
      UpperCamelCase__					:    Dict          =   '''こんばんは、㔺界。😀'''

      UpperCamelCase__					:    Optional[int]          =   len(tokenizer.encode(lowerCamelCase__			)			) - 2
      UpperCamelCase__					:    int          =   len(tokenizer.encode(lowerCamelCase__			)			) - 2

      UpperCamelCase__					:    Optional[Any]          =   [1] + [0] * (len_prefix + len_text + 1)
      UpperCamelCase__					:    Tuple          =   [1] * (len_prefix + len_text + 1) + [0]
      UpperCamelCase__					:    int          =   [1] + [1] * (len_prefix) + [0] * (len_text + 1)

      UpperCamelCase__					:    Tuple          =   tokenizer(prefix_text + input_text			).token_type_ids
      UpperCamelCase__					:    Tuple          =   tokenizer(''''''							,	prefix_text=prefix_text + input_text			).token_type_ids
      UpperCamelCase__					:    List[str]          =   tokenizer(lowerCamelCase__							,	prefix_text=lowerCamelCase__			).token_type_ids
      self.assertListEqual(lowerCamelCase__							,	lowerCamelCase__			)
      self.assertListEqual(lowerCamelCase__							,	lowerCamelCase__			)
      self.assertListEqual(lowerCamelCase__							,	lowerCamelCase__			)


     @slow
     def      UpperCAmelCase__	(      self							:     Optional[Any]			)  ->      Optional[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    Tuple          =   self.tokenizer_class.from_pretrained('''Tanrei/GPTSAN-japanese'''			)

      UpperCamelCase__					:    str          =   tokenizer.encode('''あンいワ'''			)
      UpperCamelCase__					:    Any          =   tokenizer.encode(''''''							,	prefix_text='''あンいワ'''			)
      UpperCamelCase__					:    List[str]          =   tokenizer.encode('''いワ'''							,	prefix_text='''あン'''			)

      self.assertEqual(tokenizer.decode(lowerCamelCase__			)							,	tokenizer.decode(lowerCamelCase__			)			)
      self.assertEqual(tokenizer.decode(lowerCamelCase__			)							,	tokenizer.decode(lowerCamelCase__			)			)
      self.assertNotEqual(lowerCamelCase__							,	lowerCamelCase__			)
      self.assertNotEqual(lowerCamelCase__							,	lowerCamelCase__			)
      self.assertEqual(x_token_a[1]							,	x_token_a[-1]			)  # SEG token
      self.assertEqual(x_token_a[1]							,	x_token_a[3]			)  # SEG token


     @slow
     def      UpperCAmelCase__	(      self							:     List[str]			)  ->      List[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    List[Any]          =   self.tokenizer_class.from_pretrained('''Tanrei/GPTSAN-japanese'''			)

      UpperCamelCase__					:    Any          =   [['''武田信玄''', '''は、'''], ['''織田信長''', '''の配下の、''']]
      UpperCamelCase__					:    Union[str, Any]          =   tokenizer(lowerCamelCase__							,	padding=lowerCamelCase__			)
      UpperCamelCase__					:    Any          =   tokenizer.batch_encode_plus(lowerCamelCase__							,	padding=lowerCamelCase__			)

      # fmt: off
      UpperCamelCase__					:    Union[str, Any]          =   [[35993, 8640, 25948, 35998, 30647, 35675, 35999, 35999], [35993, 10382, 9868, 35998, 30646, 9459, 30646, 35675]]
      UpperCamelCase__					:    str          =   [[1, 1, 1, 0, 0, 0, 0, 0], [1, 1, 1, 0, 0, 0, 0, 0]]
      UpperCamelCase__					:    List[str]          =   [[1, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1]]
      # fmt: on
      self.assertListEqual(x_token.input_ids							,	lowerCamelCase__			)
      self.assertListEqual(x_token.token_type_ids							,	lowerCamelCase__			)
      self.assertListEqual(x_token.attention_mask							,	lowerCamelCase__			)
      self.assertListEqual(x_token_a.input_ids							,	lowerCamelCase__			)
      self.assertListEqual(x_token_a.token_type_ids							,	lowerCamelCase__			)
      self.assertListEqual(x_token_a.attention_mask							,	lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     int			)  ->      Any:





      '''simple docstring'''






      pass


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      Any:





      '''simple docstring'''






      pass






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


     def __init__(      self							:     Dict			)  ->      str:





      '''simple docstring'''






      super().__init__()
      UpperCamelCase__					:    int          =   nn.Linear(3							,	4			)
      UpperCamelCase__					:    str          =   nn.BatchNormad(4			)
      UpperCamelCase__					:    List[str]          =   nn.Linear(4							,	5			)


     def      UpperCAmelCase__	(      self							:     int							,	lowerCamelCase__							:     Dict			)  ->      Tuple:





      '''simple docstring'''






      return self.lineara(self.batchnorm(self.lineara(lowerCamelCase__			)			)			)



class       __magic_name__	(     __lowerCAmelCase):


     def      UpperCAmelCase__	(      self							:     Union[str, Any]							,	lowerCamelCase__							:     Any							,	*lowerCamelCase__							:     int							,	**lowerCamelCase__							:     List[str]			)  ->      Any:





      '''simple docstring'''






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



class       __magic_name__	(     __lowerCAmelCase):


     def      UpperCAmelCase__	(      self							:     Any							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     Any			)  ->      Any:





      '''simple docstring'''






      return output + 1



class       __magic_name__	(     unittest.TestCase):


     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Tuple          =   ModelForTest()
      UpperCamelCase__					:    Union[str, Any]          =   ModelHook()

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


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      str:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   ModelForTest()
      UpperCamelCase__					:    List[str]          =   ModelHook()

      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__							,	append=lowerCamelCase__			)

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


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





      '''simple docstring'''






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

      UpperCamelCase__					:    Optional[int]          =   PreForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    List[Any]          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	lowerCamelCase__							,	atol=1E-5			)			)

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

      # You need to use the sequential hook to chain two or more hooks
      UpperCamelCase__					:    Tuple          =   SequentialHook(PreForwardHook()							,	PreForwardHook()			)
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)

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


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





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   ModelForTest()
      UpperCamelCase__					:    str          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Any          =   test_model(lowerCamelCase__			)

      UpperCamelCase__					:    int          =   PostForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    Dict          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	output + 1							,	atol=1E-5			)			)

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

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

      UpperCamelCase__					:    Tuple          =   test_model(lowerCamelCase__			)
      assert torch.allclose(lowerCamelCase__							,	output + 2							,	atol=1E-5			)


     def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    List[Any]          =   ModelForTest()
      UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Tuple          =   test_model(lowerCamelCase__			)

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

      UpperCamelCase__					:    Optional[int]          =   True
      UpperCamelCase__					:    Optional[int]          =   test_model(lowerCamelCase__			)
      self.assertFalse(outputa.requires_grad			)


     @require_multi_gpu
     def      UpperCAmelCase__	(      self							:     Any			)  ->      Dict:





      '''simple docstring'''






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


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

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

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

      UpperCamelCase__					:    Dict          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    List[str]          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

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

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

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

      # 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      UpperCAmelCase__	(      self							:     str			)  ->      Optional[int]:





      '''simple docstring'''






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

      # 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__					:    str          =   torch.device(lowerCamelCase__			)
      self.assertEqual(model.batchnorm.running_mean.device							,	lowerCamelCase__			)

      UpperCamelCase__					:    int          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Tuple          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

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

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

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(lowerCamelCase__			)
      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      UpperCAmelCase__	(      self							:     int			)  ->      List[Any]:





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

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

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

      # 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__					:    Union[str, Any]          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Optional[int]          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

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


