






"""simple docstring"""




import argparse
import json
import os
import sys
import tempfile
import unittest
from argparse import Namespace
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import List, Literal, Optional

import yaml

from transformers import HfArgumentParser, TrainingArguments
from transformers.hf_argparser import make_choice_type_function, string_to_bool


# Since Python 3.10, we can use the builtin `|` operator for Union types
# See PEP 604: https://peps.python.org/pep-0604
_snake_case          =   sys.version_info >= (3, 10)

def 			lowerCAmelCase__     (     UpperCamelCase__=None     ,   UpperCamelCase__=None				):



  '''simple docstring'''
  return field(default_factory=lambda: default     ,   metadata=UpperCamelCase__				)






@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       int
 UpperCamelCase					:       float
 UpperCamelCase					:       str
 UpperCamelCase					:       bool







@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       int			=					42
 UpperCamelCase					:       str			=					field(default='''toto'''		,     metadata={'''help''': '''help message'''}      )







@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       bool			=					False
 UpperCamelCase					:       bool			=					True
 UpperCamelCase					:       Optional[bool]			=					None







class        UpperCamelCase			(	snake_case_      ):
 UpperCamelCase					:       Any			=					'''titi'''
 UpperCamelCase					:       Dict			=					'''toto'''







class        UpperCamelCase			(	snake_case_      ):
 UpperCamelCase					:       str			=					'''titi'''
 UpperCamelCase					:       Union[str, Any]			=					'''toto'''
 UpperCamelCase					:       Optional[int]			=					42







@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       BasicEnum			=					"toto"


 def _lowercase	(     self		:						Dict							)					->       Union[str, Any]:
   _a  :       Union[str, Any]             =	BasicEnum(self.foo							)







@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       MixedTypeEnum			=					"toto"


 def _lowercase	(     self		:						Union[str, Any]							)					->       Optional[int]:
   _a  :       Tuple             =	MixedTypeEnum(self.foo							)







@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       Optional[int]			=					None
 UpperCamelCase					:       Optional[float]			=					field(default=snake_case_		,     metadata={'''help''': '''help message'''}      )
 UpperCamelCase					:       Optional[str]			=					None
 UpperCamelCase					:       Optional[List[str]]			=					list_field(default=[]      )
 UpperCamelCase					:       Optional[List[int]]			=					list_field(default=[]      )







@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       List[int]			=					list_field(default=[]      )
 UpperCamelCase					:       List[int]			=					list_field(default=[1, 2, 3]      )
 UpperCamelCase					:       List[str]			=					list_field(default=['''Hallo''', '''Bonjour''', '''Hello''']      )
 UpperCamelCase					:       List[float]			=					list_field(default=[0.1, 0.2, 0.3]      )







@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       List[int]			=					field()
 UpperCamelCase					:       str			=					field()
 UpperCamelCase					:       BasicEnum			=					field()


 def _lowercase	(     self		:						Optional[Any]							)					->       Optional[int]:
   _a  :       str             =	BasicEnum(self.required_enum							)







@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       int
 UpperCamelCase					:       "BasicEnum"			=					field()
 UpperCamelCase					:       "Optional[bool]"			=					None
 UpperCamelCase					:       "str"			=					field(default='''toto'''		,     metadata={'''help''': '''help message'''}      )
 UpperCamelCase					:       "List[str]"			=					list_field(default=['''Hallo''', '''Bonjour''', '''Hello''']      )


if is_python_no_less_than_3_10:






 @dataclass
 class        UpperCamelCase			:
  UpperCamelCase					:       bool			=					False
  UpperCamelCase					:       bool			=					True
  UpperCamelCase					:       bool | None			=					None







 @dataclass
 class        UpperCamelCase			:
  UpperCamelCase					:       int | None			=					None
  UpperCamelCase					:       float | None			=					field(default=snake_case_		,     metadata={'''help''': '''help message'''}      )
  UpperCamelCase					:       str | None			=					None
  UpperCamelCase					:       list[str] | None			=					list_field(default=[]      )
  UpperCamelCase					:       list[int] | None			=					list_field(default=[]      )






