





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



     '''simple docstring'''






     def __init__(					self	):
       super().__init__()
       _snake_case     		=				nn.Linear(3						,   4	)
       _snake_case     		=				nn.BatchNormad(4	)
       _snake_case     		=				nn.Linear(4						,   5	)






     def 	UpperCamelCase(					self						,   lowerCamelCase	):
       return self.lineara(self.batchnorm(self.lineara(lowerCamelCase	)	)	)







class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self						,   lowerCamelCase						,   *lowerCamelCase						,   **lowerCamelCase	):
       return (args[0] + 1,) + args[1:], kwargs







class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase	):
       return output + 1







class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _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"	)	)









'''simple docstring'''



from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__magic_name__   :						Any    =			logging.get_logger(__name__)

__magic_name__   :						Dict    =			{
    """google/bit-50""": """https://huggingface.co/google/bit-50/resolve/main/config.json""",
}


class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase					,    __UpperCamelCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						List[str]      =   '''bit'''
     UpperCAmelCase__       :						Union[str, Any]      =   ['''preactivation''', '''bottleneck''']
     UpperCAmelCase__       :						int      =   ['''SAME''', '''VALID''']

     def __init__(					self						,   lowerCamelCase=3						,   lowerCamelCase=64						,   lowerCamelCase=[256, 512, 1_024, 2_048]						,   lowerCamelCase=[3, 4, 6, 3]						,   lowerCamelCase="preactivation"						,   lowerCamelCase="relu"						,   lowerCamelCase=None						,   lowerCamelCase=32						,   lowerCamelCase=0.0						,   lowerCamelCase=False						,   lowerCamelCase=32						,   lowerCamelCase=1						,   lowerCamelCase=None						,   lowerCamelCase=None						,   **lowerCamelCase						,   ):
       super().__init__(**lowerCamelCase	)
       if layer_type not in self.layer_types:
         raise ValueError(F'''layer_type={layer_type} is not one of {",".join(self.layer_types	)}'''	)
       if global_padding is not None:
         if global_padding.upper() in self.supported_padding:
           _snake_case     		=				global_padding.upper()
         else:
           raise ValueError(F'''Padding strategy {global_padding} not supported'''	)
       _snake_case     		=				num_channels
       _snake_case     		=				embedding_size
       _snake_case     		=				hidden_sizes
       _snake_case     		=				depths
       _snake_case     		=				layer_type
       _snake_case     		=				hidden_act
       _snake_case     		=				global_padding
       _snake_case     		=				num_groups
       _snake_case     		=				drop_path_rate
       _snake_case     		=				embedding_dynamic_padding
       _snake_case     		=				output_stride
       _snake_case     		=				width_factor

       _snake_case     		=				["stem"] + [F'''stage{idx}''' for idx in range(1						,   len(lowerCamelCase	) + 1	)]
       _snake_case		,		_snake_case     		=				get_aligned_output_features_output_indices(
           out_features=lowerCamelCase						,   out_indices=lowerCamelCase						,   stage_names=self.stage_names	)


