




'''simple docstring'''
import inspect
import unittest

from transformers import DataaVecVisionConfig
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,
      DataaVecVisionForImageClassification,
      DataaVecVisionForSemanticSegmentation,
      DataaVecVisionModel,
  )
  from transformers.models.dataavec.modeling_dataavec_vision import DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import BeitImageProcessor






class 	SCREAMING_SNAKE_CASE:

       def __init__(   self ,		lowerCamelCase__ ,		lowerCamelCase__=100 ,		lowerCamelCase__=13 ,		lowerCamelCase__=30 ,		lowerCamelCase__=2 ,		lowerCamelCase__=3 ,		lowerCamelCase__=True ,		lowerCamelCase__=True ,		lowerCamelCase__=32 ,		lowerCamelCase__=4 ,		lowerCamelCase__=4 ,		lowerCamelCase__=37 ,		lowerCamelCase__="gelu" ,		lowerCamelCase__=0.1 ,		lowerCamelCase__=0.1 ,		lowerCamelCase__=10 ,		lowerCamelCase__=0.02 ,		lowerCamelCase__=3 ,		lowerCamelCase__=None ,		lowerCamelCase__=[0, 1, 2, 3] ,		)   -> str:

           """simple docstring"""



           __lowercase          =						parent
           __lowercase          =						100
           __lowercase          =						batch_size
           __lowercase          =						image_size
           __lowercase          =						patch_size
           __lowercase          =						num_channels
           __lowercase          =						is_training
           __lowercase          =						use_labels
           __lowercase          =						hidden_size
           __lowercase          =						num_hidden_layers
           __lowercase          =						num_attention_heads
           __lowercase          =						intermediate_size
           __lowercase          =						hidden_act
           __lowercase          =						hidden_dropout_prob
           __lowercase          =						attention_probs_dropout_prob
           __lowercase          =						type_sequence_label_size
           __lowercase          =						initializer_range
           __lowercase          =						scope
           __lowercase          =						out_indices
           __lowercase          =						num_labels

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

       def    snake_case__      (   self					)   -> int:

           """simple docstring"""



           __lowercase          =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

           __lowercase          =						None
           __lowercase          =						None
           if self.use_labels:
               __lowercase          =						ids_tensor([self.batch_size] ,		self.type_sequence_label_size					)
               __lowercase          =						ids_tensor([self.batch_size, self.image_size, self.image_size] ,		self.num_labels					)

           __lowercase          =						self.get_config()

           return config, pixel_values, labels, pixel_labels

       def    snake_case__      (   self					)   -> Dict:

           """simple docstring"""



           return DataaVecVisionConfig(
               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=lowerCamelCase__ ,		initializer_range=self.initializer_range ,		out_indices=self.out_indices ,		)

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__					)   -> Tuple:

           """simple docstring"""



           __lowercase          =						DataaVecVisionModel(config=lowerCamelCase__					)
           model.to(lowerCamelCase__					)
           model.eval()
           __lowercase          =						model(lowerCamelCase__					)
           # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
           __lowercase          =						(self.image_size // self.patch_size) ** 2
           self.parent.assertEqual(result.last_hidden_state.shape ,		(self.batch_size, num_patches + 1, self.hidden_size)					)

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__					)   -> int:

           """simple docstring"""



           __lowercase          =						self.type_sequence_label_size
           __lowercase          =						DataaVecVisionForImageClassification(lowerCamelCase__					)
           model.to(lowerCamelCase__					)
           model.eval()
           __lowercase          =						model(lowerCamelCase__ ,		labels=lowerCamelCase__					)
           self.parent.assertEqual(result.logits.shape ,		(self.batch_size, self.type_sequence_label_size)					)

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__					)   -> Optional[int]:

           """simple docstring"""



           __lowercase          =						self.num_labels
           __lowercase          =						DataaVecVisionForSemanticSegmentation(lowerCamelCase__					)
           model.to(lowerCamelCase__					)
           model.eval()
           __lowercase          =						model(lowerCamelCase__					)
           self.parent.assertEqual(
               result.logits.shape ,		(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)					)
           __lowercase          =						model(lowerCamelCase__ ,		labels=lowerCamelCase__					)
           self.parent.assertEqual(
               result.logits.shape ,		(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)					)




       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           __lowercase          =						self.prepare_config_and_inputs()
           __lowercase    ,__lowercase    ,__lowercase    ,__lowercase          =						config_and_inputs
           __lowercase          =						{"""pixel_values""": pixel_values}
           return config, inputs_dict






