





import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

						if is_vision_available():
												from transformers import MaskaFormerImageProcessor

if is_vision_available():
						from PIL import Image
class SCREAMING_SNAKE_CASE_						:






						def __init__(      self      :      Union[str, Any]		,   _A      :      str		,   _A      :      Union[str, Any]=2		,   _A      :      Dict=True		,   _A      :      Union[str, Any]=False		,   _A      :      Any=10		,   _A      :      str=3		,   _A      :      Optional[Any]=32 * 8		,   _A      :      Union[str, Any]=32 * 8		,   _A      :      List[str]=4		,   _A      :      int=64		,   )							->    Union[str, Any]:







							"""simple docstring"""


							snake_case_		:			Any											=	parent
							snake_case_		:			List[Any]											=	batch_size
							snake_case_		:			Optional[int]											=	is_training
							snake_case_		:			Any											=	use_auxiliary_loss
							snake_case_		:			List[Any]											=	num_queries
							snake_case_		:			List[str]											=	num_channels
							snake_case_		:			int											=	min_size
							snake_case_		:			Optional[Any]											=	max_size
							snake_case_		:			Union[str, Any]											=	num_labels
							snake_case_		:			str											=	hidden_dim
							snake_case_		:			Tuple											=	hidden_dim






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







							"""simple docstring"""


							snake_case_		:			List[Any]											=	floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]	).to(
							    A__	)

							snake_case_		:			Any											=	torch.ones([self.batch_size, self.min_size, self.max_size]		,   device=A__	)

							snake_case_		:			Dict											=	(
							    torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]		,   device=A__	) > 0.5
							).float()
							snake_case_		:			Dict											=	(torch.rand((self.batch_size, self.num_labels)		,   device=A__	) > 0.5).long()

							snake_case_		:			Tuple											=	self.get_config()
							return config, pixel_values, pixel_mask, mask_labels, class_labels






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







							"""simple docstring"""


							snake_case_		:			List[str]											=	MaskaFormerConfig(
							    hidden_size=self.hidden_dim		,   )
							snake_case_		:			int											=	self.num_queries
							snake_case_		:			List[str]											=	self.num_labels

							snake_case_		:			Dict											=	[1, 1, 1, 1]
							snake_case_		:			Optional[int]											=	self.num_channels

							snake_case_		:			List[str]											=	64
							snake_case_		:			str											=	128
							snake_case_		:			Any											=	self.hidden_dim
							snake_case_		:			Dict											=	self.hidden_dim
							snake_case_		:			str											=	self.hidden_dim
							return config






						def UpperCAmelCase_				(      self      :      List[Any]	)							->    Dict:







							"""simple docstring"""


							snake_case_						,snake_case_						,snake_case_						,snake_case_						,snake_case_		:			Optional[int]											=	self.prepare_config_and_inputs()
							snake_case_		:			Tuple											=	{'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
							return config, inputs_dict






						def UpperCAmelCase_				(      self      :      Union[str, Any]		,   _A      :      Optional[Any]		,   _A      :      Optional[int]	)							->    Union[str, Any]:







							"""simple docstring"""


							snake_case_		:			List[str]											=	output.encoder_hidden_states
							snake_case_		:			List[str]											=	output.pixel_decoder_hidden_states
							snake_case_		:			int											=	output.transformer_decoder_hidden_states

							self.parent.assertTrue(len(A__	)		,   len(config.backbone_config.depths	)	)
							self.parent.assertTrue(len(A__	)		,   len(config.backbone_config.depths	)	)
							self.parent.assertTrue(len(A__	)		,   config.decoder_layers	)






						def UpperCAmelCase_				(      self      :      List[Any]		,   _A      :      Any		,   _A      :      str		,   _A      :      Tuple		,   _A      :      int=False	)							->    str:







							"""simple docstring"""


							with torch.no_grad():
								snake_case_		:			Any											=	MaskaFormerModel(config=A__	)
								model.to(A__	)
								model.eval()

								snake_case_		:			List[Any]											=	model(pixel_values=A__		,   pixel_mask=A__	)
								snake_case_		:			Any											=	model(A__		,   output_hidden_states=A__	)

							self.parent.assertEqual(
							    output.transformer_decoder_last_hidden_state.shape		,   (self.batch_size, self.num_queries, self.hidden_dim)		,   )
							# let's ensure the other two hidden state exists
							self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None	)
							self.parent.assertTrue(output.encoder_last_hidden_state is not None	)

							if output_hidden_states:
								self.check_output_hidden_state(A__		,   A__	)






						def UpperCAmelCase_				(      self      :      str		,   _A      :      Union[str, Any]		,   _A      :      List[str]		,   _A      :      Any		,   _A      :      int		,   _A      :      int	)							->    List[str]:







							"""simple docstring"""


							snake_case_		:			Optional[Any]											=	MaskaFormerForUniversalSegmentation(config=A__	)
							model.to(A__	)
							model.eval()

							def comm_check_on_output(_A      :      List[Any]	):
								# let's still check that all the required stuff is there
								self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None	)
								self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None	)
								self.parent.assertTrue(result.encoder_last_hidden_state is not None	)
								# okay, now we need to check the logits shape
								# due to the encoder compression, masks have a //4 spatial size
								self.parent.assertEqual(
								    result.masks_queries_logits.shape		,   (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)		,   )
								# + 1 for null class
								self.parent.assertEqual(
								    result.class_queries_logits.shape		,   (self.batch_size, self.num_queries, self.num_labels + 1)	)

							with torch.no_grad():
								snake_case_		:			int											=	model(pixel_values=A__		,   pixel_mask=A__	)
								snake_case_		:			Union[str, Any]											=	model(A__	)

								comm_check_on_output(A__	)

								snake_case_		:			List[str]											=	model(
								    pixel_values=A__		,   pixel_mask=A__		,   mask_labels=A__		,   class_labels=A__	)

							comm_check_on_output(A__	)

							self.parent.assertTrue(result.loss is not None	)
							self.parent.assertEqual(result.loss.shape		,   torch.Size([1]	)	)
