

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



 """simple docstring"""



 def __init__(		self					:       Optional[Any]			)  -> Optional[int]:


      """simple docstring"""




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


 def       UpperCAmelCase__		(		self					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]			)  -> Dict:


      """simple docstring"""




      return self.lineara(self.batchnorm(self.lineara(__SCREAMING_SNAKE_CASE			)			)			)







class 						lowerCAmelCase__   (      a					):



 """simple docstring"""



 def       UpperCAmelCase__		(		self					:       Optional[int]					,							__SCREAMING_SNAKE_CASE					:       Dict					,							*__SCREAMING_SNAKE_CASE					:       str					,							**__SCREAMING_SNAKE_CASE					:       Optional[int]			)  -> Optional[Any]:


      """simple docstring"""




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







class 						lowerCAmelCase__   (      a					):



 """simple docstring"""



 def       UpperCAmelCase__		(		self					:       Tuple					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]			)  -> Union[str, Any]:


      """simple docstring"""




      return output + 1







class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



 def       UpperCAmelCase__		(		self					:       str			)  -> Union[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      ModelForTest()
      __SCREAMING_SNAKE_CASE     =      ModelHook()

      add_hook_to_module(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
      self.assertEqual(test_model._hf_hook					,							__SCREAMING_SNAKE_CASE			)
      self.assertTrue(hasattr(__SCREAMING_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(__SCREAMING_SNAKE_CASE			)
      self.assertFalse(hasattr(__SCREAMING_SNAKE_CASE					,							"""_hf_hook"""			)			)
      self.assertFalse(hasattr(__SCREAMING_SNAKE_CASE					,							"""_old_forward"""			)			)
 def       UpperCAmelCase__		(		self					:       Tuple			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      ModelForTest()
      __SCREAMING_SNAKE_CASE     =      ModelHook()

      add_hook_to_module(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
      add_hook_to_module(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							append=__SCREAMING_SNAKE_CASE			)

      self.assertEqual(isinstance(test_model._hf_hook					,							__SCREAMING_SNAKE_CASE			)					,							__SCREAMING_SNAKE_CASE			)
      self.assertEqual(len(test_model._hf_hook.hooks			)					,							2			)
      self.assertTrue(hasattr(__SCREAMING_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(__SCREAMING_SNAKE_CASE			)
      self.assertFalse(hasattr(__SCREAMING_SNAKE_CASE					,							"""_hf_hook"""			)			)
      self.assertFalse(hasattr(__SCREAMING_SNAKE_CASE					,							"""_old_forward"""			)			)
 def       UpperCAmelCase__		(		self					:       int			)  -> List[Any]:


      """simple docstring"""




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

      __SCREAMING_SNAKE_CASE     =      PreForwardHook()
      add_hook_to_module(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)
      self.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							atol=1E-5			)			)

      # Attaching a hook to a model when it already has one replaces, does not chain
      __SCREAMING_SNAKE_CASE     =      PreForwardHook()
      add_hook_to_module(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)
      self.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							atol=1E-5			)			)

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

      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)
      assert torch.allclose(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							atol=1E-5			)
 def       UpperCAmelCase__		(		self					:       Optional[int]			)  -> Optional[int]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      ModelForTest()
      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      PostForwardHook()
      add_hook_to_module(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)
      self.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE					,							output + 1					,							atol=1E-5			)			)

      # Attaching a hook to a model when it already has one replaces, does not chain
      __SCREAMING_SNAKE_CASE     =      PostForwardHook()
      add_hook_to_module(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)
      self.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE					,							output + 1					,							atol=1E-5			)			)

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

      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)
      assert torch.allclose(__SCREAMING_SNAKE_CASE					,							output + 2					,							atol=1E-5			)
 def       UpperCAmelCase__		(		self					:       List[str]			)  -> Union[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      ModelForTest()
      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)

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

      __SCREAMING_SNAKE_CASE     =      True
      __SCREAMING_SNAKE_CASE     =      test_model(__SCREAMING_SNAKE_CASE			)
      self.assertFalse(outputa.requires_grad			)
 @require_multi_gpu
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> str:


      """simple docstring"""




      __SCREAMING_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
      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_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(__SCREAMING_SNAKE_CASE					,							AlignDevicesHook(io_same_device=__SCREAMING_SNAKE_CASE			)			)
      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			).to(0			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(output.device					,							torch.device(0			)			)
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> Tuple:


      """simple docstring"""




      __SCREAMING_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
      __SCREAMING_SNAKE_CASE     =      {"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

      add_hook_to_module(model.lineara					,							AlignDevicesHook(**__SCREAMING_SNAKE_CASE			)			)
      add_hook_to_module(model.batchnorm					,							AlignDevicesHook(**__SCREAMING_SNAKE_CASE			)			)
      add_hook_to_module(model.lineara					,							AlignDevicesHook(**__SCREAMING_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
      __SCREAMING_SNAKE_CASE     =      torch.device(hook_kwargs["""execution_device"""]			)
      self.assertEqual(model.batchnorm.running_mean.device					,							__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(output.device					,							__SCREAMING_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
      __SCREAMING_SNAKE_CASE     =      {
          """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
          """offload""": True,
          """offload_buffers""": True,
      }

      add_hook_to_module(model.lineara					,							AlignDevicesHook(**__SCREAMING_SNAKE_CASE			)			)
      add_hook_to_module(model.batchnorm					,							AlignDevicesHook(**__SCREAMING_SNAKE_CASE			)			)
      add_hook_to_module(model.lineara					,							AlignDevicesHook(**__SCREAMING_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"""			)			)

      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(output.device					,							__SCREAMING_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       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> Union[str, Any]:


      """simple docstring"""




      __SCREAMING_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
      __SCREAMING_SNAKE_CASE     =      0 if torch.cuda.is_available() else """cpu"""
      attach_align_device_hook(__SCREAMING_SNAKE_CASE					,							execution_device=__SCREAMING_SNAKE_CASE					,							offload=__SCREAMING_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
      __SCREAMING_SNAKE_CASE     =      torch.device(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(model.batchnorm.running_mean.device					,							__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(output.device					,							__SCREAMING_SNAKE_CASE			)

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

      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(output.device					,							__SCREAMING_SNAKE_CASE			)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(__SCREAMING_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       UpperCAmelCase__		(		self					:       Any			)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_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
      __SCREAMING_SNAKE_CASE     =      0 if torch.cuda.is_available() else """cpu"""
      attach_align_device_hook(
          __SCREAMING_SNAKE_CASE					,							execution_device=__SCREAMING_SNAKE_CASE					,							offload=__SCREAMING_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
      __SCREAMING_SNAKE_CASE     =      torch.device(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(model.batchnorm.running_mean.device					,							__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(output.device					,							__SCREAMING_SNAKE_CASE			)

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

      __SCREAMING_SNAKE_CASE     =      torch.randn(2					,							3			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE			)
      self.assertEqual(output.device					,							__SCREAMING_SNAKE_CASE			)

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



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


UpperCAmelCase :       List[str]			   =     logging.get_logger(__name__)







class 						lowerCAmelCase__   (      a ,					a					):



 """simple docstring"""



 lowerCAmelCase__  					=      "maskformer-swin"

 lowerCAmelCase__  					=      {
     "num_attention_heads": "num_heads",
     "num_hidden_layers": "num_layers",
 }
 def __init__(		self					:       str					,							__SCREAMING_SNAKE_CASE					:       Tuple=224					,							__SCREAMING_SNAKE_CASE					:       str=4					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]=3					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]=96					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]=[2, 2, 6, 2]					,							__SCREAMING_SNAKE_CASE					:       Any=[3, 6, 12, 24]					,							__SCREAMING_SNAKE_CASE					:       Dict=7					,							__SCREAMING_SNAKE_CASE					:       Dict=4.0					,							__SCREAMING_SNAKE_CASE					:       Dict=True					,							__SCREAMING_SNAKE_CASE					:       str=0.0					,							__SCREAMING_SNAKE_CASE					:       int=0.0					,							__SCREAMING_SNAKE_CASE					:       str=0.1					,							__SCREAMING_SNAKE_CASE					:       List[Any]="gelu"					,							__SCREAMING_SNAKE_CASE					:       str=False					,							__SCREAMING_SNAKE_CASE					:       Optional[int]=0.02					,							__SCREAMING_SNAKE_CASE					:       Optional[int]=1E-5					,							__SCREAMING_SNAKE_CASE					:       Dict=None					,							__SCREAMING_SNAKE_CASE					:       Dict=None					,							**__SCREAMING_SNAKE_CASE					:       Tuple					,							)  -> Tuple:


      """simple docstring"""




      super().__init__(**__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      image_size
      __SCREAMING_SNAKE_CASE     =      patch_size
      __SCREAMING_SNAKE_CASE     =      num_channels
      __SCREAMING_SNAKE_CASE     =      embed_dim
      __SCREAMING_SNAKE_CASE     =      depths
      __SCREAMING_SNAKE_CASE     =      len(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      num_heads
      __SCREAMING_SNAKE_CASE     =      window_size
      __SCREAMING_SNAKE_CASE     =      mlp_ratio
      __SCREAMING_SNAKE_CASE     =      qkv_bias
      __SCREAMING_SNAKE_CASE     =      hidden_dropout_prob
      __SCREAMING_SNAKE_CASE     =      attention_probs_dropout_prob
      __SCREAMING_SNAKE_CASE     =      drop_path_rate
      __SCREAMING_SNAKE_CASE     =      hidden_act
      __SCREAMING_SNAKE_CASE     =      use_absolute_embeddings
      __SCREAMING_SNAKE_CASE     =      layer_norm_eps
      __SCREAMING_SNAKE_CASE     =      initializer_range
      # we set the hidden_size attribute in order to make Swin work with VisionEncoderDecoderModel
      # this indicates the channel dimension after the last stage of the model
      __SCREAMING_SNAKE_CASE     =      int(embed_dim * 2 ** (len(__SCREAMING_SNAKE_CASE			) - 1)			)
      __SCREAMING_SNAKE_CASE     =      ["""stem"""] + [f'stage{idx}' for idx in range(1					,							len(__SCREAMING_SNAKE_CASE			) + 1			)]
      __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      get_aligned_output_features_output_indices(
          out_features=__SCREAMING_SNAKE_CASE					,							out_indices=__SCREAMING_SNAKE_CASE					,							stage_names=self.stage_names			)






