


'''simple docstring'''






from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


__A :  Tuple    							=				"Create a default config file for Accelerate with only a few flags set."







def     UpperCamelCase_      (   A__ :       Dict="no"							,   A__ :       Any = default_json_config_file							,   A__ :       Any = False ):
  lowerCAmelCase_      :				Dict      =       Path(lowerCAmelCase_ )
  path.parent.mkdir(parents=lowerCAmelCase_							,   exist_ok=lowerCAmelCase_ )
  if path.exists():
    print(
        f'Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.' )
    return False
  lowerCAmelCase_      :				List[str]      =       mixed_precision.lower()
  if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
    raise ValueError(
        f'`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}' )
  lowerCAmelCase_      :				Dict      =       {
      'compute_environment': 'LOCAL_MACHINE',
      'mixed_precision': mixed_precision,
  }
  if torch.cuda.is_available():
    lowerCAmelCase_      :				List[Any]      =       torch.cuda.device_count()
    lowerCAmelCase_      :				List[Any]      =       num_gpus
    lowerCAmelCase_      :				Optional[Any]      =       False
    if num_gpus > 1:
      lowerCAmelCase_      :				int      =       'MULTI_GPU'
    else:
      lowerCAmelCase_      :				Tuple      =       'NO'
  elif is_xpu_available() and use_xpu:
    lowerCAmelCase_      :				Optional[Any]      =       torch.xpu.device_count()
    lowerCAmelCase_      :				Optional[int]      =       num_xpus
    lowerCAmelCase_      :				Union[str, Any]      =       False
    if num_xpus > 1:
      lowerCAmelCase_      :				Tuple      =       'MULTI_XPU'
    else:
      lowerCAmelCase_      :				List[str]      =       'NO'
  elif is_npu_available():
    lowerCAmelCase_      :				str      =       torch.npu.device_count()
    lowerCAmelCase_      :				Union[str, Any]      =       num_npus
    lowerCAmelCase_      :				Tuple      =       False
    if num_npus > 1:
      lowerCAmelCase_      :				List[str]      =       'MULTI_NPU'
    else:
      lowerCAmelCase_      :				List[str]      =       'NO'
  else:
    lowerCAmelCase_      :				Optional[int]      =       0
    lowerCAmelCase_      :				Optional[int]      =       True
    lowerCAmelCase_      :				List[str]      =       1
    lowerCAmelCase_      :				Union[str, Any]      =       'NO'
  lowerCAmelCase_      :				Union[str, Any]      =       ClusterConfig(**lowerCAmelCase_ )
  config.to_json_file(lowerCAmelCase_ )
  return path







def     UpperCamelCase_      (   A__ :       int							,   A__ :       Optional[Any] ):
  lowerCAmelCase_      :				List[str]      =       parser.add_parser("""default"""							,   parents=lowerCAmelCase_							,   help=lowerCAmelCase_							,   formatter_class=lowerCAmelCase_ )
  parser.add_argument(
      """--config_file"""							,   default=lowerCAmelCase_							,   help=(
          """The path to use to store the config file. Will default to a file named default_config.yaml in the cache """
          """location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have """
          """such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed """
          """with \'huggingface\'."""
      )							,   dest="""save_location"""							,   )

  parser.add_argument(
      """--mixed_precision"""							,   choices=["""no""", """fp16""", """bf16"""]							,   type=lowerCAmelCase_							,   help="""Whether or not to use mixed precision training. """
      """Choose between FP16 and BF16 (bfloat16) training. """
      """BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later."""							,   default="""no"""							,   )
  parser.set_defaults(func=lowerCAmelCase_ )
  return parser







def     UpperCamelCase_      (   A__ :       Optional[Any] ):
  lowerCAmelCase_      :				str      =       write_basic_config(args.mixed_precision							,   args.save_location )
  if config_file:
    print(f'accelerate configuration saved at {config_file}' )







'''simple docstring'''






import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, 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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import ViTMAEForPreTraining, ViTMAEModel
  from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import ViTImageProcessor






