




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




