


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



     """simple docstring"""






     def __init__( self   ):

            '''simple docstring'''

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







     def a	( self      ,	snake_case__   ):

            '''simple docstring'''

            return self.lineara(self.batchnorm(self.lineara(snake_case__   )   )   )
class 	UpperCamelCase__					( SCREAMING_SNAKE_CASE_	):



     """simple docstring"""






     def a	( self      ,	snake_case__      ,	*snake_case__      ,	**snake_case__   ):

            '''simple docstring'''

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



     """simple docstring"""






     def a	( self      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            return output + 1
class 	UpperCamelCase__					( unittest.TestCase	):



     """simple docstring"""






     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[Any]   			=       ModelForTest()
            _lowerCAmelCase    :							Any   			=       ModelHook()

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







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       ModelForTest()
            _lowerCAmelCase    :							Dict   			=       ModelHook()

            add_hook_to_module(snake_case__      ,	snake_case__   )
            add_hook_to_module(snake_case__      ,	snake_case__      ,	append=snake_case__   )

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







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       ModelForTest()
            _lowerCAmelCase    :							Union[str, Any]   			=       torch.randn(2      ,	3   )
            _lowerCAmelCase    :							Tuple   			=       test_model(x + 1   )
            _lowerCAmelCase    :							Dict   			=       test_model(x + 2   )

            _lowerCAmelCase    :							List[str]   			=       PreForwardHook()
            add_hook_to_module(snake_case__      ,	snake_case__   )
            _lowerCAmelCase    :							Dict   			=       test_model(snake_case__   )
            self.assertTrue(torch.allclose(snake_case__      ,	snake_case__      ,	atol=1E-5   )   )

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

            # You need to use the sequential hook to chain two or more hooks
            _lowerCAmelCase    :							str   			=       SequentialHook(PreForwardHook()      ,	PreForwardHook()   )
            add_hook_to_module(snake_case__      ,	snake_case__   )

            _lowerCAmelCase    :							Dict   			=       test_model(snake_case__   )
            assert torch.allclose(snake_case__      ,	snake_case__      ,	atol=1E-5   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       ModelForTest()
            _lowerCAmelCase    :							List[Any]   			=       torch.randn(2      ,	3   )
            _lowerCAmelCase    :							List[Any]   			=       test_model(snake_case__   )

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

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

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

            _lowerCAmelCase    :							List[Any]   			=       test_model(snake_case__   )
            assert torch.allclose(snake_case__      ,	output + 2      ,	atol=1E-5   )







     def a	( self   ):

            '''simple docstring'''

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

            _lowerCAmelCase    :							List[Any]   			=       PostForwardHook()
            add_hook_to_module(snake_case__      ,	snake_case__   )
            _lowerCAmelCase    :							Optional[int]   			=       test_model(snake_case__   )
            self.assertTrue(torch.allclose(snake_case__      ,	output + 1   )   )
            self.assertTrue(outputa.requires_grad   )

            _lowerCAmelCase    :							Dict   			=       True
            _lowerCAmelCase    :							List[str]   			=       test_model(snake_case__   )
            self.assertFalse(outputa.requires_grad   )







     @require_multi_gpu
     def a	( self   ):

            '''simple docstring'''

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







     def a	( self   ):

            '''simple docstring'''

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

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

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

            _lowerCAmelCase    :							Tuple   			=       torch.randn(2      ,	3   )
            _lowerCAmelCase    :							List[str]   			=       model(snake_case__   )
            self.assertEqual(output.device      ,	snake_case__   )

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

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

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

            _lowerCAmelCase    :							Optional[Any]   			=       torch.randn(2      ,	3   )
            _lowerCAmelCase    :							List[Any]   			=       model(snake_case__   )
            self.assertEqual(output.device      ,	snake_case__   )

            # 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 a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							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
            _lowerCAmelCase    :							Union[str, Any]   			=       0 if torch.cuda.is_available() else 'cpu'
            attach_align_device_hook(snake_case__      ,	execution_device=snake_case__      ,	offload=snake_case__   )

            # 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
            _lowerCAmelCase    :							int   			=       torch.device(snake_case__   )
            self.assertEqual(model.batchnorm.running_mean.device      ,	snake_case__   )

            _lowerCAmelCase    :							Dict   			=       torch.randn(2      ,	3   )
            _lowerCAmelCase    :							Tuple   			=       model(snake_case__   )
            self.assertEqual(output.device      ,	snake_case__   )

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

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

            _lowerCAmelCase    :							List[Any]   			=       torch.randn(2      ,	3   )
            _lowerCAmelCase    :							Any   			=       model(snake_case__   )
            self.assertEqual(output.device      ,	snake_case__   )

            # Removing hooks loads back the weights in the model.
            remove_hook_from_submodules(snake_case__   )
            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 a	( self   ):

            '''simple docstring'''

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

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

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

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

            _lowerCAmelCase    :							List[str]   			=       torch.randn(2      ,	3   )
            _lowerCAmelCase    :							int   			=       model(snake_case__   )
            self.assertEqual(output.device      ,	snake_case__   )

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







import socket






def 			lowercase				():







       """simple docstring"""







       _lowerCAmelCase    :							Tuple   			=       socket.socket(socket.AF_INET   ,						socket.SOCK_STREAM					)
       _lowerCAmelCase    :							Optional[int]   			=       socket.gethostname()
       _lowerCAmelCase    :							Any   			=       1_2_3_1_2

       sock.connect((host, port)					)
       sock.send(B'Hello server!'					)

       with open('Received_file'   ,						'wb'					) as out_file:
              print('File opened'					)
              print('Receiving data...'					)
              while True:
                     _lowerCAmelCase    :							Union[str, Any]   			=       sock.recv(1_0_2_4					)
                     if not data:
                            break
                     out_file.write(_A					)

       print('Successfully received the file'					)
       sock.close()
       print('Connection closed'					)


if __name__ == "__main__":
       main()