class        UpperCamelCase			(	unittest.TestCase      ):


 def _lowercase	(     self		:						str							,						UpperCAmelCase__		:						argparse.ArgumentParser							,						UpperCAmelCase__		:						argparse.ArgumentParser							)					->       List[Any]:
   self.assertEqual(len(a._actions							)							,						len(b._actions							)							)
   for x, y in zip(a._actions							,						b._actions							):
     _a  :       Optional[int]             =	{k: v for k, v in vars(UpperCAmelCase__							).items() if k != """container"""}
     _a  :       Union[str, Any]             =	{k: v for k, v in vars(UpperCAmelCase__							).items() if k != """container"""}

     # Choices with mixed type have custom function as "type"
     # So we need to compare results directly for equality
     if xx.get("""choices"""							,						UpperCAmelCase__							) and yy.get("""choices"""							,						UpperCAmelCase__							):
       for expected_choice in yy["choices"] + xx["choices"]:
         self.assertEqual(xx["""type"""](UpperCAmelCase__							)							,						yy["""type"""](UpperCAmelCase__							)							)
       del xx["type"], yy["type"]

     self.assertEqual(UpperCAmelCase__							,						UpperCAmelCase__							)


 def _lowercase	(     self		:						Tuple							)					->       Dict:
   _a  :       str             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	argparse.ArgumentParser()
   expected.add_argument("""--foo"""							,						type=UpperCAmelCase__							,						required=UpperCAmelCase__							)
   expected.add_argument("""--bar"""							,						type=UpperCAmelCase__							,						required=UpperCAmelCase__							)
   expected.add_argument("""--baz"""							,						type=UpperCAmelCase__							,						required=UpperCAmelCase__							)
   expected.add_argument("""--flag"""							,						type=UpperCAmelCase__							,						default=UpperCAmelCase__							,						const=UpperCAmelCase__							,						nargs="""?"""							)
   self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	["""--foo""", """1""", """--baz""", """quux""", """--bar""", """0.5"""]
   ((_a)      ,  )  :       Any             =	parser.parse_args_into_dataclasses(UpperCAmelCase__							,						look_for_args_file=UpperCAmelCase__							)
   self.assertFalse(example.flag							)


 def _lowercase	(     self		:						List[Any]							)					->       Optional[int]:
   _a  :       int             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       Optional[int]             =	argparse.ArgumentParser()
   expected.add_argument("""--foo"""							,						default=42							,						type=UpperCAmelCase__							)
   expected.add_argument("""--baz"""							,						default="""toto"""							,						type=UpperCAmelCase__							,						help="""help message"""							)
   self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)


 def _lowercase	(     self		:						Any							)					->       List[str]:
   _a  :       List[Any]             =	argparse.ArgumentParser()
   expected.add_argument("""--foo"""							,						type=UpperCAmelCase__							,						default=UpperCAmelCase__							,						const=UpperCAmelCase__							,						nargs="""?"""							)
   expected.add_argument("""--baz"""							,						type=UpperCAmelCase__							,						default=UpperCAmelCase__							,						const=UpperCAmelCase__							,						nargs="""?"""							)
   # A boolean no_* argument always has to come after its "default: True" regular counter-part
   # and its default must be set to False
   expected.add_argument("""--no_baz"""							,						action="""store_false"""							,						default=UpperCAmelCase__							,						dest="""baz"""							)
   expected.add_argument("""--opt"""							,						type=UpperCAmelCase__							,						default=UpperCAmelCase__							)

   _a  :       Any             =	[WithDefaultBoolExample]
   if is_python_no_less_than_3_10:
     dataclass_types.append(UpperCAmelCase__							)

   for dataclass_type in dataclass_types:
     _a  :       List[Any]             =	HfArgumentParser(UpperCAmelCase__							)
     self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)

     _a  :       int             =	parser.parse_args([]							)
     self.assertEqual(UpperCAmelCase__							,						Namespace(foo=UpperCAmelCase__							,						baz=UpperCAmelCase__							,						opt=UpperCAmelCase__							)							)

     _a  :       Tuple             =	parser.parse_args(["""--foo""", """--no_baz"""]							)
     self.assertEqual(UpperCAmelCase__							,						Namespace(foo=UpperCAmelCase__							,						baz=UpperCAmelCase__							,						opt=UpperCAmelCase__							)							)

     _a  :       Optional[int]             =	parser.parse_args(["""--foo""", """--baz"""]							)
     self.assertEqual(UpperCAmelCase__							,						Namespace(foo=UpperCAmelCase__							,						baz=UpperCAmelCase__							,						opt=UpperCAmelCase__							)							)

     _a  :       List[Any]             =	parser.parse_args(["""--foo""", """True""", """--baz""", """True""", """--opt""", """True"""]							)
     self.assertEqual(UpperCAmelCase__							,						Namespace(foo=UpperCAmelCase__							,						baz=UpperCAmelCase__							,						opt=UpperCAmelCase__							)							)

     _a  :       Tuple             =	parser.parse_args(["""--foo""", """False""", """--baz""", """False""", """--opt""", """False"""]							)
     self.assertEqual(UpperCAmelCase__							,						Namespace(foo=UpperCAmelCase__							,						baz=UpperCAmelCase__							,						opt=UpperCAmelCase__							)							)


 def _lowercase	(     self		:						Dict							)					->       Optional[int]:
   _a  :       Optional[Any]             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       Tuple             =	argparse.ArgumentParser()
   expected.add_argument(
       """--foo"""							,						default="""toto"""							,						choices=["""titi""", """toto""", 42]							,						type=make_choice_type_function(["""titi""", """toto""", 42]							)							,						)
   self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       str             =	parser.parse_args([]							)
   self.assertEqual(args.foo							,						"""toto"""							)
   _a  :       Optional[Any]             =	parser.parse_args_into_dataclasses([]							)[0]
   self.assertEqual(enum_ex.foo							,						MixedTypeEnum.toto							)

   _a  :       Dict             =	parser.parse_args(["""--foo""", """titi"""]							)
   self.assertEqual(args.foo							,						"""titi"""							)
   _a  :       Optional[Any]             =	parser.parse_args_into_dataclasses(["""--foo""", """titi"""]							)[0]
   self.assertEqual(enum_ex.foo							,						MixedTypeEnum.titi							)

   _a  :       List[Any]             =	parser.parse_args(["""--foo""", """42"""]							)
   self.assertEqual(args.foo							,						42							)
   _a  :       Dict             =	parser.parse_args_into_dataclasses(["""--foo""", """42"""]							)[0]
   self.assertEqual(enum_ex.foo							,						MixedTypeEnum.fourtytwo							)


 def _lowercase	(     self		:						int							)					->       Any:






   @dataclass
   class        UpperCamelCase			:
    UpperCamelCase					:       Literal["titi", "toto", 42]			=					"toto"







   _a  :       Tuple             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       Dict             =	argparse.ArgumentParser()
   expected.add_argument(
       """--foo"""							,						default="""toto"""							,						choices=("""titi""", """toto""", 42)							,						type=make_choice_type_function(["""titi""", """toto""", 42]							)							,						)
   self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Tuple             =	parser.parse_args([]							)
   self.assertEqual(args.foo							,						"""toto"""							)

   _a  :       Optional[Any]             =	parser.parse_args(["""--foo""", """titi"""]							)
   self.assertEqual(args.foo							,						"""titi"""							)

   _a  :       List[str]             =	parser.parse_args(["""--foo""", """42"""]							)
   self.assertEqual(args.foo							,						42							)


 def _lowercase	(     self		:						Tuple							)					->       str:
   _a  :       Tuple             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       List[str]             =	argparse.ArgumentParser()
   expected.add_argument("""--foo_int"""							,						nargs="""+"""							,						default=[]							,						type=UpperCAmelCase__							)
   expected.add_argument("""--bar_int"""							,						nargs="""+"""							,						default=[1, 2, 3]							,						type=UpperCAmelCase__							)
   expected.add_argument("""--foo_str"""							,						nargs="""+"""							,						default=["""Hallo""", """Bonjour""", """Hello"""]							,						type=UpperCAmelCase__							)
   expected.add_argument("""--foo_float"""							,						nargs="""+"""							,						default=[0.1, 0.2, 0.3]							,						type=UpperCAmelCase__							)

   self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Tuple             =	parser.parse_args([]							)
   self.assertEqual(
       UpperCAmelCase__							,						Namespace(foo_int=[]							,						bar_int=[1, 2, 3]							,						foo_str=["""Hallo""", """Bonjour""", """Hello"""]							,						foo_float=[0.1, 0.2, 0.3]							)							,						)

   _a  :       Any             =	parser.parse_args("""--foo_int 1 --bar_int 2 3 --foo_str a b c --foo_float 0.1 0.7""".split()							)
   self.assertEqual(UpperCAmelCase__							,						Namespace(foo_int=[1]							,						bar_int=[2, 3]							,						foo_str=["""a""", """b""", """c"""]							,						foo_float=[0.1, 0.7]							)							)


 def _lowercase	(     self		:						List[str]							)					->       Optional[Any]:
   _a  :       Any             =	argparse.ArgumentParser()
   expected.add_argument("""--foo"""							,						default=UpperCAmelCase__							,						type=UpperCAmelCase__							)
   expected.add_argument("""--bar"""							,						default=UpperCAmelCase__							,						type=UpperCAmelCase__							,						help="""help message"""							)
   expected.add_argument("""--baz"""							,						default=UpperCAmelCase__							,						type=UpperCAmelCase__							)
   expected.add_argument("""--ces"""							,						nargs="""+"""							,						default=[]							,						type=UpperCAmelCase__							)
   expected.add_argument("""--des"""							,						nargs="""+"""							,						default=[]							,						type=UpperCAmelCase__							)

   _a  :       List[Any]             =	[OptionalExample]
   if is_python_no_less_than_3_10:
     dataclass_types.append(UpperCAmelCase__							)

   for dataclass_type in dataclass_types:
     _a  :       List[str]             =	HfArgumentParser(UpperCAmelCase__							)

     self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)

     _a  :       str             =	parser.parse_args([]							)
     self.assertEqual(UpperCAmelCase__							,						Namespace(foo=UpperCAmelCase__							,						bar=UpperCAmelCase__							,						baz=UpperCAmelCase__							,						ces=[]							,						des=[]							)							)

     _a  :       str             =	parser.parse_args("""--foo 12 --bar 3.14 --baz 42 --ces a b c --des 1 2 3""".split()							)
     self.assertEqual(UpperCAmelCase__							,						Namespace(foo=12							,						bar=3.1_4							,						baz="""42"""							,						ces=["""a""", """b""", """c"""]							,						des=[1, 2, 3]							)							)


 def _lowercase	(     self		:						int							)					->       Tuple:
   _a  :       Optional[int]             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       List[Any]             =	argparse.ArgumentParser()
   expected.add_argument("""--required_list"""							,						nargs="""+"""							,						type=UpperCAmelCase__							,						required=UpperCAmelCase__							)
   expected.add_argument("""--required_str"""							,						type=UpperCAmelCase__							,						required=UpperCAmelCase__							)
   expected.add_argument(
       """--required_enum"""							,						type=make_choice_type_function(["""titi""", """toto"""]							)							,						choices=["""titi""", """toto"""]							,						required=UpperCAmelCase__							,						)
   self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)


 def _lowercase	(     self		:						int							)					->       Optional[int]:
   _a  :       Any             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       Any             =	argparse.ArgumentParser()
   expected.add_argument("""--foo"""							,						type=UpperCAmelCase__							,						required=UpperCAmelCase__							)
   expected.add_argument(
       """--required_enum"""							,						type=make_choice_type_function(["""titi""", """toto"""]							)							,						choices=["""titi""", """toto"""]							,						required=UpperCAmelCase__							,						)
   expected.add_argument("""--opt"""							,						type=UpperCAmelCase__							,						default=UpperCAmelCase__							)
   expected.add_argument("""--baz"""							,						default="""toto"""							,						type=UpperCAmelCase__							,						help="""help message"""							)
   expected.add_argument("""--foo_str"""							,						nargs="""+"""							,						default=["""Hallo""", """Bonjour""", """Hello"""]							,						type=UpperCAmelCase__							)
   self.argparsersEqual(UpperCAmelCase__							,						UpperCAmelCase__							)


 def _lowercase	(     self		:						Any							)					->       Optional[Any]:
   _a  :       Optional[int]             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	{
       """foo""": 12,
       """bar""": 3.1_4,
       """baz""": """42""",
       """flag""": True,
   }

   _a  :       Union[str, Any]             =	parser.parse_dict(UpperCAmelCase__							)[0]
   _a  :       str             =	BasicExample(**UpperCAmelCase__							)
   self.assertEqual(UpperCAmelCase__							,						UpperCAmelCase__							)


 def _lowercase	(     self		:						List[Any]							)					->       Union[str, Any]:
   _a  :       Any             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       Tuple             =	{
       """foo""": 12,
       """bar""": 3.1_4,
       """baz""": """42""",
       """flag""": True,
       """extra""": 42,
   }

   self.assertRaises(UpperCAmelCase__							,						parser.parse_dict							,						UpperCAmelCase__							,						allow_extra_keys=UpperCAmelCase__							)


 def _lowercase	(     self		:						Dict							)					->       List[str]:
   _a  :       List[str]             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       str             =	{
       """foo""": 12,
       """bar""": 3.1_4,
       """baz""": """42""",
       """flag""": True,
   }
   with tempfile.TemporaryDirectory() as tmp_dir:
     _a  :       str             =	os.path.join(UpperCAmelCase__							,						"""temp_json"""							)
     os.mkdir(UpperCAmelCase__							)
     with open(temp_local_path + """.json"""							,						"""w+"""							) as f:
       json.dump(UpperCAmelCase__							,						UpperCAmelCase__							)
     _a  :       Union[str, Any]             =	parser.parse_yaml_file(Path(temp_local_path + """.json"""							)							)[0]

   _a  :       Any             =	BasicExample(**UpperCAmelCase__							)
   self.assertEqual(UpperCAmelCase__							,						UpperCAmelCase__							)


 def _lowercase	(     self		:						Union[str, Any]							)					->       Optional[Any]:
   _a  :       Dict             =	HfArgumentParser(UpperCAmelCase__							)

   _a  :       Dict             =	{
       """foo""": 12,
       """bar""": 3.1_4,
       """baz""": """42""",
       """flag""": True,
   }
   with tempfile.TemporaryDirectory() as tmp_dir:
     _a  :       Dict             =	os.path.join(UpperCAmelCase__							,						"""temp_yaml"""							)
     os.mkdir(UpperCAmelCase__							)
     with open(temp_local_path + """.yaml"""							,						"""w+"""							) as f:
       yaml.dump(UpperCAmelCase__							,						UpperCAmelCase__							)
     _a  :       Any             =	parser.parse_yaml_file(Path(temp_local_path + """.yaml"""							)							)[0]
   _a  :       Union[str, Any]             =	BasicExample(**UpperCAmelCase__							)
   self.assertEqual(UpperCAmelCase__							,						UpperCAmelCase__							)







 def _lowercase	(     self		:						int							)					->       Optional[Any]:
   _a  :       List[str]             =	HfArgumentParser(UpperCAmelCase__							)
   self.assertIsNotNone(UpperCAmelCase__							)









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


 def __init__(     self		:						Union[str, Any]							)					->       int:
   super().__init__()
   _a  :       Optional[Any]             =	nn.Linear(3							,						4							)
   _a  :       Tuple             =	nn.BatchNormad(4							)
   _a  :       Dict             =	nn.Linear(4							,						5							)







 def _lowercase	(     self		:						Optional[int]							,						UpperCAmelCase__		:						List[str]							)					->       int:
   return self.lineara(self.batchnorm(self.lineara(UpperCAmelCase__							)							)							)







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Any							,						UpperCAmelCase__		:						Any							,						*UpperCAmelCase__		:						List[str]							,						**UpperCAmelCase__		:						Optional[int]							)					->       Optional[int]:
   return (args[0] + 1,) + args[1:], kwargs







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Optional[Any]							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Optional[Any]							)					->       List[str]:
   return output + 1