@require_torch
class 	SCREAMING_SNAKE_CASE(    __A       ,				__A       ,				unittest.TestCase					):
       snake_case_			:					Dict   = (
           (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
           if is_torch_available()
           else ()
       )
       snake_case_			:					Optional[int]   = (
           {
               """feature-extraction""": DataaVecVisionModel,
               """image-classification""": DataaVecVisionForImageClassification,
               """image-segmentation""": DataaVecVisionForSemanticSegmentation,
           }
           if is_torch_available()
           else {}
       )

       snake_case_			:					int   = False
       snake_case_			:					Union[str, Any]   = False
       snake_case_			:					Tuple   = False

       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           __lowercase          =						DataaVecVisionModelTester(self					)
           __lowercase          =						ConfigTester(
               self ,		config_class=lowerCamelCase__ ,		has_text_modality=lowerCamelCase__ ,		hidden_size=37					)

       def    snake_case__      (   self					)   -> Optional[int]:

           """simple docstring"""



           self.config_tester.run_common_tests()

       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           pass

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

           """simple docstring"""



           pass

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

           """simple docstring"""



           __lowercase    ,__lowercase          =						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               __lowercase          =						model_class(lowerCamelCase__					)
               self.assertIsInstance(model.get_input_embeddings() ,		(nn.Module)					)
               __lowercase          =						model.get_output_embeddings()
               self.assertTrue(x is None or isinstance(lowerCamelCase__ ,		nn.Linear					)					)

       def    snake_case__      (   self					)   -> Optional[int]:

           """simple docstring"""



           __lowercase    ,__lowercase          =						self.model_tester.prepare_config_and_inputs_for_common()

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

               __lowercase          =						["""pixel_values"""]
               self.assertListEqual(arg_names[:1] ,		lowerCamelCase__					)

       def    snake_case__      (   self					)   -> List[Any]:

           """simple docstring"""



           __lowercase          =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase__					)

       def    snake_case__      (   self					)   -> int:

           """simple docstring"""



           __lowercase          =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_segmentation(*lowerCamelCase__					)

       def    snake_case__      (   self					)   -> int:

           """simple docstring"""



           if not self.model_tester.is_training:
               return

           __lowercase    ,__lowercase          =						self.model_tester.prepare_config_and_inputs_for_common()
           __lowercase          =						True

           for model_class in self.all_model_classes:
               if model_class in [*get_values(lowerCamelCase__					)]:
                   continue

               __lowercase          =						model_class(lowerCamelCase__					)
               model.to(lowerCamelCase__					)
               model.train()
               __lowercase          =						self._prepare_for_class(lowerCamelCase__ ,		lowerCamelCase__ ,		return_labels=lowerCamelCase__					)
               __lowercase          =						model(**lowerCamelCase__					).loss
               loss.backward()

       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



           __lowercase    ,__lowercase          =						self.model_tester.prepare_config_and_inputs_for_common()
           if not self.model_tester.is_training:
               return

           __lowercase          =						False
           __lowercase          =						True

           for model_class in self.all_model_classes:
               if model_class in [*get_values(lowerCamelCase__					)] or not model_class.supports_gradient_checkpointing:
                   continue
               # TODO: remove the following 3 lines once we have a MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING
               # this can then be incorporated into _prepare_for_class in test_modeling_common.py
               elif model_class.__name__ == "Data2VecVisionForSemanticSegmentation":
                   __lowercase    ,__lowercase    ,__lowercase    ,__lowercase          =						inputs_dict["""pixel_values"""].shape
                   __lowercase          =						torch.zeros(
                       [self.model_tester.batch_size, height, width] ,		device=lowerCamelCase__					).long()
               __lowercase          =						model_class(lowerCamelCase__					)
               model.gradient_checkpointing_enable()
               model.to(lowerCamelCase__					)
               model.train()
               __lowercase          =						self._prepare_for_class(lowerCamelCase__ ,		lowerCamelCase__ ,		return_labels=lowerCamelCase__					)
               __lowercase          =						model(**lowerCamelCase__					).loss
               loss.backward()

       def    snake_case__      (   self					)   -> Optional[Any]:

           """simple docstring"""



           __lowercase    ,__lowercase          =						self.model_tester.prepare_config_and_inputs_for_common()

           __lowercase          =						_config_zero_init(lowerCamelCase__					)
           for model_class in self.all_model_classes:
               __lowercase          =						model_class(config=lowerCamelCase__					)
               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' ,		)

       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__=2E-4 ,		lowerCamelCase__="outputs" ,		lowerCamelCase__=None					)   -> Union[str, Any]:

           """simple docstring"""



           super().check_pt_tf_outputs(lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__ ,		lowerCamelCase__					)

       def    snake_case__      (   self					)   -> Optional[Any]:

           """simple docstring"""



           __lowercase          =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__					)




       @slow
       def    snake_case__      (   self					)   -> Optional[int]:

           """simple docstring"""



           for model_name in DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __lowercase          =						DataaVecVisionModel.from_pretrained(lowerCamelCase__					)
               self.assertIsNotNone(lowerCamelCase__					)

