


"""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    lowerCamelCase__						(		nn.Module    ):
      def __init__(		self	:	List[Any]							):






       '''simple docstring'''


       super().__init__()
       __lowercase									=	nn.Linear(3      ,       4							)
       __lowercase									=	nn.BatchNormad(4							)
       __lowercase									=	nn.Linear(4      ,       5							)






      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple      ,       A_	:	Union[str, Any]							):






       '''simple docstring'''


       return self.lineara(self.batchnorm(self.lineara(_UpperCAmelCase							)							)							)






class    lowerCamelCase__						(		_UpperCamelCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	int      ,       A_	:	Optional[Any]      ,       *A_	:	Optional[int]      ,       **A_	:	str							):






       '''simple docstring'''


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






class    lowerCamelCase__						(		_UpperCamelCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple      ,       A_	:	str      ,       A_	:	str							):






       '''simple docstring'''


       return output + 1






class    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


       __lowercase									=	ModelForTest()
       __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       __lowercase									=	ModelForTest()
       __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	ModelForTest()
       __lowercase									=	torch.randn(2      ,       3							)
       __lowercase									=	test_model(x + 1							)
       __lowercase									=	test_model(x + 2							)

       __lowercase									=	PreForwardHook()
       add_hook_to_module(_UpperCAmelCase      ,       _UpperCAmelCase							)
       __lowercase									=	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
       __lowercase									=	PreForwardHook()
       add_hook_to_module(_UpperCAmelCase      ,       _UpperCAmelCase							)
       __lowercase									=	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
       __lowercase									=	SequentialHook(PreForwardHook()      ,       PreForwardHook()							)
       add_hook_to_module(_UpperCAmelCase      ,       _UpperCAmelCase							)

       __lowercase									=	test_model(_UpperCAmelCase							)
       assert torch.allclose(_UpperCAmelCase      ,       _UpperCAmelCase      ,       atol=1e-5							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       __lowercase									=	ModelForTest()
       __lowercase									=	torch.randn(2      ,       3							)
       __lowercase									=	test_model(_UpperCAmelCase							)

       __lowercase									=	PostForwardHook()
       add_hook_to_module(_UpperCAmelCase      ,       _UpperCAmelCase							)
       __lowercase									=	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
       __lowercase									=	PostForwardHook()
       add_hook_to_module(_UpperCAmelCase      ,       _UpperCAmelCase							)
       __lowercase									=	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
       __lowercase									=	SequentialHook(PostForwardHook()      ,       PostForwardHook()							)
       add_hook_to_module(_UpperCAmelCase      ,       _UpperCAmelCase							)

       __lowercase									=	test_model(_UpperCAmelCase							)
       assert torch.allclose(_UpperCAmelCase      ,       output + 2      ,       atol=1e-5							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       __lowercase									=	ModelForTest()
       __lowercase									=	torch.randn(2      ,       3							)
       __lowercase									=	test_model(_UpperCAmelCase							)

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

       __lowercase									=	True
       __lowercase									=	test_model(_UpperCAmelCase							)
       self.assertFalse(outputa.requires_grad							)
      @require_multi_gpu
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	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
       __lowercase									=	torch.randn(2      ,       3							)
       __lowercase									=	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							)							)
       __lowercase									=	torch.randn(2      ,       3							).to(0							)
       __lowercase									=	model(_UpperCAmelCase							)
       self.assertEqual(output.device      ,       torch.device(0							)							)
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


       __lowercase									=	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
       __lowercase									=	{'''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
       __lowercase									=	torch.device(hook_kwargs["""execution_device"""]							)
       self.assertEqual(model.batchnorm.running_mean.device      ,       _UpperCAmelCase							)

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

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






       '''simple docstring'''


       __lowercase									=	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
       __lowercase									=	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
       __lowercase									=	torch.device(_UpperCAmelCase							)
       self.assertEqual(model.batchnorm.running_mean.device      ,       _UpperCAmelCase							)

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

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






       '''simple docstring'''


       __lowercase									=	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
       __lowercase									=	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
       __lowercase									=	torch.device(_UpperCAmelCase							)
       self.assertEqual(model.batchnorm.running_mean.device      ,       _UpperCAmelCase							)

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

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




"""simple docstring"""

import warnings
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


UpperCAmelCase__  							=logging.get_logger(__name__)

UpperCAmelCase__  							={
    "nvidia/segformer-b0-finetuned-ade-512-512": (
        "https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512/resolve/main/config.json"
    ),
    # See all SegFormer models at https://huggingface.co/models?filter=segformer
}





class    lowerCamelCase__						(		_a    ):
      a :     Optional[int]											=		"""segformer"""
      def __init__(		self	:	str      ,       A_	:	Optional[int]=3      ,       A_	:	Any=4      ,       A_	:	Optional[int]=[2, 2, 2, 2]      ,       A_	:	str=[8, 4, 2, 1]      ,       A_	:	Any=[3_2, 6_4, 1_6_0, 2_5_6]      ,       A_	:	Union[str, Any]=[7, 3, 3, 3]      ,       A_	:	Any=[4, 2, 2, 2]      ,       A_	:	int=[1, 2, 5, 8]      ,       A_	:	Union[str, Any]=[4, 4, 4, 4]      ,       A_	:	Optional[Any]="gelu"      ,       A_	:	List[str]=0.0      ,       A_	:	Any=0.0      ,       A_	:	Optional[Any]=0.1      ,       A_	:	List[str]=0.02      ,       A_	:	Tuple=0.1      ,       A_	:	Optional[Any]=1e-6      ,       A_	:	Optional[int]=2_5_6      ,       A_	:	Optional[Any]=2_5_5      ,       **A_	:	Union[str, Any]      ,       ):






       '''simple docstring'''


       super().__init__(**A_							)

       if "reshape_last_stage" in kwargs and kwargs["reshape_last_stage"] is False:
        warnings.warn(
            """Reshape_last_stage is set to False in this config. This argument is deprecated and will soon be"""
            """ removed, as the behaviour will default to that of reshape_last_stage = True."""      ,       A_      ,       )

       __lowercase									=	num_channels
       __lowercase									=	num_encoder_blocks
       __lowercase									=	depths
       __lowercase									=	sr_ratios
       __lowercase									=	hidden_sizes
       __lowercase									=	patch_sizes
       __lowercase									=	strides
       __lowercase									=	mlp_ratios
       __lowercase									=	num_attention_heads
       __lowercase									=	hidden_act
       __lowercase									=	hidden_dropout_prob
       __lowercase									=	attention_probs_dropout_prob
       __lowercase									=	classifier_dropout_prob
       __lowercase									=	initializer_range
       __lowercase									=	drop_path_rate
       __lowercase									=	layer_norm_eps
       __lowercase									=	decoder_hidden_size
       __lowercase									=	kwargs.get("""reshape_last_stage"""      ,       A_							)
       __lowercase									=	semantic_loss_ignore_index






class    lowerCamelCase__						(		_a    ):
      a :     List[str]											=		version.parse("""1.11"""    )
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''simple docstring'''


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






       '''simple docstring'''


       return 1e-4






      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       return 1_2
