




'''simple docstring'''


import inspect
import os

import torch
from transformers import AutoModel
from transformers.testing_utils import mockenv_context
from transformers.trainer_utils import set_seed

import accelerate
from accelerate.accelerator import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils.testing import (
    AccelerateTestCase,
    TempDirTestCase,
    execute_subprocess_async,
    require_cuda,
    require_fsdp,
    require_multi_gpu,
    slow,
)
from accelerate.utils.constants import (
    FSDP_AUTO_WRAP_POLICY,
    FSDP_BACKWARD_PREFETCH,
    FSDP_SHARDING_STRATEGY,
    FSDP_STATE_DICT_TYPE,
)
from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
from accelerate.utils.other import patch_environment


set_seed(4_2)

_lowerCAmelCase  		=      "bert-base-cased"
_lowerCAmelCase  		=      "fp16"
_lowerCAmelCase  		=      "bf16"
_lowerCAmelCase  		=      [FPaa, BFaa]
@require_fsdp
@require_cuda
class   __A (							SCREAMING_SNAKE_CASE__  ):



   """simple docstring"""







   def snake_case_(      self						)->       Optional[int]:
      super().setUp()

      lowercase__       							=							dict(
          ACCELERATE_USE_FSDP='''true'''				, MASTER_ADDR='''localhost'''				, MASTER_PORT='''10999'''				, RANK='''0'''				, LOCAL_RANK='''0'''				, WORLD_SIZE='''1'''				, )

   def snake_case_(      self						)->       str:
      from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

      for i, strategy in enumerate(_lowercase						):
         lowercase__       							=							self.dist_env.copy()
         lowercase__       							=							f'''{i + 1}'''
         lowercase__       							=							strategy
         with mockenv_context(**_lowercase						):
            lowercase__       							=							FullyShardedDataParallelPlugin()
            self.assertEqual(fsdp_plugin.sharding_strategy				, ShardingStrategy(i + 1						)						)

   def snake_case_(      self						)->       Any:
      from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

      for i, prefetch_policy in enumerate(_lowercase						):
         lowercase__       							=							self.dist_env.copy()
         lowercase__       							=							prefetch_policy
         with mockenv_context(**_lowercase						):
            lowercase__       							=							FullyShardedDataParallelPlugin()
            if prefetch_policy == "NO_PREFETCH":
               self.assertIsNone(fsdp_plugin.backward_prefetch						)
            else:
               self.assertEqual(fsdp_plugin.backward_prefetch				, BackwardPrefetch(i + 1						)						)

   def snake_case_(      self						)->       int:
      from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

      for i, state_dict_type in enumerate(_lowercase						):
         lowercase__       							=							self.dist_env.copy()
         lowercase__       							=							state_dict_type
         with mockenv_context(**_lowercase						):
            lowercase__       							=							FullyShardedDataParallelPlugin()
            self.assertEqual(fsdp_plugin.state_dict_type				, StateDictType(i + 1						)						)
            if state_dict_type == "FULL_STATE_DICT":
               self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu						)
               self.assertTrue(fsdp_plugin.state_dict_config.ranka_only						)

   def snake_case_(      self						)->       str:
      lowercase__       							=							AutoModel.from_pretrained(_lowercase						)
      for policy in FSDP_AUTO_WRAP_POLICY:
         lowercase__       							=							self.dist_env.copy()
         lowercase__       							=							policy
         if policy == "TRANSFORMER_BASED_WRAP":
            lowercase__       							=							"BertLayer"
         elif policy == "SIZE_BASED_WRAP":
            lowercase__       							=							"2000"
         with mockenv_context(**_lowercase						):
            lowercase__       							=							FullyShardedDataParallelPlugin()
            fsdp_plugin.set_auto_wrap_policy(_lowercase						)
            if policy == "NO_WRAP":
               self.assertIsNone(fsdp_plugin.auto_wrap_policy						)
            else:
               self.assertIsNotNone(fsdp_plugin.auto_wrap_policy						)

      lowercase__       							=							self.dist_env.copy()
      lowercase__       							=							"TRANSFORMER_BASED_WRAP"
      lowercase__       							=							"T5Layer"
      with mockenv_context(**_lowercase						):
         lowercase__       							=							FullyShardedDataParallelPlugin()
         with self.assertRaises(_lowercase						) as cm:
            fsdp_plugin.set_auto_wrap_policy(_lowercase						)
         self.assertTrue('''Could not find the transformer layer class to wrap in the model.''' in str(cm.exception						)						)

      lowercase__       							=							self.dist_env.copy()
      lowercase__       							=							"SIZE_BASED_WRAP"
      lowercase__       							=							"0"
      with mockenv_context(**_lowercase						):
         lowercase__       							=							FullyShardedDataParallelPlugin()
         fsdp_plugin.set_auto_wrap_policy(_lowercase						)
         self.assertIsNone(fsdp_plugin.auto_wrap_policy						)

   def snake_case_(      self						)->       str:
      from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
      from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

      for mp_dtype in dtypes:
         lowercase__       							=							self.dist_env.copy()
         lowercase__       							=							mp_dtype
         with mockenv_context(**_lowercase						):
            lowercase__       							=							Accelerator()
            if mp_dtype == "fp16":
               lowercase__       							=							torch.floataa
            elif mp_dtype == "bf16":
               lowercase__       							=							torch.bfloataa
            lowercase__       							=							MixedPrecision(param_dtype=_lowercase				, reduce_dtype=_lowercase				, buffer_dtype=_lowercase						)
            self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy				, _lowercase						)
            if mp_dtype == FPaa:
               self.assertTrue(isinstance(accelerator.scaler				, _lowercase						)						)
            elif mp_dtype == BFaa:
               self.assertIsNone(accelerator.scaler						)
            AcceleratorState._reset_state(_lowercase						)

   def snake_case_(      self						)->       Union[str, Any]:
      from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

      for flag in [True, False]:
         lowercase__       							=							self.dist_env.copy()
         lowercase__       							=							str(_lowercase						).lower()
         with mockenv_context(**_lowercase						):
            lowercase__       							=							FullyShardedDataParallelPlugin()
            self.assertEqual(fsdp_plugin.cpu_offload				, CPUOffload(offload_params=_lowercase						)						)



