






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


							def __init__(			self      :						int							):
													'''simple docstring'''





													super().__init__()
													A__						:  int		  =  nn.Linear(3			,    4							)
													A__						:  str		  =  nn.BatchNormad(4							)
													A__						:  Tuple		  =  nn.Linear(4			,    5							)


							def 						_UpperCamelCase       (			self      :						Dict			,    snake_case      :						Optional[Any]							):
													'''simple docstring'''





													return self.lineara(self.batchnorm(self.lineara(__UpperCamelCase							)							)							)






class 							__SCREAMING_SNAKE_CASE						(				UpperCamelCase		):


							def 						_UpperCamelCase       (			self      :						Dict			,    snake_case      :						Union[str, Any]			,    *snake_case      :						Optional[Any]			,    **snake_case      :						str							):
													'''simple docstring'''





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






class 							__SCREAMING_SNAKE_CASE						(				UpperCamelCase		):


							def 						_UpperCamelCase       (			self      :						Dict			,    snake_case      :						int			,    snake_case      :						int							):
													'''simple docstring'''





													return output + 1






class 							__SCREAMING_SNAKE_CASE						(				unittest.TestCase		):


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





													A__						:  Union[str, Any]		  =  ModelForTest()
													A__						:  Tuple		  =  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 						_UpperCamelCase       (			self      :						List[str]							):
													'''simple docstring'''





													A__						:  Union[str, Any]		  =  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 						_UpperCamelCase       (			self      :						int							):
													'''simple docstring'''





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

													A__						:  Optional[int]		  =  PreForwardHook()
													add_hook_to_module(__UpperCamelCase			,    __UpperCamelCase							)
													A__						:  Optional[int]		  =  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__						:  Union[str, Any]		  =  PreForwardHook()
													add_hook_to_module(__UpperCamelCase			,    __UpperCamelCase							)
													A__						:  Union[str, Any]		  =  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__						:  Dict		  =  SequentialHook(PreForwardHook()			,    PreForwardHook()							)
													add_hook_to_module(__UpperCamelCase			,    __UpperCamelCase							)

													A__						:  int		  =  test_model(__UpperCamelCase							)
													assert torch.allclose(__UpperCamelCase			,    __UpperCamelCase			,    atol=1e-5							)


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





													A__						:  int		  =  ModelForTest()
													A__						:  Any		  =  torch.randn(2			,    3							)
													A__						:  Tuple		  =  test_model(__UpperCamelCase							)

													A__						:  Optional[Any]		  =  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__						:  Optional[int]		  =  PostForwardHook()
													add_hook_to_module(__UpperCamelCase			,    __UpperCamelCase							)
													A__						:  Union[str, Any]		  =  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__						:  Optional[int]		  =  SequentialHook(PostForwardHook()			,    PostForwardHook()							)
													add_hook_to_module(__UpperCamelCase			,    __UpperCamelCase							)

													A__						:  Union[str, Any]		  =  test_model(__UpperCamelCase							)
													assert torch.allclose(__UpperCamelCase			,    output + 2			,    atol=1e-5							)


							def 						_UpperCamelCase       (			self      :						str							):
													'''simple docstring'''





													A__						:  int		  =  ModelForTest()
													A__						:  Dict		  =  torch.randn(2			,    3							)
													A__						:  int		  =  test_model(__UpperCamelCase							)

													A__						:  Optional[int]		  =  PostForwardHook()
													add_hook_to_module(__UpperCamelCase			,    __UpperCamelCase							)
													A__						:  int		  =  test_model(__UpperCamelCase							)
													self.assertTrue(torch.allclose(__UpperCamelCase			,    output + 1							)							)
													self.assertTrue(outputa.requires_grad							)

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


							@require_multi_gpu
							def 						_UpperCamelCase       (			self      :						Optional[int]							):
													'''simple docstring'''





													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__						:  Dict		  =  torch.randn(2			,    3							)
													A__						:  Union[str, 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__						:  List[str]		  =  model(__UpperCamelCase							)
													self.assertEqual(output.device			,    torch.device(0							)							)


							def 						_UpperCamelCase       (			self      :						int							):
													'''simple docstring'''





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

													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_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__						:  Tuple		  =  {
													    """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__						:  Optional[int]		  =  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_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							):
													'''simple docstring'''





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

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


							def 						_UpperCamelCase       (			self      :						int							):
													'''simple docstring'''





													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__						:  Optional[Any]		  =  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__						:  Tuple		  =  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__						:  List[str]		  =  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_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"""							)							)














"""simple docstring"""



def      _lowerCAmelCase		(							UpperCAmelCase__       :					int,						UpperCAmelCase__       :					int      )		->str:
						if a < 0 or b < 0:
												raise ValueError("""the value of both inputs must be positive"""      )

						A__						:  str		  =  str(bin(UpperCAmelCase__      )      )[2:]  # remove the leading "0b"
						A__						:  List[Any]		  =  str(bin(UpperCAmelCase__      )      )[2:]  # remove the leading "0b"

						A__						:  List[str]		  =  max(len(UpperCAmelCase__      ),						len(UpperCAmelCase__      )      )

						return "0b" + "".join(
						    str(int(char_a == """1""" and char_b == """1"""      )      )
						    for char_a, char_b in zip(a_binary.zfill(UpperCAmelCase__      ),						b_binary.zfill(UpperCAmelCase__      )      )      )


if __name__ == "__main__":
						import doctest

						doctest.testmod()