class        __snake_case	:







    """simple docstring"""



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

      # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
      # (we add 1 for the [CLS] token)
      lowerCAmelCase_      :				Union[str, Any]      =       (image_size // patch_size) ** 2
      lowerCAmelCase_      :				Any      =       int(math.ceil((1 - mask_ratio) * (num_patches + 1)		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  str:
      lowerCAmelCase_      :				str      =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

      lowerCAmelCase_      :				Optional[int]      =       None
      if self.use_labels:
        lowerCAmelCase_      :				int      =       ids_tensor([self.batch_size]   ,						self.type_sequence_label_size		)

      lowerCAmelCase_      :				str      =       self.get_config()

      return config, pixel_values, labels

    def __lowercase       (       self   :  Optional[int]		)       ->  Optional[int]:
      return ViTMAEConfig(
          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=lowerCamelCase   ,						initializer_range=self.initializer_range   ,						mask_ratio=self.mask_ratio   ,						)

    def __lowercase       (       self   :  Any   ,						lowerCamelCase   :  Any   ,						lowerCamelCase   :  Union[str, Any]   ,						lowerCamelCase   :  Dict		)       ->  Tuple:
      lowerCAmelCase_      :				Tuple      =       ViTMAEModel(config=lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.eval()
      lowerCAmelCase_      :				Dict      =       model(lowerCamelCase		)
      self.parent.assertEqual(result.last_hidden_state.shape   ,						(self.batch_size, self.seq_length, self.hidden_size)		)

    def __lowercase       (       self   :  List[str]   ,						lowerCamelCase   :  Union[str, Any]   ,						lowerCamelCase   :  List[str]   ,						lowerCamelCase   :  Union[str, Any]		)       ->  Dict:
      lowerCAmelCase_      :				Tuple      =       ViTMAEForPreTraining(lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.eval()
      lowerCAmelCase_      :				List[str]      =       model(lowerCamelCase		)
      lowerCAmelCase_      :				int      =       (self.image_size // self.patch_size) ** 2
      lowerCAmelCase_      :				int      =       self.patch_size**2 * self.num_channels
      self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, num_patches, expected_num_channels)		)

      # test greyscale images
      lowerCAmelCase_      :				List[Any]      =       1
      lowerCAmelCase_      :				List[str]      =       ViTMAEForPreTraining(lowerCamelCase		)
      model.to(lowerCamelCase		)
      model.eval()
      lowerCAmelCase_      :				Optional[Any]      =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
      lowerCAmelCase_      :				Tuple      =       model(lowerCamelCase		)
      lowerCAmelCase_      :				List[Any]      =       self.patch_size**2
      self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, num_patches, expected_num_channels)		)





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







@require_torch
class        __snake_case	(       _SCREAMING_SNAKE_CASE     ,_SCREAMING_SNAKE_CASE     ,unittest.TestCase):







    """simple docstring"""



    lowercase       			= (ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
    lowercase       			= {'feature-extraction': ViTMAEModel} if is_torch_available() else {}

    lowercase       			= False
    lowercase       			= False
    lowercase       			= False
    lowercase       			= False

    def __lowercase       (       self   :  Optional[Any]		)       ->  List[Any]:
      lowerCAmelCase_      :				Optional[int]      =       ViTMAEModelTester(self		)
      lowerCAmelCase_      :				Optional[int]      =       ConfigTester(self   ,						config_class=lowerCamelCase   ,						has_text_modality=lowerCamelCase   ,						hidden_size=37		)

    def __lowercase       (       self   :  Dict		)       ->  Tuple:
      self.config_tester.run_common_tests()

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

    def __lowercase       (       self   :  List[str]		)       ->  Tuple:
      lowerCAmelCase_,							lowerCAmelCase_      :				Optional[Any]      =       self.model_tester.prepare_config_and_inputs_for_common()

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

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

      for model_class in self.all_model_classes:
        lowerCAmelCase_      :				Optional[int]      =       model_class(lowerCamelCase		)
        lowerCAmelCase_      :				Optional[Any]      =       inspect.signature(model.forward		)
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        lowerCAmelCase_      :				Any      =       [*signature.parameters.keys()]

        lowerCAmelCase_      :				Optional[Any]      =       ["""pixel_values"""]
        self.assertListEqual(arg_names[:1]   ,						lowerCamelCase		)

    def __lowercase       (       self   :  Tuple		)       ->  str:
      lowerCAmelCase_      :				Any      =       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCamelCase		)

    def __lowercase       (       self   :  Optional[int]		)       ->  str:
      lowerCAmelCase_      :				Tuple      =       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_pretraining(*lowerCamelCase		)

    def __lowercase       (       self   :  Optional[int]   ,						lowerCamelCase   :  Dict   ,						lowerCamelCase   :  List[Any]   ,						lowerCamelCase   :  Optional[Any]		)       ->  str:
      # make masks reproducible
      np.random.seed(2		)

      lowerCAmelCase_      :				Tuple      =       int((pt_model.config.image_size // pt_model.config.patch_size) ** 2		)
      lowerCAmelCase_      :				Optional[int]      =       np.random.uniform(size=(self.model_tester.batch_size, num_patches)		)
      lowerCAmelCase_      :				Optional[Any]      =       torch.from_numpy(lowerCamelCase		)

      # Add `noise` argument.
      # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
      lowerCAmelCase_      :				int      =       pt_noise

      super().check_pt_tf_models(lowerCamelCase   ,						lowerCamelCase   ,						lowerCamelCase		)

    def __lowercase       (       self   :  int		)       ->  Dict:
      lowerCAmelCase_,							lowerCAmelCase_      :				Optional[Any]      =       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCAmelCase_      :				Optional[int]      =       model_class(lowerCamelCase		)
        model.to(lowerCamelCase		)
        model.eval()
        # make random mask reproducible
        torch.manual_seed(2		)
        with torch.no_grad():
          lowerCAmelCase_      :				Any      =       model(**self._prepare_for_class(lowerCamelCase   ,						lowerCamelCase		)		)

        lowerCAmelCase_      :				Any      =       outputs[0].cpu().numpy()
        lowerCAmelCase_      :				List[str]      =       0

        with tempfile.TemporaryDirectory() as tmpdirname:
          model.save_pretrained(lowerCamelCase		)
          lowerCAmelCase_      :				int      =       model_class.from_pretrained(lowerCamelCase		)
          model.to(lowerCamelCase		)
          # make random mask reproducible
          torch.manual_seed(2		)
          with torch.no_grad():
            lowerCAmelCase_      :				str      =       model(**self._prepare_for_class(lowerCamelCase   ,						lowerCamelCase		)		)

          # Make sure we don't have nans
          lowerCAmelCase_      :				Optional[Any]      =       after_outputs[0].cpu().numpy()
          lowerCAmelCase_      :				str      =       0
          lowerCAmelCase_      :				List[Any]      =       np.amax(np.abs(out_a - out_a		)		)
          self.assertLessEqual(lowerCamelCase   ,						1E-5		)

    @unittest.skip(
        reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""		)
    def __lowercase       (       self   :  Optional[int]		)       ->  List[Any]:
      pass

    @unittest.skip(
        reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""		)
    def __lowercase       (       self   :  Union[str, Any]		)       ->  str:
      pass

    @unittest.skip(
        reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""		)
    def __lowercase       (       self   :  Optional[Any]		)       ->  Union[str, Any]:
      pass

    @unittest.skip(reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"""		)
    def __lowercase       (       self   :  Tuple		)       ->  Optional[Any]:
      pass

    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""		)
    def __lowercase       (       self   :  List[Any]		)       ->  str:
      pass





    @slow
    def __lowercase       (       self   :  List[str]		)       ->  List[Any]:
      for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowerCAmelCase_      :				List[Any]      =       ViTMAEModel.from_pretrained(lowerCamelCase		)
        self.assertIsNotNone(lowerCamelCase		)







def     UpperCamelCase_      (   ):
  '''simple docstring'''





  lowerCAmelCase_      :				List[str]      =       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
  return image






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







    """simple docstring"""



    @cached_property
    def __lowercase       (       self   :  Union[str, Any]		)       ->  Union[str, Any]:
      return ViTImageProcessor.from_pretrained("""facebook/vit-mae-base"""		) if is_vision_available() else None





    @slow
    def __lowercase       (       self   :  int		)       ->  List[Any]:
      # make random mask reproducible across the PT and TF model
      np.random.seed(2		)

      lowerCAmelCase_      :				Dict      =       ViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base"""		).to(lowerCamelCase		)

      lowerCAmelCase_      :				Union[str, Any]      =       self.default_image_processor
      lowerCAmelCase_      :				Union[str, Any]      =       prepare_img()
      lowerCAmelCase_      :				Dict      =       image_processor(images=lowerCamelCase   ,						return_tensors="""pt"""		).to(lowerCamelCase		)

      # prepare a noise vector that will be also used for testing the TF model
      # (this way we can ensure that the PT and TF models operate on the same inputs)
      lowerCAmelCase_      :				Optional[int]      =       ViTMAEConfig()
      lowerCAmelCase_      :				Optional[Any]      =       int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2		)
      lowerCAmelCase_      :				Optional[int]      =       np.random.uniform(size=(1, num_patches)		)

      # forward pass
      with torch.no_grad():
        lowerCAmelCase_      :				str      =       model(**lowerCamelCase   ,						noise=torch.from_numpy(lowerCamelCase		).to(device=lowerCamelCase		)		)

      # verify the logits
      lowerCAmelCase_      :				str      =       torch.Size((1, 1_96, 7_68)		)
      self.assertEqual(outputs.logits.shape   ,						lowerCamelCase		)

      lowerCAmelCase_      :				str      =       torch.tensor(
          [[-0.0_548, -1.7_023, -0.9_325], [0.3_721, -0.5_670, -0.2_233], [0.8_235, -1.3_878, -0.3_524]]		)

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