@require_fsdp
@require_multi_gpu
@slow
class   __A (							SCREAMING_SNAKE_CASE__  ):



   """simple docstring"""







   def snake_case_(      self						)->       Dict:
      super().setUp()
      lowercase__       							=							0.8_2
      lowercase__       							=							[
          "fsdp_shard_grad_op_transformer_based_wrap",
          "fsdp_full_shard_transformer_based_wrap",
      ]
      lowercase__       							=							{
          "multi_gpu_fp16": 3_2_0_0,
          "fsdp_shard_grad_op_transformer_based_wrap_fp16": 2_0_0_0,
          "fsdp_full_shard_transformer_based_wrap_fp16": 1_9_0_0,
          # Disabling below test as it overwhelms the RAM memory usage
          # on CI self-hosted runner leading to tests getting killed.
          # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
      }
      lowercase__       							=							1_6_0
      lowercase__       							=							1_6_0

      lowercase__       							=							inspect.getfile(accelerate.test_utils						)
      lowercase__       							=							os.path.sep.join(mod_file.split(os.path.sep						)[:-1] + ['''scripts''', '''external_deps''']						)

   def snake_case_(      self						)->       Dict:
      lowercase__       							=							os.path.join(self.test_scripts_folder				, '''test_performance.py'''						)
      lowercase__       							=							["accelerate", "launch", "--num_processes=2", "--num_machines=1", "--machine_rank=0", "--use_fsdp"]
      for config in self.performance_configs:
         lowercase__       							=							cmd.copy()
         for i, strategy in enumerate(_lowercase						):
            if strategy.lower() in config:
               cmd_config.append(f'''--fsdp_sharding_strategy={i+1}'''						)
               break

         if "fp32" in config:
            cmd_config.append('''--mixed_precision=no'''						)
         else:
            cmd_config.append('''--mixed_precision=fp16'''						)

         if "cpu_offload" in config:
            cmd_config.append('''--fsdp_offload_params=True'''						)

         for policy in FSDP_AUTO_WRAP_POLICY:
            if policy.lower() in config:
               cmd_config.append(f'''--fsdp_auto_wrap_policy={policy}'''						)
               break

         if policy == "TRANSFORMER_BASED_WRAP":
            cmd_config.append('''--fsdp_transformer_layer_cls_to_wrap=BertLayer'''						)
         elif policy == "SIZE_BASED_WRAP":
            cmd_config.append('''--fsdp_min_num_params=2000'''						)

         cmd_config.extend(
             [
                 self.test_file_path,
                 f'''--output_dir={self.tmpdir}''',
                 f'''--performance_lower_bound={self.performance_lower_bound}''',
             ]						)
         with patch_environment(omp_num_threads=1						):
            execute_subprocess_async(_lowercase				, env=os.environ.copy()						)

   def snake_case_(      self						)->       Any:
      lowercase__       							=							os.path.join(self.test_scripts_folder				, '''test_checkpointing.py'''						)
      lowercase__       							=							[
          "accelerate",
          "launch",
          "--num_processes=2",
          "--num_machines=1",
          "--machine_rank=0",
          "--use_fsdp",
          "--mixed_precision=fp16",
          "--fsdp_transformer_layer_cls_to_wrap=BertLayer",
      ]

      for i, strategy in enumerate(_lowercase						):
         lowercase__       							=							cmd.copy()
         cmd_config.append(f'''--fsdp_sharding_strategy={i+1}'''						)
         if strategy != "FULL_SHARD":
            continue
         lowercase__       							=							len(_lowercase						)
         for state_dict_type in FSDP_STATE_DICT_TYPE:
            lowercase__       							=							cmd_config[:state_dict_config_index]
            cmd_config.append(f'''--fsdp_state_dict_type={state_dict_type}'''						)
            cmd_config.extend(
                [
                    self.test_file_path,
                    f'''--output_dir={self.tmpdir}''',
                    '''--partial_train_epoch=1''',
                ]						)
            with patch_environment(omp_num_threads=1						):
               execute_subprocess_async(_lowercase				, env=os.environ.copy()						)

            lowercase__       							=							cmd_config[:-1]
            lowercase__       							=							os.path.join(self.tmpdir				, '''epoch_0'''						)
            cmd_config.extend(
                [
                    f'''--resume_from_checkpoint={resume_from_checkpoint}''',
                ]						)
            with patch_environment(omp_num_threads=1						):
               execute_subprocess_async(_lowercase				, env=os.environ.copy()						)

   def snake_case_(      self						)->       Union[str, Any]:
      lowercase__       							=							os.path.join(self.test_scripts_folder				, '''test_peak_memory_usage.py'''						)
      lowercase__       							=							[
          "accelerate",
          "launch",
          "--num_processes=2",
          "--num_machines=1",
          "--machine_rank=0",
      ]
      for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
         lowercase__       							=							cmd.copy()
         if "fp16" in spec:
            cmd_config.extend(['''--mixed_precision=fp16''']						)
         else:
            cmd_config.extend(['''--mixed_precision=no''']						)

         if "multi_gpu" in spec:
            continue
         else:
            cmd_config.extend(['''--use_fsdp''']						)
            for i, strategy in enumerate(_lowercase						):
               if strategy.lower() in spec:
                  cmd_config.append(f'''--fsdp_sharding_strategy={i+1}'''						)
                  break

            if "cpu_offload" in spec:
               cmd_config.append('''--fsdp_offload_params=True'''						)

            for policy in FSDP_AUTO_WRAP_POLICY:
               if policy.lower() in spec:
                  cmd_config.append(f'''--fsdp_auto_wrap_policy={policy}'''						)
                  break

            if policy == "TRANSFORMER_BASED_WRAP":
               cmd_config.append('''--fsdp_transformer_layer_cls_to_wrap=BertLayer'''						)
            elif policy == "SIZE_BASED_WRAP":
               cmd_config.append('''--fsdp_min_num_params=2000'''						)

         cmd_config.extend(
             [
                 self.test_file_path,
                 f'''--output_dir={self.tmpdir}''',
                 f'''--peak_memory_upper_bound={peak_mem_upper_bound}''',
                 f'''--n_train={self.n_train}''',
                 f'''--n_val={self.n_val}''',
             ]						)
         with patch_environment(omp_num_threads=1						):
            execute_subprocess_async(_lowercase				, env=os.environ.copy()						)