def     snake_case_  (       ):



    """simple docstring"""




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






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

       @cached_property
       def    snake_case__      (   self					)   -> Union[str, Any]:

           """simple docstring"""



           return (
               BeitImageProcessor.from_pretrained("""facebook/data2vec-vision-base-ft1k"""					) if is_vision_available() else None
           )




       @slow
       def    snake_case__      (   self					)   -> str:

           """simple docstring"""



           __lowercase          =						DataaVecVisionForImageClassification.from_pretrained("""facebook/data2vec-vision-base-ft1k"""					).to(
               lowerCamelCase__					)

           __lowercase          =						self.default_image_processor
           __lowercase          =						prepare_img()
           __lowercase          =						image_processor(images=lowerCamelCase__ ,		return_tensors="""pt"""					).to(lowerCamelCase__					)

           # forward pass
           with torch.no_grad():
               __lowercase          =						model(**lowerCamelCase__					)
           __lowercase          =						outputs.logits

           # verify the logits
           __lowercase          =						torch.Size((1, 1000)					)
           self.assertEqual(logits.shape ,		lowerCamelCase__					)

           __lowercase          =						torch.tensor([0.32_77, -0.13_95, 0.09_11]					).to(lowerCamelCase__					)

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

           __lowercase          =						[model.config.labelaid[i] for i in ["""remote control, remote""", """tabby, tabby cat"""]]
           self.assertEqual(logits[0].topk(2					).indices.cpu().tolist() ,		lowerCamelCase__					)










'''simple docstring'''
import random
from typing import List, Optional, Union

import numpy as np
import torch

from ..state import AcceleratorState
from .constants import CUDA_DISTRIBUTED_TYPES
from .dataclasses import DistributedType, RNGType
from .imports import is_npu_available, is_tpu_available, is_xpu_available


if is_tpu_available(check_device=False):
  import torch_xla.core.xla_model as xm
def     snake_case_  (       a__		:     int   ,a__		:     bool = False      ):



    """simple docstring"""




    if device_specific:
        seed += AcceleratorState().process_index
    random.seed(a__      )
    np.random.seed(a__      )
    torch.manual_seed(a__      )
    if is_xpu_available():
        torch.xpu.manual_seed_all(a__      )
    elif is_npu_available():
        torch.npu.manual_seed_all(a__      )
    else:
        torch.cuda.manual_seed_all(a__      )
    # ^^ safe to call this function even if cuda is not available
    if is_tpu_available():
        xm.set_rng_state(a__      )
def     snake_case_  (       a__		:     Optional[RNGType] = None   ,a__		:     Optional[torch.Generator] = None      ):



    """simple docstring"""




    if rng_type == RNGType.TORCH:
        __lowercase          =						torch.get_rng_state()
    elif rng_type == RNGType.CUDA:
        __lowercase          =						torch.cuda.get_rng_state()
    elif rng_type == RNGType.XLA:
        assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
        __lowercase          =						torch.tensor(xm.get_rng_state()      )
    elif rng_type == RNGType.NPU:
        assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs."
        __lowercase          =						torch.npu.get_rng_state()
    elif rng_type == RNGType.XPU:
        assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
        __lowercase          =						torch.xpu.get_rng_state()
    elif rng_type == RNGType.GENERATOR:
        assert generator is not None, "Need a generator to synchronize its seed."
        __lowercase          =						generator.get_state()

    # Broadcast the rng state from device 0 to other devices
    __lowercase          =						AcceleratorState()
    if state.distributed_type == DistributedType.TPU:
        __lowercase          =						rng_state.to(xm.xla_device()      )
        xm.collective_broadcast([rng_state]      )
        xm.mark_step()
        __lowercase          =						rng_state.cpu()
    elif (
        state.distributed_type in CUDA_DISTRIBUTED_TYPES
        or state.distributed_type == DistributedType.MULTI_NPU
        or state.distributed_type == DistributedType.MULTI_XPU
    ):
        __lowercase          =						rng_state.to(state.device      )
        torch.distributed.broadcast(a__   ,0      )
        __lowercase          =						rng_state.cpu()
    elif state.distributed_type == DistributedType.MULTI_CPU:
        torch.distributed.broadcast(a__   ,0      )

    # Set the broadcast rng state
    if rng_type == RNGType.TORCH:
        torch.set_rng_state(a__      )
    elif rng_type == RNGType.CUDA:
        torch.cuda.set_rng_state(a__      )
    elif rng_type == RNGType.NPU:
        torch.npu.set_rng_state(a__      )
    elif rng_type == RNGType.XPU:
        torch.xpu.set_rng_state(a__      )
    elif rng_type == RNGType.XLA:
        xm.set_rng_state(rng_state.item()      )
    elif rng_type == RNGType.GENERATOR:
        generator.set_state(a__      )



def     snake_case_  (       a__		:     List[Union[str, RNGType]]   ,a__		:     Optional[torch.Generator] = None      ):



    """simple docstring"""




    for rng_type in rng_types:
        synchronize_rng_state(RNGType(a__      )   ,generator=a__      )




