






"""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      :				str		)		->    Dict:
				super().__init__()
				UpperCAmelCase_  :				Optional[int]      =      nn.Linear(3			,			4		)
				UpperCAmelCase_  :				List[Any]      =      nn.BatchNormad(4		)
				UpperCAmelCase_  :				Any      =      nn.Linear(4			,			5		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Dict			,			lowerCAmelCase_      :				List[Any]		)		->    Optional[Any]:
				return self.lineara(self.batchnorm(self.lineara(lowerCAmelCase_		)		)		)





class     UpperCamelCase_ (__A  ):





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]			,			lowerCAmelCase_      :				Union[str, Any]			,			*lowerCAmelCase_      :				str			,			**lowerCAmelCase_      :				int		)		->    Optional[int]:
				return (args[0] + 1,) + args[1:], kwargs





class     UpperCamelCase_ (__A  ):





	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]			,			lowerCAmelCase_      :				Any			,			lowerCAmelCase_      :				Optional[Any]		)		->    Dict:
				return output + 1





class     UpperCamelCase_ (unittest.TestCase  ):





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

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

				UpperCAmelCase_  :				Optional[Any]      =      test_model(lowerCAmelCase_		)
				assert torch.allclose(lowerCAmelCase_			,			lowerCAmelCase_			,			atol=1e-5		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]		)		->    str:
				UpperCAmelCase_  :				str      =      ModelForTest()
				UpperCAmelCase_  :				str      =      torch.randn(2			,			3		)
				UpperCAmelCase_  :				Optional[Any]      =      test_model(lowerCAmelCase_		)

				UpperCAmelCase_  :				Dict      =      PostForwardHook()
				add_hook_to_module(lowerCAmelCase_			,			lowerCAmelCase_		)
				UpperCAmelCase_  :				Union[str, Any]      =      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
				UpperCAmelCase_  :				Any      =      PostForwardHook()
				add_hook_to_module(lowerCAmelCase_			,			lowerCAmelCase_		)
				UpperCAmelCase_  :				Optional[Any]      =      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
				UpperCAmelCase_  :				Any      =      SequentialHook(PostForwardHook()			,			PostForwardHook()		)
				add_hook_to_module(lowerCAmelCase_			,			lowerCAmelCase_		)

				UpperCAmelCase_  :				List[Any]      =      test_model(lowerCAmelCase_		)
				assert torch.allclose(lowerCAmelCase_			,			output + 2			,			atol=1e-5		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]		)		->    Dict:
				UpperCAmelCase_  :				Optional[Any]      =      ModelForTest()
				UpperCAmelCase_  :				int      =      torch.randn(2			,			3		)
				UpperCAmelCase_  :				Tuple      =      test_model(lowerCAmelCase_		)

				UpperCAmelCase_  :				str      =      PostForwardHook()
				add_hook_to_module(lowerCAmelCase_			,			lowerCAmelCase_		)
				UpperCAmelCase_  :				Any      =      test_model(lowerCAmelCase_		)
				self.assertTrue(torch.allclose(lowerCAmelCase_			,			output + 1		)		)
				self.assertTrue(outputa.requires_grad		)

				UpperCAmelCase_  :				Dict      =      True
				UpperCAmelCase_  :				Optional[Any]      =      test_model(lowerCAmelCase_		)
				self.assertFalse(outputa.requires_grad		)





	@require_multi_gpu
	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]		)		->    str:
				UpperCAmelCase_  :				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
				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
				UpperCAmelCase_  :				List[str]      =      torch.randn(2			,			3		)
				UpperCAmelCase_  :				int      =      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_		)		)
				UpperCAmelCase_  :				Optional[int]      =      torch.randn(2			,			3		).to(0		)
				UpperCAmelCase_  :				Optional[int]      =      model(lowerCAmelCase_		)
				self.assertEqual(output.device			,			torch.device(0		)		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Dict		)		->    Dict:
				UpperCAmelCase_  :				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
				UpperCAmelCase_  :				List[str]      =      {"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
				UpperCAmelCase_  :				Union[str, Any]      =      torch.device(hook_kwargs["execution_device"]		)
				self.assertEqual(model.batchnorm.running_mean.device			,			lowerCAmelCase_		)

				UpperCAmelCase_  :				Dict      =      torch.randn(2			,			3		)
				UpperCAmelCase_  :				int      =      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
				UpperCAmelCase_  :				Optional[Any]      =      {
				    "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"		)		)

				UpperCAmelCase_  :				Any      =      torch.randn(2			,			3		)
				UpperCAmelCase_  :				Union[str, Any]      =      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   _SCREAMING_SNAKE_CASE   (   self      :				Tuple		)		->    Optional[Any]:
				UpperCAmelCase_  :				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
				UpperCAmelCase_  :				str      =      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
				UpperCAmelCase_  :				Tuple      =      torch.device(lowerCAmelCase_		)
				self.assertEqual(model.batchnorm.running_mean.device			,			lowerCAmelCase_		)

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

				UpperCAmelCase_  :				str      =      torch.randn(2			,			3		)
				UpperCAmelCase_  :				Union[str, Any]      =      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   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    Tuple:
				UpperCAmelCase_  :				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
				UpperCAmelCase_  :				int      =      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
				UpperCAmelCase_  :				int      =      torch.device(lowerCAmelCase_		)
				self.assertEqual(model.batchnorm.running_mean.device			,			lowerCAmelCase_		)

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

				UpperCAmelCase_  :				List[Any]      =      torch.randn(2			,			3		)
				UpperCAmelCase_  :				Tuple      =      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 collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCamelCase_       			=   logging.get_logger(__name__)

lowerCamelCase_       			=   {
    '''facebook/levit-128S''': '''https://huggingface.co/facebook/levit-128S/resolve/main/config.json''',
    # See all LeViT models at https://huggingface.co/models?filter=levit
}







class     UpperCamelCase_ (__A  ):
	__magic_name__			   =       '''levit'''





	def __init__(   self      :				List[str]			,			lowerCAmelCase_      :				int=224			,			lowerCAmelCase_      :				Union[str, Any]=3			,			lowerCAmelCase_      :				Tuple=3			,			lowerCAmelCase_      :				int=2			,			lowerCAmelCase_      :				str=1			,			lowerCAmelCase_      :				Optional[Any]=16			,			lowerCAmelCase_      :				Tuple=[128, 256, 384]			,			lowerCAmelCase_      :				Optional[int]=[4, 8, 12]			,			lowerCAmelCase_      :				str=[4, 4, 4]			,			lowerCAmelCase_      :				Dict=[16, 16, 16]			,			lowerCAmelCase_      :				int=0			,			lowerCAmelCase_      :				Optional[int]=[2, 2, 2]			,			lowerCAmelCase_      :				Any=[2, 2, 2]			,			lowerCAmelCase_      :				int=0.0_2			,			**lowerCAmelCase_      :				List[Any]			,			)		->    List[str]:
				super().__init__(**lowerCAmelCase_		)
				UpperCAmelCase_  :				int      =      image_size
				UpperCAmelCase_  :				List[Any]      =      num_channels
				UpperCAmelCase_  :				str      =      kernel_size
				UpperCAmelCase_  :				List[Any]      =      stride
				UpperCAmelCase_  :				List[str]      =      padding
				UpperCAmelCase_  :				Any      =      hidden_sizes
				UpperCAmelCase_  :				List[str]      =      num_attention_heads
				UpperCAmelCase_  :				List[str]      =      depths
				UpperCAmelCase_  :				int      =      key_dim
				UpperCAmelCase_  :				List[str]      =      drop_path_rate
				UpperCAmelCase_  :				str      =      patch_size
				UpperCAmelCase_  :				Tuple      =      attention_ratio
				UpperCAmelCase_  :				Optional[int]      =      mlp_ratio
				UpperCAmelCase_  :				Union[str, Any]      =      initializer_range
				UpperCAmelCase_  :				List[Any]      =      [
				    ["Subsample", key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
				    ["Subsample", key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
				]







class     UpperCamelCase_ (__A  ):
	__magic_name__			   =       version.parse('''1.11'''  )





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				int		)		->    Mapping[str, Mapping[int, str]]:
				return OrderedDict(
				    [
				        ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
				    ]		)





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				Dict		)		->    float:
				return 1e-4