"""simple docstring"""







import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch
			from torch import nn

			from transformers import (
			    MODEL_MAPPING,
			    BeitForImageClassification,
			    BeitForMaskedImageModeling,
			    BeitForSemanticSegmentation,
			    BeitModel,
			)
			from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
			import PIL
			from PIL import Image

			from transformers import BeitImageProcessor

class       _a  :



			def __init__(			self							:				str   ,	_lowercase							:				Tuple   ,	_lowercase							:				List[Any]=100   ,	_lowercase							:				Union[str, Any]=13   ,	_lowercase							:				Tuple=30   ,	_lowercase							:				Dict=2   ,	_lowercase							:				Union[str, Any]=3   ,	_lowercase							:				Union[str, Any]=True   ,	_lowercase							:				Dict=True   ,	_lowercase							:				Dict=32   ,	_lowercase							:				Tuple=4   ,	_lowercase							:				List[Any]=4   ,	_lowercase							:				List[Any]=37   ,	_lowercase							:				List[Any]="gelu"   ,	_lowercase							:				int=0.1   ,	_lowercase							:				Optional[Any]=0.1   ,	_lowercase							:				Dict=10   ,	_lowercase							:				List[str]=0.02   ,	_lowercase							:				Optional[int]=3   ,	_lowercase							:				Tuple=None   ,	_lowercase							:				Optional[Any]=[0, 1, 2, 3]   ,	)							->					Any:
									snake_case   :			str						    = parent
									snake_case   :			List[Any]						    = 100
									snake_case   :			Any						    = batch_size
									snake_case   :			str						    = image_size
									snake_case   :			Optional[Any]						    = patch_size
									snake_case   :			Union[str, Any]						    = num_channels
									snake_case   :			str						    = is_training
									snake_case   :			List[str]						    = use_labels
									snake_case   :			Optional[Any]						    = hidden_size
									snake_case   :			Any						    = num_hidden_layers
									snake_case   :			Union[str, Any]						    = num_attention_heads
									snake_case   :			Any						    = intermediate_size
									snake_case   :			List[Any]						    = hidden_act
									snake_case   :			Dict						    = hidden_dropout_prob
									snake_case   :			Union[str, Any]						    = attention_probs_dropout_prob
									snake_case   :			Optional[Any]						    = type_sequence_label_size
									snake_case   :			Any						    = initializer_range
									snake_case   :			Tuple						    = scope
									snake_case   :			Optional[int]						    = out_indices
									snake_case   :			List[Any]						    = num_labels

									# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
									snake_case   :			Optional[Any]						    = (image_size // patch_size) ** 2
									snake_case   :			Any						    = num_patches + 1



			def 				__lowercase					(			self							:				Tuple		)							->					Union[str, Any]:
									snake_case   :			Any						    = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

									snake_case   :			Optional[int]						    = None
									snake_case   :			Optional[int]						    = None
									if self.use_labels:
															snake_case   :			List[str]						    = ids_tensor([self.batch_size]   ,	self.type_sequence_label_size		)
															snake_case   :			Tuple						    = ids_tensor([self.batch_size, self.image_size, self.image_size]   ,	self.num_labels		)

									snake_case   :			str						    = self.get_config()

									return config, pixel_values, labels, pixel_labels



			def 				__lowercase					(			self							:				str		)							->					List[Any]:
									return BeitConfig(
									    vocab_size=self.vocab_size   ,	image_size=self.image_size   ,	patch_size=self.patch_size   ,	num_channels=self.num_channels   ,	hidden_size=self.hidden_size   ,	num_hidden_layers=self.num_hidden_layers   ,	num_attention_heads=self.num_attention_heads   ,	intermediate_size=self.intermediate_size   ,	hidden_act=self.hidden_act   ,	hidden_dropout_prob=self.hidden_dropout_prob   ,	attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,	is_decoder=_lowercase   ,	initializer_range=self.initializer_range   ,	out_indices=self.out_indices   ,	)



			def 				__lowercase					(			self							:				Union[str, Any]   ,	_lowercase							:				int   ,	_lowercase							:				Any   ,	_lowercase							:				List[Any]   ,	_lowercase							:				List[Any]		)							->					Any:
									snake_case   :			int						    = BeitModel(config=_lowercase		)
									model.to(_lowercase		)
									model.eval()
									snake_case   :			List[str]						    = model(_lowercase		)
									self.parent.assertEqual(result.last_hidden_state.shape   ,	(self.batch_size, self.seq_length, self.hidden_size)		)



			def 				__lowercase					(			self							:				Any   ,	_lowercase							:				Any   ,	_lowercase							:				Optional[int]   ,	_lowercase							:				Tuple   ,	_lowercase							:				List[Any]		)							->					Tuple:
									snake_case   :			str						    = BeitForMaskedImageModeling(config=_lowercase		)
									model.to(_lowercase		)
									model.eval()
									snake_case   :			Tuple						    = model(_lowercase		)
									self.parent.assertEqual(result.logits.shape   ,	(self.batch_size, self.seq_length - 1, self.vocab_size)		)



			def 				__lowercase					(			self							:				Optional[Any]   ,	_lowercase							:				List[Any]   ,	_lowercase							:				Tuple   ,	_lowercase							:				Any   ,	_lowercase							:				List[Any]		)							->					str:
									snake_case   :			str						    = self.type_sequence_label_size
									snake_case   :			Union[str, Any]						    = BeitForImageClassification(_lowercase		)
									model.to(_lowercase		)
									model.eval()
									snake_case   :			Union[str, Any]						    = model(_lowercase   ,	labels=_lowercase		)
									self.parent.assertEqual(result.logits.shape   ,	(self.batch_size, self.type_sequence_label_size)		)

									# test greyscale images
									snake_case   :			Dict						    = 1
									snake_case   :			Tuple						    = BeitForImageClassification(_lowercase		)
									model.to(_lowercase		)
									model.eval()

									snake_case   :			Optional[Any]						    = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
									snake_case   :			List[str]						    = model(_lowercase   ,	labels=_lowercase		)
									self.parent.assertEqual(result.logits.shape   ,	(self.batch_size, self.type_sequence_label_size)		)



			def 				__lowercase					(			self							:				str   ,	_lowercase							:				Dict   ,	_lowercase							:				Dict   ,	_lowercase							:				Optional[Any]   ,	_lowercase							:				List[Any]		)							->					List[Any]:
									snake_case   :			List[Any]						    = self.num_labels
									snake_case   :			Union[str, Any]						    = BeitForSemanticSegmentation(_lowercase		)
									model.to(_lowercase		)
									model.eval()
									snake_case   :			Tuple						    = model(_lowercase		)
									self.parent.assertEqual(
									    result.logits.shape   ,	(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)		)
									snake_case   :			Any						    = model(_lowercase   ,	labels=_lowercase		)
									self.parent.assertEqual(
									    result.logits.shape   ,	(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)		)



			def 				__lowercase					(			self							:				Any		)							->					Dict:
									snake_case   :			List[Any]						    = self.prepare_config_and_inputs()
									snake_case    ,			snake_case    ,			snake_case    ,			snake_case   :			Optional[int]						    = config_and_inputs
									snake_case   :			Dict						    = {"pixel_values": pixel_values}
									return config, inputs_dict



@require_torch
class       _a  ( SCREAMING_SNAKE_CASE__     ,	SCREAMING_SNAKE_CASE__     ,	unittest.TestCase):
			__magic_name__     					=     (
			    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
			    if is_torch_available()
			    else ()
			)
			__magic_name__     					=     (
			    {
			        """feature-extraction""": BeitModel,
			        """image-classification""": BeitForImageClassification,
			        """image-segmentation""": BeitForSemanticSegmentation,
			    }
			    if is_torch_available()
			    else {}
			)

			__magic_name__     					=     False
			__magic_name__     					=     False
			__magic_name__     					=     False



			def 				__lowercase					(			self							:				Optional[int]		)							->					Tuple:
									snake_case   :			Optional[Any]						    = BeitModelTester(self		)
									snake_case   :			Union[str, Any]						    = ConfigTester(self   ,	config_class=_lowercase   ,	has_text_modality=_lowercase   ,	hidden_size=37		)



			def 				__lowercase					(			self							:				int		)							->					str:
									self.config_tester.run_common_tests()



			@unittest.skip(reason="BEiT does not use inputs_embeds"		)
			def 				__lowercase					(			self							:				Optional[Any]		)							->					Tuple:
									pass



			@require_torch_multi_gpu
			@unittest.skip(reason="BEiT has some layers using `add_module` which doesn't work well with `nn.DataParallel`"		)
			def 				__lowercase					(			self							:				List[str]		)							->					str:
									pass



			def 				__lowercase					(			self							:				List[str]		)							->					Union[str, Any]:
									snake_case    ,			snake_case   :			Any						    = self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															snake_case   :			List[str]						    = model_class(_lowercase		)
															self.assertIsInstance(model.get_input_embeddings()   ,	(nn.Module)		)
															snake_case   :			Dict						    = model.get_output_embeddings()
															self.assertTrue(x is None or isinstance(_lowercase   ,	nn.Linear		)		)



			def 				__lowercase					(			self							:				List[Any]		)							->					List[Any]:
									snake_case    ,			snake_case   :			List[Any]						    = self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															snake_case   :			str						    = model_class(_lowercase		)
															snake_case   :			Union[str, Any]						    = inspect.signature(model.forward		)
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															snake_case   :			List[Any]						    = [*signature.parameters.keys()]

															snake_case   :			int						    = ["pixel_values"]
															self.assertListEqual(arg_names[:1]   ,	_lowercase		)



			def 				__lowercase					(			self							:				Any		)							->					List[str]:
									snake_case   :			List[str]						    = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_lowercase		)



			def 				__lowercase					(			self							:				Optional[int]		)							->					str:
									snake_case   :			Optional[Any]						    = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*_lowercase		)



			def 				__lowercase					(			self							:				List[Any]		)							->					Tuple:
									snake_case   :			Optional[Any]						    = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*_lowercase		)



			def 				__lowercase					(			self							:				Union[str, Any]		)							->					Optional[int]:
									snake_case   :			List[str]						    = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_semantic_segmentation(*_lowercase		)



			def 				__lowercase					(			self							:				Optional[int]		)							->					Optional[Any]:
									if not self.model_tester.is_training:
															return

									snake_case    ,			snake_case   :			Dict						    = self.model_tester.prepare_config_and_inputs_for_common()
									snake_case   :			str						    = True

									for model_class in self.all_model_classes:
															# we don't test BeitForMaskedImageModeling
															if model_class in [*get_values(_lowercase		), BeitForMaskedImageModeling]:
																					continue

															snake_case   :			List[str]						    = model_class(_lowercase		)
															model.to(_lowercase		)
															model.train()
															snake_case   :			Dict						    = self._prepare_for_class(_lowercase   ,	_lowercase   ,	return_labels=_lowercase		)
															snake_case   :			Tuple						    = model(**_lowercase		).loss
															loss.backward()



			def 				__lowercase					(			self							:				Dict		)							->					str:
									snake_case    ,			snake_case   :			List[str]						    = self.model_tester.prepare_config_and_inputs_for_common()
									if not self.model_tester.is_training:
															return

									snake_case   :			Dict						    = False
									snake_case   :			Dict						    = True

									for model_class in self.all_model_classes:
															# we don't test BeitForMaskedImageModeling
															if (
															    model_class in [*get_values(_lowercase		), BeitForMaskedImageModeling]
															    or not model_class.supports_gradient_checkpointing
															):
																					continue

															snake_case   :			str						    = model_class(_lowercase		)
															model.gradient_checkpointing_enable()
															model.to(_lowercase		)
															model.train()
															snake_case   :			Dict						    = self._prepare_for_class(_lowercase   ,	_lowercase   ,	return_labels=_lowercase		)
															snake_case   :			str						    = model(**_lowercase		).loss
															loss.backward()



			def 				__lowercase					(			self							:				List[Any]		)							->					Optional[int]:
									snake_case    ,			snake_case   :			Union[str, Any]						    = self.model_tester.prepare_config_and_inputs_for_common()

									snake_case   :			List[Any]						    = _config_zero_init(_lowercase		)
									for model_class in self.all_model_classes:
															snake_case   :			Union[str, Any]						    = model_class(config=_lowercase		)
															for name, param in model.named_parameters():
																					# we skip lambda parameters as these require special initial values
																					# determined by config.layer_scale_init_value
																					if "lambda" in name:
																											continue
																					if param.requires_grad:
																											self.assertIn(
																											    ((param.data.mean() * 1E9).round() / 1E9).item()   ,	[0.0, 1.0]   ,	msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''   ,	)



			@slow
			def 				__lowercase					(			self							:				Optional[Any]		)							->					List[Any]:
									for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															snake_case   :			Tuple						    = BeitModel.from_pretrained(_lowercase		)
															self.assertIsNotNone(_lowercase		)