@require_torch
class SCREAMING_SNAKE_CASE_						(					__a       ,				__a       ,				unittest.TestCase					):
						__magic_name__:			Tuple       =       (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
						__magic_name__:			Optional[int]       =       {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

						__magic_name__:			Optional[Any]       =       False
						__magic_name__:			int       =       False
						__magic_name__:			Any       =       False
						__magic_name__:			Dict       =       False






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







							"""simple docstring"""


							snake_case_		:			List[Any]											=	MaskaFormerModelTester(self	)
							snake_case_		:			Tuple											=	ConfigTester(self		,   config_class=A__		,   has_text_modality=A__	)






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







							"""simple docstring"""


							self.config_tester.run_common_tests()






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







							"""simple docstring"""


							snake_case_						,snake_case_		:			Union[str, Any]											=	self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.create_and_check_maskaformer_model(A__		,   **A__		,   output_hidden_states=A__	)






						def UpperCAmelCase_				(      self      :      int	)							->    str:







							"""simple docstring"""


							snake_case_		:			List[Any]											=	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*A__	)






						@unittest.skip(reason='Mask2Former does not use inputs_embeds'	)
						def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    Dict:







							"""simple docstring"""


							pass






						@unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'	)
						def UpperCAmelCase_				(      self      :      List[Any]	)							->    List[Any]:







							"""simple docstring"""


							pass






						@unittest.skip(reason='Mask2Former is not a generative model'	)
						def UpperCAmelCase_				(      self      :      List[str]	)							->    Optional[Any]:







							"""simple docstring"""


							pass






						@unittest.skip(reason='Mask2Former does not use token embeddings'	)
						def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    List[str]:







							"""simple docstring"""


							pass






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







							"""simple docstring"""


							pass






						@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'	)
						def UpperCAmelCase_				(      self      :      Union[str, Any]	)							->    str:







							"""simple docstring"""


							pass






						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    str:







							"""simple docstring"""


							snake_case_						,snake_case_		:			Optional[int]											=	self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								snake_case_		:			Union[str, Any]											=	model_class(A__	)
								snake_case_		:			List[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_		:			Tuple											=	['pixel_values']
								self.assertListEqual(arg_names[:1]		,   A__	)






						@slow
						def UpperCAmelCase_				(      self      :      Any	)							->    List[str]:







							"""simple docstring"""


							for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
								snake_case_		:			Union[str, Any]											=	MaskaFormerModel.from_pretrained(A__	)
								self.assertIsNotNone(A__	)






						def UpperCAmelCase_				(      self      :      int	)							->    Dict:







							"""simple docstring"""


							snake_case_		:			Optional[Any]											=	(self.model_tester.min_size,) * 2
							snake_case_		:			List[Any]											=	{
							    'pixel_values': torch.randn((2, 3, *size)		,   device=A__	),
							    'mask_labels': torch.randn((2, 10, *size)		,   device=A__	),
							    'class_labels': torch.zeros(2		,   10		,   device=A__	).long(),
							}
							snake_case_		:			int											=	self.model_tester.get_config()

							snake_case_		:			Any											=	MaskaFormerForUniversalSegmentation(A__	).to(A__	)
							snake_case_		:			Tuple											=	model(**A__	)
							self.assertTrue(outputs.loss is not None	)






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







							"""simple docstring"""


							snake_case_						,snake_case_		:			Optional[int]											=	self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.create_and_check_maskaformer_model(A__		,   **A__		,   output_hidden_states=A__	)






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







							"""simple docstring"""


							snake_case_						,snake_case_		:			Any											=	self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								snake_case_		:			Optional[Any]											=	model_class(A__	).to(A__	)
								snake_case_		:			str											=	model(**A__		,   output_attentions=A__	)
								self.assertTrue(outputs.attentions is not None	)






						def UpperCAmelCase_				(      self      :      List[str]	)							->    int:







							"""simple docstring"""


							if not self.model_tester.is_training:
								return

							snake_case_		:			Tuple											=	self.all_model_classes[1]
							snake_case_						,snake_case_						,snake_case_						,snake_case_						,snake_case_		:			Union[str, Any]											=	self.model_tester.prepare_config_and_inputs()

							snake_case_		:			int											=	model_class(A__	)
							model.to(A__	)
							model.train()

							snake_case_		:			Dict											=	model(A__		,   mask_labels=A__		,   class_labels=A__	).loss
							loss.backward()






						def UpperCAmelCase_				(      self      :      str	)							->    Optional[Any]:







							"""simple docstring"""


							snake_case_		:			Optional[Any]											=	self.all_model_classes[1]
							snake_case_						,snake_case_						,snake_case_						,snake_case_						,snake_case_		:			str											=	self.model_tester.prepare_config_and_inputs()
							snake_case_		:			int											=	True
							snake_case_		:			Dict											=	True

							snake_case_		:			int											=	model_class(A__	).to(A__	)
							model.train()

							snake_case_		:			List[str]											=	model(A__		,   mask_labels=A__		,   class_labels=A__	)

							snake_case_		:			Union[str, Any]											=	outputs.encoder_hidden_states[0]
							encoder_hidden_states.retain_grad()

							snake_case_		:			Union[str, Any]											=	outputs.pixel_decoder_hidden_states[0]
							pixel_decoder_hidden_states.retain_grad()

							snake_case_		:			Optional[Any]											=	outputs.transformer_decoder_hidden_states[0]
							transformer_decoder_hidden_states.retain_grad()

							snake_case_		:			List[Any]											=	outputs.attentions[0]
							attentions.retain_grad()

							outputs.loss.backward(retain_graph=A__	)

							self.assertIsNotNone(encoder_hidden_states.grad	)
							self.assertIsNotNone(pixel_decoder_hidden_states.grad	)
							self.assertIsNotNone(transformer_decoder_hidden_states.grad	)
							self.assertIsNotNone(attentions.grad	)
_SCREAMING_SNAKE_CASE       	=		1e-4


def 							SCREAMING_SNAKE_CASE__						(    ):
	snake_case_		:			Union[str, Any]											=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
	return image


@require_vision
@slow
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






						@cached_property
						def UpperCAmelCase_				(      self      :      Dict	)							->    Tuple:







							"""simple docstring"""


							return "facebook/mask2former-swin-small-coco-instance"






						@cached_property
						def UpperCAmelCase_				(      self      :      Dict	)							->    str:







							"""simple docstring"""


							return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints	) if is_vision_available() else None






						def UpperCAmelCase_				(      self      :      Tuple	)							->    Optional[int]:







							"""simple docstring"""


							snake_case_		:			Union[str, Any]											=	MaskaFormerModel.from_pretrained(self.model_checkpoints	).to(A__	)
							snake_case_		:			int											=	self.default_image_processor
							snake_case_		:			str											=	prepare_img()
							snake_case_		:			List[Any]											=	image_processor(A__		,   return_tensors='pt'	).to(A__	)
							snake_case_		:			Union[str, Any]											=	inputs['pixel_values'].shape
							# check size is divisible by 32
							self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0	)
							# check size
							self.assertEqual(A__		,   (1, 3, 384, 384)	)

							with torch.no_grad():
								snake_case_		:			Optional[int]											=	model(**A__	)

							snake_case_		:			str											=	torch.tensor(
							    [[-0.2_7_9_0, -1.0_7_1_7, -1.1_6_6_8], [-0.5_1_2_8, -0.3_1_2_8, -0.4_9_8_7], [-0.5_8_3_2, 0.1_9_7_1, -0.0_1_9_7]]	).to(A__	)
							self.assertTrue(
							    torch.allclose(
							        outputs.encoder_last_hidden_state[0, 0, :3, :3]		,   A__		,   atol=A__	)	)

							snake_case_		:			List[str]											=	torch.tensor(
							    [[0.8_9_7_3, 1.1_8_4_7, 1.1_7_7_6], [1.1_9_3_4, 1.5_0_4_0, 1.5_1_2_8], [1.1_1_5_3, 1.4_4_8_6, 1.4_9_5_1]]	).to(A__	)
							self.assertTrue(
							    torch.allclose(
							        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]		,   A__		,   atol=A__	)	)

							snake_case_		:			Optional[int]											=	torch.tensor(
							    [[2.1_1_5_2, 1.7_0_0_0, -0.8_6_0_3], [1.5_8_0_8, 1.8_0_0_4, -0.9_3_5_3], [1.6_0_4_3, 1.7_4_9_5, -0.5_9_9_9]]	).to(A__	)
							self.assertTrue(
							    torch.allclose(
							        outputs.transformer_decoder_last_hidden_state[0, :3, :3]		,   A__		,   atol=A__	)	)






						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    str:







							"""simple docstring"""


							snake_case_		:			List[str]											=	MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints	).to(A__	).eval()
							snake_case_		:			List[Any]											=	self.default_image_processor
							snake_case_		:			List[str]											=	prepare_img()
							snake_case_		:			List[Any]											=	image_processor(A__		,   return_tensors='pt'	).to(A__	)
							snake_case_		:			List[str]											=	inputs['pixel_values'].shape
							# check size is divisible by 32
							self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0	)
							# check size
							self.assertEqual(A__		,   (1, 3, 384, 384)	)

							with torch.no_grad():
								snake_case_		:			List[Any]											=	model(**A__	)
							# masks_queries_logits
							snake_case_		:			Any											=	outputs.masks_queries_logits
							self.assertEqual(
							    masks_queries_logits.shape		,   (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)	)
							snake_case_		:			List[str]											=	[
							    [-8.7_8_3_9, -9.0_0_5_6, -8.8_1_2_1],
							    [-7.4_1_0_4, -7.0_3_1_3, -6.5_4_0_1],
							    [-6.6_1_0_5, -6.3_4_2_7, -6.4_6_7_5],
							]
							snake_case_		:			Any											=	torch.tensor(A__	).to(A__	)
							self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]		,   A__		,   atol=A__	)	)
							# class_queries_logits
							snake_case_		:			List[str]											=	outputs.class_queries_logits
							self.assertEqual(class_queries_logits.shape		,   (1, model.config.num_queries, model.config.num_labels + 1)	)
							snake_case_		:			int											=	torch.tensor(
							    [
							        [1.8_3_2_4, -8.0_8_3_5, -4.1_9_2_2],
							        [0.8_4_5_0, -9.0_0_5_0, -3.6_0_5_3],
							        [0.3_0_4_5, -7.7_2_9_3, -3.0_2_7_5],
							    ]	).to(A__	)
							self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]		,   A__		,   atol=A__	)	)






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







							"""simple docstring"""


							snake_case_		:			Any											=	MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints	).to(A__	).eval()
							snake_case_		:			List[Any]											=	self.default_image_processor

							snake_case_		:			str											=	image_processor(
							    [np.zeros((3, 800, 1333)	), np.zeros((3, 800, 1333)	)]		,   segmentation_maps=[np.zeros((384, 384)	).astype(np.floataa	), np.zeros((384, 384)	).astype(np.floataa	)]		,   return_tensors='pt'		,   )

							snake_case_		:			List[Any]											=	inputs['pixel_values'].to(A__	)
							snake_case_		:			str											=	[el.to(A__	) for el in inputs['mask_labels']]
							snake_case_		:			Optional[int]											=	[el.to(A__	) for el in inputs['class_labels']]

							with torch.no_grad():
								snake_case_		:			Union[str, Any]											=	model(**A__	)

							self.assertTrue(outputs.loss is not None	)