class        UpperCamelCase			(	unittest.TestCase      ):


 def _lowercase	(     self		:						Dict							)					->       str:
   _a  :       List[Any]             =	ModelForTest()
   _a  :       str             =	ModelHook()

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


 def _lowercase	(     self		:						Optional[int]							)					->       Optional[int]:
   _a  :       Dict             =	ModelForTest()
   _a  :       Dict             =	ModelHook()

   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							,						append=UpperCAmelCase__							)

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


 def _lowercase	(     self		:						Dict							)					->       int:
   _a  :       str             =	ModelForTest()
   _a  :       List[Any]             =	torch.randn(2							,						3							)
   _a  :       Optional[Any]             =	test_model(x + 1							)
   _a  :       str             =	test_model(x + 2							)

   _a  :       Union[str, Any]             =	PreForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       Tuple             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						UpperCAmelCase__							,						atol=1E-5							)							)

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

   # You need to use the sequential hook to chain two or more hooks
   _a  :       int             =	SequentialHook(PreForwardHook()							,						PreForwardHook()							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Tuple             =	test_model(UpperCAmelCase__							)
   assert torch.allclose(UpperCAmelCase__							,						UpperCAmelCase__							,						atol=1E-5							)


 def _lowercase	(     self		:						Tuple							)					->       int:
   _a  :       Tuple             =	ModelForTest()
   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       Optional[int]             =	test_model(UpperCAmelCase__							)

   _a  :       int             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       List[str]             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						output + 1							,						atol=1E-5							)							)

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

   # You need to use the sequential hook to chain two or more hooks
   _a  :       Any             =	SequentialHook(PostForwardHook()							,						PostForwardHook()							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Optional[int]             =	test_model(UpperCAmelCase__							)
   assert torch.allclose(UpperCAmelCase__							,						output + 2							,						atol=1E-5							)


 def _lowercase	(     self		:						Dict							)					->       Optional[Any]:
   _a  :       Any             =	ModelForTest()
   _a  :       List[Any]             =	torch.randn(2							,						3							)
   _a  :       Dict             =	test_model(UpperCAmelCase__							)

   _a  :       Any             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       List[str]             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						output + 1							)							)
   self.assertTrue(outputa.requires_grad							)

   _a  :       Any             =	True
   _a  :       Union[str, Any]             =	test_model(UpperCAmelCase__							)
   self.assertFalse(outputa.requires_grad							)


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


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

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

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

   _a  :       int             =	torch.randn(2							,						3							)
   _a  :       str             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

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

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

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

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

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

   # 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
   _a  :       Dict             =	torch.device(UpperCAmelCase__							)
   self.assertEqual(model.batchnorm.running_mean.device							,						UpperCAmelCase__							)

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

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

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

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

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(UpperCAmelCase__							)
   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:
   _a  :       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
   _a  :       str             =	0 if torch.cuda.is_available() else """cpu"""
   attach_align_device_hook(
       UpperCAmelCase__							,						execution_device=UpperCAmelCase__							,						offload=UpperCAmelCase__							,						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
   _a  :       Union[str, Any]             =	torch.device(UpperCAmelCase__							)
   self.assertEqual(model.batchnorm.running_mean.device							,						UpperCAmelCase__							)

   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       int             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

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

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

   _a  :       Any             =	torch.randn(2							,						3							)
   _a  :       int             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						UpperCAmelCase__							)

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