def  __SCREAMING_SNAKE_CASE		(			):
						"""simple docstring"""






						snake_case   :			Tuple						    = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png" )
						return image

@require_torch
@require_vision
class       _a  ( unittest.TestCase):



			@cached_property
			def 				__lowercase					(			self							:				Any		)							->					List[str]:
									return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"		) if is_vision_available() else None



			@slow
			def 				__lowercase					(			self							:				Dict		)							->					Tuple:
									snake_case   :			int						    = BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"		).to(_lowercase		)

									snake_case   :			str						    = self.default_image_processor
									snake_case   :			int						    = prepare_img()
									snake_case   :			int						    = image_processor(images=_lowercase   ,	return_tensors="pt"		).pixel_values.to(_lowercase		)

									# prepare bool_masked_pos
									snake_case   :			str						    = torch.ones((1, 196)   ,	dtype=torch.bool		).to(_lowercase		)

									# forward pass
									with torch.no_grad():
															snake_case   :			Optional[Any]						    = model(pixel_values=_lowercase   ,	bool_masked_pos=_lowercase		)
									snake_case   :			Any						    = outputs.logits

									# verify the logits
									snake_case   :			int						    = torch.Size((1, 196, 8192)		)
									self.assertEqual(logits.shape   ,	_lowercase		)

									snake_case   :			List[Any]						    = torch.tensor(
									    [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]		).to(_lowercase		)

									self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]   ,	_lowercase   ,	atol=1E-2		)		)



			@slow
			def 				__lowercase					(			self							:				List[str]		)							->					Tuple:
									snake_case   :			List[Any]						    = BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"		).to(_lowercase		)

									snake_case   :			Tuple						    = self.default_image_processor
									snake_case   :			Dict						    = prepare_img()
									snake_case   :			int						    = image_processor(images=_lowercase   ,	return_tensors="pt"		).to(_lowercase		)

									# forward pass
									with torch.no_grad():
															snake_case   :			Tuple						    = model(**_lowercase		)
									snake_case   :			Union[str, Any]						    = outputs.logits

									# verify the logits
									snake_case   :			Dict						    = torch.Size((1, 1000)		)
									self.assertEqual(logits.shape   ,	_lowercase		)

									snake_case   :			Optional[Any]						    = torch.tensor([-1.2385, -1.0987, -1.0108]		).to(_lowercase		)

									self.assertTrue(torch.allclose(logits[0, :3]   ,	_lowercase   ,	atol=1E-4		)		)

									snake_case   :			Optional[int]						    = 281
									self.assertEqual(logits.argmax(-1		).item()   ,	_lowercase		)



			@slow
			def 				__lowercase					(			self							:				Union[str, Any]		)							->					List[str]:
									snake_case   :			str						    = BeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"		).to(
									    _lowercase		)

									snake_case   :			int						    = self.default_image_processor
									snake_case   :			List[str]						    = prepare_img()
									snake_case   :			Tuple						    = image_processor(images=_lowercase   ,	return_tensors="pt"		).to(_lowercase		)

									# forward pass
									with torch.no_grad():
															snake_case   :			Optional[Any]						    = model(**_lowercase		)
									snake_case   :			str						    = outputs.logits

									# verify the logits
									snake_case   :			int						    = torch.Size((1, 21841)		)
									self.assertEqual(logits.shape   ,	_lowercase		)

									snake_case   :			Optional[Any]						    = torch.tensor([1.6881, -0.2787, 0.5901]		).to(_lowercase		)

									self.assertTrue(torch.allclose(logits[0, :3]   ,	_lowercase   ,	atol=1E-4		)		)

									snake_case   :			Dict						    = 2396
									self.assertEqual(logits.argmax(-1		).item()   ,	_lowercase		)



			@slow
			def 				__lowercase					(			self							:				int		)							->					Union[str, Any]:
									snake_case   :			Optional[Any]						    = BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"		)
									snake_case   :			Optional[int]						    = model.to(_lowercase		)

									snake_case   :			Optional[Any]						    = BeitImageProcessor(do_resize=_lowercase   ,	size=640   ,	do_center_crop=_lowercase		)

									snake_case   :			List[str]						    = load_dataset("hf-internal-testing/fixtures_ade20k"   ,	split="test"		)
									snake_case   :			Tuple						    = Image.open(ds[0]["file"]		)
									snake_case   :			List[Any]						    = image_processor(images=_lowercase   ,	return_tensors="pt"		).to(_lowercase		)

									# forward pass
									with torch.no_grad():
															snake_case   :			Union[str, Any]						    = model(**_lowercase		)
									snake_case   :			Any						    = outputs.logits

									# verify the logits
									snake_case   :			Union[str, Any]						    = torch.Size((1, 150, 160, 160)		)
									self.assertEqual(logits.shape   ,	_lowercase		)

									snake_case   :			Dict						    = version.parse(PIL.__version__		) < version.parse("9.0.0"		)

									if is_pillow_less_than_a:
															snake_case   :			Dict						    = torch.tensor(
															    [
															        [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
															        [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
															        [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
															    ]   ,	device=_lowercase   ,	)
									else:
															snake_case   :			List[Any]						    = torch.tensor(
															    [
															        [[-4.8960, -2.3688, -3.0355], [-2.8478, -0.9836, -1.7418], [-2.9449, -1.3332, -2.1456]],
															        [[-5.8081, -3.4124, -4.1006], [-3.8561, -2.2081, -3.0323], [-3.8365, -2.4601, -3.3669]],
															        [[-0.0309, 3.9868, 4.0540], [2.9640, 4.6877, 4.9976], [3.2081, 4.7690, 4.9942]],
															    ]   ,	device=_lowercase   ,	)

									self.assertTrue(torch.allclose(logits[0, :3, :3, :3]   ,	_lowercase   ,	atol=1E-4		)		)



			@slow
			def 				__lowercase					(			self							:				List[str]		)							->					Tuple:
									snake_case   :			str						    = BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"		)
									snake_case   :			Tuple						    = model.to(_lowercase		)

									snake_case   :			List[str]						    = BeitImageProcessor(do_resize=_lowercase   ,	size=640   ,	do_center_crop=_lowercase		)

									snake_case   :			Optional[Any]						    = load_dataset("hf-internal-testing/fixtures_ade20k"   ,	split="test"		)
									snake_case   :			Union[str, Any]						    = Image.open(ds[0]["file"]		)
									snake_case   :			List[str]						    = image_processor(images=_lowercase   ,	return_tensors="pt"		).to(_lowercase		)

									# forward pass
									with torch.no_grad():
															snake_case   :			Tuple						    = model(**_lowercase		)

									snake_case   :			Tuple						    = outputs.logits.detach().cpu()

									snake_case   :			Optional[Any]						    = image_processor.post_process_semantic_segmentation(outputs=_lowercase   ,	target_sizes=[(500, 300)]		)
									snake_case   :			List[str]						    = torch.Size((500, 300)		)
									self.assertEqual(segmentation[0].shape   ,	_lowercase		)

									snake_case   :			Optional[Any]						    = image_processor.post_process_semantic_segmentation(outputs=_lowercase		)
									snake_case   :			List[str]						    = torch.Size((160, 160)		)
									self.assertEqual(segmentation[0].shape   ,	_lowercase		)