import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
						import jax

						from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






						def __init__(      self      :      Dict		,   _A      :      List[str]		,   _A      :      Dict=13		,   _A      :      str=30		,   _A      :      Any=2		,   _A      :      Dict=3		,   _A      :      Optional[Any]=True		,   _A      :      Tuple=True		,   _A      :      List[str]=32		,   _A      :      int=5		,   _A      :      Optional[int]=4		,   _A      :      Dict=37		,   _A      :      Tuple="gelu"		,   _A      :      Dict=0.1		,   _A      :      str=0.1		,   _A      :      int=10		,   _A      :      Union[str, Any]=0.0_2		,   )							->    int:







							"""simple docstring"""


							snake_case_		:			int											=	parent
							snake_case_		:			str											=	batch_size
							snake_case_		:			List[str]											=	image_size
							snake_case_		:			Tuple											=	patch_size
							snake_case_		:			str											=	num_channels
							snake_case_		:			List[str]											=	is_training
							snake_case_		:			List[str]											=	use_labels
							snake_case_		:			str											=	hidden_size
							snake_case_		:			Union[str, Any]											=	num_hidden_layers
							snake_case_		:			Any											=	num_attention_heads
							snake_case_		:			List[Any]											=	intermediate_size
							snake_case_		:			Any											=	hidden_act
							snake_case_		:			List[Any]											=	hidden_dropout_prob
							snake_case_		:			List[Any]											=	attention_probs_dropout_prob
							snake_case_		:			Tuple											=	type_sequence_label_size
							snake_case_		:			Any											=	initializer_range

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






						def UpperCAmelCase_				(      self      :      Dict	)							->    List[str]:







							"""simple docstring"""


							snake_case_		:			int											=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

							snake_case_		:			Optional[Any]											=	ViTConfig(
							    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=_A		,   initializer_range=self.initializer_range		,   )

							return config, pixel_values






						def UpperCAmelCase_				(      self      :      int		,   _A      :      Union[str, Any]		,   _A      :      List[Any]	)							->    Dict:







							"""simple docstring"""


							snake_case_		:			Union[str, Any]											=	FlaxViTModel(config=_A	)
							snake_case_		:			Optional[Any]											=	model(_A	)
							# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
							snake_case_		:			Union[str, Any]											=	(self.image_size, self.image_size)
							snake_case_		:			int											=	(self.patch_size, self.patch_size)
							snake_case_		:			Optional[Any]											=	(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
							self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, num_patches + 1, self.hidden_size)	)






						def UpperCAmelCase_				(      self      :      str		,   _A      :      Any		,   _A      :      List[str]	)							->    Dict:







							"""simple docstring"""


							snake_case_		:			Dict											=	self.type_sequence_label_size
							snake_case_		:			Any											=	FlaxViTForImageClassification(config=_A	)
							snake_case_		:			Optional[int]											=	model(_A	)
							self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.type_sequence_label_size)	)

							# test greyscale images
							snake_case_		:			List[str]											=	1
							snake_case_		:			Dict											=	FlaxViTForImageClassification(_A	)

							snake_case_		:			Optional[int]											=	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
							snake_case_		:			Any											=	model(_A	)






						def UpperCAmelCase_				(      self      :      Any	)							->    Any:







							"""simple docstring"""


							snake_case_		:			Tuple											=	self.prepare_config_and_inputs()
							(
							    (
							    snake_case_
							)						,(
							    snake_case_
							)						,
							)		:			Any											=	config_and_inputs
							snake_case_		:			Tuple											=	{'pixel_values': pixel_values}
							return config, inputs_dict
