






import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu


class 							__a							(  nn.Module						):





		'''simple docstring'''



		def __init__(						self       ):
						super().__init__()
						SCREAMING_SNAKE_CASE_   :	int													=   nn.Linear(3 ,    4       )
						SCREAMING_SNAKE_CASE_   :	Tuple													=   nn.BatchNormad(4       )
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   nn.Linear(4 ,    5       )






		def   __snake_case			(						self ,    UpperCamelCase__       ):
						return self.lineara(self.batchnorm(self.lineara(UpperCamelCase__       )       )       )






class 							__a							(  __A						):





		'''simple docstring'''



		def   __snake_case			(						self ,    UpperCamelCase__ ,    *UpperCamelCase__ ,    **UpperCamelCase__       ):
						return (args[0] + 1,) + args[1:], kwargs






class 							__a							(  __A						):





		'''simple docstring'''



		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__       ):
						return output + 1






class 							__a							(  unittest.TestCase						):





		'''simple docstring'''



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

						SCREAMING_SNAKE_CASE_   :	List[Any]													=   PreForwardHook()
						add_hook_to_module(UpperCamelCase__ ,    UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	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
						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   PreForwardHook()
						add_hook_to_module(UpperCamelCase__ ,    UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	List[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
						SCREAMING_SNAKE_CASE_   :	Any													=   SequentialHook(PreForwardHook() ,    PreForwardHook()       )
						add_hook_to_module(UpperCamelCase__ ,    UpperCamelCase__       )

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


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Tuple													=   ModelForTest()
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   torch.randn(2 ,    3       )
						SCREAMING_SNAKE_CASE_   :	Tuple													=   test_model(UpperCamelCase__       )

						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   PostForwardHook()
						add_hook_to_module(UpperCamelCase__ ,    UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	Any													=   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
						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   PostForwardHook()
						add_hook_to_module(UpperCamelCase__ ,    UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	int													=   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
						SCREAMING_SNAKE_CASE_   :	Tuple													=   SequentialHook(PostForwardHook() ,    PostForwardHook()       )
						add_hook_to_module(UpperCamelCase__ ,    UpperCamelCase__       )

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


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   ModelForTest()
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   torch.randn(2 ,    3       )
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   test_model(UpperCamelCase__       )

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

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


		@require_multi_gpu
		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Tuple													=   ModelForTest()
						# Everything is on CPU
						self.assertEqual(model.lineara.weight.device ,    torch.device('cpu'       )       )
						self.assertEqual(model.batchnorm.weight.device ,    torch.device('cpu'       )       )
						self.assertEqual(model.lineara.weight.device ,    torch.device('cpu'       )       )

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

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

						# We can still make a forward pass. The input does not need to be on any particular device
						SCREAMING_SNAKE_CASE_   :	int													=   torch.randn(2 ,    3       )
						SCREAMING_SNAKE_CASE_   :	Dict													=   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__       )       )
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   torch.randn(2 ,    3       ).to(0       )
						SCREAMING_SNAKE_CASE_   :	int													=   model(UpperCamelCase__       )
						self.assertEqual(output.device ,    torch.device(0       )       )


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	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
						SCREAMING_SNAKE_CASE_   :	str													=   {'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
						SCREAMING_SNAKE_CASE_   :	Dict													=   torch.device(hook_kwargs['execution_device']       )
						self.assertEqual(model.batchnorm.running_mean.device ,    UpperCamelCase__       )

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

						# Now test with buffers included in the offload
						SCREAMING_SNAKE_CASE_   :	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'       )       )

						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   torch.randn(2 ,    3       )
						SCREAMING_SNAKE_CASE_   :	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   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Union[str, 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
						SCREAMING_SNAKE_CASE_   :	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
						SCREAMING_SNAKE_CASE_   :	Any													=   torch.device(UpperCamelCase__       )
						self.assertEqual(model.batchnorm.running_mean.device ,    UpperCamelCase__       )

						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   torch.randn(2 ,    3       )
						SCREAMING_SNAKE_CASE_   :	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'       )       )

						SCREAMING_SNAKE_CASE_   :	int													=   torch.randn(2 ,    3       )
						SCREAMING_SNAKE_CASE_   :	List[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'       )       )






		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Any													=   ModelForTest()

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

						# This will move each submodule on different devices
						SCREAMING_SNAKE_CASE_   :	Union[str, 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
						SCREAMING_SNAKE_CASE_   :	List[str]													=   torch.device(UpperCamelCase__       )
						self.assertEqual(model.batchnorm.running_mean.device ,    UpperCamelCase__       )

						SCREAMING_SNAKE_CASE_   :	Any													=   torch.randn(2 ,    3       )
						SCREAMING_SNAKE_CASE_   :	List[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'       )       )

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

						SCREAMING_SNAKE_CASE_   :	Tuple													=   torch.randn(2 ,    3       )
						SCREAMING_SNAKE_CASE_   :	Dict													=   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'       )       )







import colorsys

from PIL import Image  # type: ignore



def 					_lowerCamelCase(			lowerCAmelCase__				:		float    ,    lowerCAmelCase__				:		float    ,    lowerCAmelCase__				:		int      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Any													=   x
				SCREAMING_SNAKE_CASE_   :	Tuple													=   y
				for step in range(lowerCAmelCase__      ):  # noqa: B007
								SCREAMING_SNAKE_CASE_   :	str													=   a * a - b * b + x
								SCREAMING_SNAKE_CASE_   :	Any													=   2 * a * b + y
								SCREAMING_SNAKE_CASE_   :	int													=   a_new

								# divergence happens for all complex number with an absolute value
								# greater than 4
								if a * a + b * b > 4:
												break
				return step / (max_step - 1)



def 					_lowerCamelCase(			lowerCAmelCase__				:		float      ):





				'''simple docstring'''







				if distance == 1:
								return (0, 0, 0)
				else:
								return (255, 255, 255)



def 					_lowerCamelCase(			lowerCAmelCase__				:		float      ):





				'''simple docstring'''







				if distance == 1:
								return (0, 0, 0)
				else:
								return tuple(round(i * 255      ) for i in colorsys.hsv_to_rgb(lowerCAmelCase__    ,    1    ,    1      )      )



def 					_lowerCamelCase(			lowerCAmelCase__				:		int = 800    ,    lowerCAmelCase__				:		int = 600    ,    lowerCAmelCase__				:		float = -0.6    ,    lowerCAmelCase__				:		float = 0    ,    lowerCAmelCase__				:		float = 3.2    ,    lowerCAmelCase__				:		int = 50    ,    lowerCAmelCase__				:		bool = True    ,    ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   Image.new('RGB'    ,    (image_width, image_height)      )
				SCREAMING_SNAKE_CASE_   :	List[str]													=   img.load()

				# loop through the image-coordinates
				for image_x in range(lowerCAmelCase__      ):
								for image_y in range(lowerCAmelCase__      ):
												# determine the figure-coordinates based on the image-coordinates
												SCREAMING_SNAKE_CASE_   :	List[str]													=   figure_width / image_width * image_height
												SCREAMING_SNAKE_CASE_   :	List[Any]													=   figure_center_x + (image_x / image_width - 0.5) * figure_width
												SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   figure_center_y + (image_y / image_height - 0.5) * figure_height

												SCREAMING_SNAKE_CASE_   :	str													=   get_distance(lowerCAmelCase__    ,    lowerCAmelCase__    ,    lowerCAmelCase__      )

												# color the corresponding pixel based on the selected coloring-function
												if use_distance_color_coding:
																SCREAMING_SNAKE_CASE_   :	int													=   get_color_coded_rgb(lowerCAmelCase__      )
												else:
																SCREAMING_SNAKE_CASE_   :	Dict													=   get_black_and_white_rgb(lowerCAmelCase__      )

				return img


if __name__ == "__main__":
				import doctest

				doctest.testmod()

				# colored version, full figure
				A        =			get_image()

				# uncomment for colored version, different section, zoomed in
				# img = get_image(figure_center_x = -0.6, figure_center_y = -0.4,
				# figure_width = 0.8)

				# uncomment for black and white version, full figure
				# img = get_image(use_distance_color_coding = False)

				# uncomment to save the image
				# img.save("mandelbrot.png")

				img.show()