





"""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					)						->							Dict:
										super().__init__()
										_snake_case													=			nn.Linear(3			,  4					)
										_snake_case													=			nn.BatchNormad(4					)
										_snake_case													=			nn.Linear(4			,  5					)






					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_					)						->							Tuple:
										return self.lineara(self.batchnorm(self.lineara(lowerCAmelCase_					)					)					)





class      UpperCamelCase_						(					_lowerCamelCase   ):




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  *lowerCAmelCase_			,  **lowerCAmelCase_					)						->							Dict:
										return (args[0] + 1,) + args[1:], kwargs





class      UpperCamelCase_						(					_lowerCamelCase   ):




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							int:
										return output + 1





class      UpperCamelCase_						(					unittest.TestCase   ):




					def 		lowerCAmelCase     (						self					)						->							List[Any]:
										_snake_case													=			ModelForTest()
										_snake_case													=			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 		lowerCAmelCase     (						self					)						->							Dict:
										_snake_case													=			ModelForTest()
										_snake_case													=			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 		lowerCAmelCase     (						self					)						->							Union[str, Any]:
										_snake_case													=			ModelForTest()
										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			test_model(x + 1					)
										_snake_case													=			test_model(x + 2					)

										_snake_case													=			PreForwardHook()
										add_hook_to_module(lowerCAmelCase_			,  lowerCAmelCase_					)
										_snake_case													=			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
										_snake_case													=			PreForwardHook()
										add_hook_to_module(lowerCAmelCase_			,  lowerCAmelCase_					)
										_snake_case													=			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
										_snake_case													=			SequentialHook(PreForwardHook()			,  PreForwardHook()					)
										add_hook_to_module(lowerCAmelCase_			,  lowerCAmelCase_					)

										_snake_case													=			test_model(lowerCAmelCase_					)
										assert torch.allclose(lowerCAmelCase_			,  lowerCAmelCase_			,  atol=1E-5					)




					def 		lowerCAmelCase     (						self					)						->							Any:
										_snake_case													=			ModelForTest()
										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			test_model(lowerCAmelCase_					)

										_snake_case													=			PostForwardHook()
										add_hook_to_module(lowerCAmelCase_			,  lowerCAmelCase_					)
										_snake_case													=			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
										_snake_case													=			PostForwardHook()
										add_hook_to_module(lowerCAmelCase_			,  lowerCAmelCase_					)
										_snake_case													=			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
										_snake_case													=			SequentialHook(PostForwardHook()			,  PostForwardHook()					)
										add_hook_to_module(lowerCAmelCase_			,  lowerCAmelCase_					)

										_snake_case													=			test_model(lowerCAmelCase_					)
										assert torch.allclose(lowerCAmelCase_			,  output + 2			,  atol=1E-5					)




					def 		lowerCAmelCase     (						self					)						->							int:
										_snake_case													=			ModelForTest()
										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			test_model(lowerCAmelCase_					)

										_snake_case													=			PostForwardHook()
										add_hook_to_module(lowerCAmelCase_			,  lowerCAmelCase_					)
										_snake_case													=			test_model(lowerCAmelCase_					)
										self.assertTrue(torch.allclose(lowerCAmelCase_			,  output + 1					)					)
										self.assertTrue(outputa.requires_grad					)

										_snake_case													=			True
										_snake_case													=			test_model(lowerCAmelCase_					)
										self.assertFalse(outputa.requires_grad					)




					@require_multi_gpu
					def 		lowerCAmelCase     (						self					)						->							List[str]:
										_snake_case													=			ModelForTest()
										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device			,  torch.device('cpu'					)					)
										self.assertEqual(model.batchnorm.weight.device			,  torch.device('cpu'					)					)
										self.assertEqual(model.lineara.weight.device			,  torch.device('cpu'					)					)

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

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

										# We can still make a forward pass. The input does not need to be on any particular device
										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			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_					)					)
										_snake_case													=			torch.randn(2			,  3					).to(0					)
										_snake_case													=			model(lowerCAmelCase_					)
										self.assertEqual(output.device			,  torch.device(0					)					)




					def 		lowerCAmelCase     (						self					)						->							Any:
										_snake_case													=			ModelForTest()

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

										# This will move each submodule on different devices
										_snake_case													=			{'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
										_snake_case													=			torch.device(hook_kwargs['execution_device']					)
										self.assertEqual(model.batchnorm.running_mean.device			,  lowerCAmelCase_					)

										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			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
										_snake_case													=			{
										    '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'					)					)

										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			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 		lowerCAmelCase     (						self					)						->							Any:
										_snake_case													=			ModelForTest()

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

										# This will move each submodule on different devices
										_snake_case													=			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
										_snake_case													=			torch.device(lowerCAmelCase_					)
										self.assertEqual(model.batchnorm.running_mean.device			,  lowerCAmelCase_					)

										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			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'					)					)

										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			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 		lowerCAmelCase     (						self					)						->							Optional[Any]:
										_snake_case													=			ModelForTest()

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

										# This will move each submodule on different devices
										_snake_case													=			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
										_snake_case													=			torch.device(lowerCAmelCase_					)
										self.assertEqual(model.batchnorm.running_mean.device			,  lowerCAmelCase_					)

										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			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'					)					)

										_snake_case													=			torch.randn(2			,  3					)
										_snake_case													=			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   lowerCamelCase__       (       UpperCamelCase__     :    list[list[int]]    ,      UpperCamelCase__     :    int    ,      UpperCamelCase__     :    int    ,      UpperCamelCase__     :    list[int]  )     ->	bool:



					'''simple docstring'''




					if graph[path[curr_ind - 1]][next_ver] == 0:
										return False

					# 2. Validate that next vertex is not already in path
					return not any(vertex == next_ver for vertex in path  )





def   lowerCamelCase__       (       UpperCamelCase__     :    list[list[int]]    ,      UpperCamelCase__     :    list[int]    ,      UpperCamelCase__     :    int  )     ->	bool:



					'''simple docstring'''




					if curr_ind == len(UpperCamelCase__  ):
										# return whether path exists between current and starting vertices
										return graph[path[curr_ind - 1]][path[0]] == 1

					# Recursive Step
					for next_ver in range(0    ,      len(UpperCamelCase__  )  ):
										if valid_connection(UpperCamelCase__    ,      UpperCamelCase__    ,      UpperCamelCase__    ,      UpperCamelCase__  ):
															# Insert current vertex  into path as next transition
															_snake_case													=			next_ver
															# Validate created path
															if util_hamilton_cycle(UpperCamelCase__    ,      UpperCamelCase__    ,      curr_ind + 1  ):
																				return True
															# Backtrack
															_snake_case													=			-1
					return False





def   lowerCamelCase__       (       UpperCamelCase__     :    list[list[int]]    ,      UpperCamelCase__     :    int = 0  )     ->	list[int]:



					'''simple docstring'''




					_snake_case													=			[-1] * (len(UpperCamelCase__  ) + 1)
					# initialize start and end of path with starting index
					_snake_case													=			_snake_case													=			start_index
					# evaluate and if we find answer return path either return empty array
					return path if util_hamilton_cycle(UpperCamelCase__    ,      UpperCamelCase__    ,      1  ) else []