@require_flax
class SCREAMING_SNAKE_CASE_						(					snake_case_       ,				unittest.TestCase					):
						__magic_name__:			Tuple       =       (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()






						def UpperCAmelCase_				(      self      :      List[str]	)							->    None:







							"""simple docstring"""


							snake_case_		:			Union[str, Any]											=	FlaxViTModelTester(self	)
							snake_case_		:			Any											=	ConfigTester(self		,   config_class=_A		,   has_text_modality=_A		,   hidden_size=37	)






						def UpperCAmelCase_				(      self      :      Any	)							->    Dict:







							"""simple docstring"""


							self.config_tester.run_common_tests()






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







							"""simple docstring"""


							snake_case_		:			str											=	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_A	)






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







							"""simple docstring"""


							snake_case_		:			Tuple											=	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_A	)






						def UpperCAmelCase_				(      self      :      int	)							->    List[str]:







							"""simple docstring"""


							snake_case_						,snake_case_		:			Tuple											=	self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								snake_case_		:			str											=	model_class(_A	)
								snake_case_		:			Dict											=	inspect.signature(model.__call__	)
								# signature.parameters is an OrderedDict => so arg_names order is deterministic
								snake_case_		:			Optional[Any]											=	[*signature.parameters.keys()]

								snake_case_		:			Dict											=	['pixel_values']
								self.assertListEqual(arg_names[:1]		,   _A	)






						def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    str:







							"""simple docstring"""


							snake_case_						,snake_case_		:			Optional[int]											=	self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								with self.subTest(model_class.__name__	):
									snake_case_		:			Dict											=	self._prepare_for_class(_A		,   _A	)
									snake_case_		:			Tuple											=	model_class(_A	)

									@jax.jit
									def model_jitted(_A      :      Optional[int]		,   **_A      :      Any	):
										return model(pixel_values=_A		,   **_A	)

									with self.subTest('JIT Enabled'	):
										snake_case_		:			Optional[Any]											=	model_jitted(**_A	).to_tuple()

									with self.subTest('JIT Disabled'	):
										with jax.disable_jit():
											snake_case_		:			Optional[Any]											=	model_jitted(**_A	).to_tuple()

									self.assertEqual(len(_A	)		,   len(_A	)	)
									for jitted_output, output in zip(_A		,   _A	):
										self.assertEqual(jitted_output.shape		,   output.shape	)






						@slow
						def UpperCAmelCase_				(      self      :      List[str]	)							->    Any:







							"""simple docstring"""


							for model_class_name in self.all_model_classes:
								snake_case_		:			List[Any]											=	model_class_name.from_pretrained('google/vit-base-patch16-224'	)
								snake_case_		:			Optional[Any]											=	model(np.ones((1, 3, 224, 224)	)	)
								self.assertIsNotNone(_A	)

