



import unittest

from transformers import MraConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
				import torch

				from transformers import (
				    MraForMaskedLM,
				    MraForMultipleChoice,
				    MraForQuestionAnswering,
				    MraForSequenceClassification,
				    MraForTokenClassification,
				    MraModel,
				)
				from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST

class 				snake_case__:

					"""simple docstring"""






					def __init__(							self :     Dict				,       SCREAMING_SNAKE_CASE :     int				,       SCREAMING_SNAKE_CASE :     int=2				,       SCREAMING_SNAKE_CASE :     List[str]=8				,       SCREAMING_SNAKE_CASE :     Tuple=True				,       SCREAMING_SNAKE_CASE :     Any=True				,       SCREAMING_SNAKE_CASE :     int=True				,       SCREAMING_SNAKE_CASE :     Any=True				,       SCREAMING_SNAKE_CASE :     Optional[Any]=99				,       SCREAMING_SNAKE_CASE :     Dict=16				,       SCREAMING_SNAKE_CASE :     List[str]=5				,       SCREAMING_SNAKE_CASE :     int=2				,       SCREAMING_SNAKE_CASE :     Optional[int]=36				,       SCREAMING_SNAKE_CASE :     List[Any]="gelu"				,       SCREAMING_SNAKE_CASE :     Dict=0.0				,       SCREAMING_SNAKE_CASE :     List[Any]=0.0				,       SCREAMING_SNAKE_CASE :     List[Any]=512				,       SCREAMING_SNAKE_CASE :     Union[str, Any]=16				,       SCREAMING_SNAKE_CASE :     Optional[Any]=2				,       SCREAMING_SNAKE_CASE :     Optional[int]=0.02				,       SCREAMING_SNAKE_CASE :     Any=3				,       SCREAMING_SNAKE_CASE :     Union[str, Any]=4				,       SCREAMING_SNAKE_CASE :     Optional[int]=None				,       ):
									lowercase__					:   Any													=					parent
									lowercase__					:   Union[str, Any]													=					batch_size
									lowercase__					:   Optional[int]													=					seq_length
									lowercase__					:   Optional[int]													=					is_training
									lowercase__					:   Any													=					use_input_mask
									lowercase__					:   Dict													=					use_token_type_ids
									lowercase__					:   List[Any]													=					use_labels
									lowercase__					:   Any													=					vocab_size
									lowercase__					:   int													=					hidden_size
									lowercase__					:   str													=					num_hidden_layers
									lowercase__					:   Any													=					num_attention_heads
									lowercase__					:   Optional[Any]													=					intermediate_size
									lowercase__					:   Optional[Any]													=					hidden_act
									lowercase__					:   Optional[int]													=					hidden_dropout_prob
									lowercase__					:   str													=					attention_probs_dropout_prob
									lowercase__					:   Optional[int]													=					max_position_embeddings
									lowercase__					:   Optional[int]													=					type_vocab_size
									lowercase__					:   Tuple													=					type_sequence_label_size
									lowercase__					:   Tuple													=					initializer_range
									lowercase__					:   str													=					num_labels
									lowercase__					:   List[str]													=					num_choices
									lowercase__					:   Any													=					scope







					def 			snake_case					(							self :     Optional[Any]					):
									lowercase__					:   str													=					ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size					)

									lowercase__					:   Any													=					None
									if self.use_input_mask:
													lowercase__					:   Any													=					random_attention_mask([self.batch_size, self.seq_length]					)

									lowercase__					:   int													=					None
									if self.use_token_type_ids:
													lowercase__					:   Tuple													=					ids_tensor([self.batch_size, self.seq_length]				,       self.type_vocab_size					)

									lowercase__					:   str													=					None
									lowercase__					:   Optional[int]													=					None
									lowercase__					:   Optional[Any]													=					None
									if self.use_labels:
													lowercase__					:   Any													=					ids_tensor([self.batch_size]				,       self.type_sequence_label_size					)
													lowercase__					:   int													=					ids_tensor([self.batch_size, self.seq_length]				,       self.num_labels					)
													lowercase__					:   Dict													=					ids_tensor([self.batch_size]				,       self.num_choices					)

									lowercase__					:   List[str]													=					self.get_config()

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







					def 			snake_case					(							self :     List[Any]					):
									return MraConfig(
									    vocab_size=self.vocab_size				,       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				,       max_position_embeddings=self.max_position_embeddings				,       type_vocab_size=self.type_vocab_size				,       is_decoder=SCREAMING_SNAKE_CASE				,       initializer_range=self.initializer_range				,       )







					def 			snake_case					(							self :     Optional[int]					):
									lowercase__					:   str													=					self.get_config()
									lowercase__					:   str													=					300
									return config







					def 			snake_case					(							self :     Optional[int]					):
									(
									    lowercase__							
									)					:   Optional[Any]													=					self.prepare_config_and_inputs()

									lowercase__					:   Dict													=					True
									lowercase__					:   Dict													=					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
									lowercase__					:   Any													=					ids_tensor([self.batch_size, self.seq_length]				,       vocab_size=2					)

									return (
									    config,
									    input_ids,
									    token_type_ids,
									    input_mask,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									    encoder_hidden_states,
									    encoder_attention_mask,
									)







					def 			snake_case					(							self :     str				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Any				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     int				,       SCREAMING_SNAKE_CASE :     List[str]				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     Optional[Any]					):
									lowercase__					:   Optional[int]													=					MraModel(config=SCREAMING_SNAKE_CASE					)
									model.to(SCREAMING_SNAKE_CASE					)
									model.eval()
									lowercase__					:   Union[str, Any]													=					model(SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE					)
									lowercase__					:   Optional[Any]													=					model(SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE					)
									lowercase__					:   Tuple													=					model(SCREAMING_SNAKE_CASE					)
									self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)					)







					def 			snake_case					(							self :     Optional[int]				,       SCREAMING_SNAKE_CASE :     Optional[Any]				,       SCREAMING_SNAKE_CASE :     int				,       SCREAMING_SNAKE_CASE :     Optional[int]				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Optional[Any]				,       SCREAMING_SNAKE_CASE :     Optional[Any]				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Optional[Any]				,       SCREAMING_SNAKE_CASE :     List[Any]				,       ):
									lowercase__					:   str													=					True
									lowercase__					:   Optional[int]													=					MraModel(SCREAMING_SNAKE_CASE					)
									model.to(SCREAMING_SNAKE_CASE					)
									model.eval()
									lowercase__					:   int													=					model(
									    SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE				,       encoder_hidden_states=SCREAMING_SNAKE_CASE				,       encoder_attention_mask=SCREAMING_SNAKE_CASE				,       )
									lowercase__					:   Tuple													=					model(
									    SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE				,       encoder_hidden_states=SCREAMING_SNAKE_CASE				,       )
									lowercase__					:   Tuple													=					model(SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE					)
									self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)					)







					def 			snake_case					(							self :     Union[str, Any]				,       SCREAMING_SNAKE_CASE :     str				,       SCREAMING_SNAKE_CASE :     List[str]				,       SCREAMING_SNAKE_CASE :     Any				,       SCREAMING_SNAKE_CASE :     Dict				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     List[str]				,       SCREAMING_SNAKE_CASE :     List[Any]					):
									lowercase__					:   str													=					MraForMaskedLM(config=SCREAMING_SNAKE_CASE					)
									model.to(SCREAMING_SNAKE_CASE					)
									model.eval()
									lowercase__					:   Tuple													=					model(SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE				,       labels=SCREAMING_SNAKE_CASE					)
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)					)







					def 			snake_case					(							self :     Dict				,       SCREAMING_SNAKE_CASE :     Optional[int]				,       SCREAMING_SNAKE_CASE :     Any				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Optional[int]				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Union[str, Any]				,       SCREAMING_SNAKE_CASE :     str					):
									lowercase__					:   Union[str, Any]													=					MraForQuestionAnswering(config=SCREAMING_SNAKE_CASE					)
									model.to(SCREAMING_SNAKE_CASE					)
									model.eval()
									lowercase__					:   Optional[int]													=					model(
									    SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE				,       start_positions=SCREAMING_SNAKE_CASE				,       end_positions=SCREAMING_SNAKE_CASE				,       )
									self.parent.assertEqual(result.start_logits.shape				,       (self.batch_size, self.seq_length)					)
									self.parent.assertEqual(result.end_logits.shape				,       (self.batch_size, self.seq_length)					)







					def 			snake_case					(							self :     Any				,       SCREAMING_SNAKE_CASE :     int				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Optional[int]				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     Any				,       SCREAMING_SNAKE_CASE :     str					):
									lowercase__					:   List[str]													=					self.num_labels
									lowercase__					:   Union[str, Any]													=					MraForSequenceClassification(SCREAMING_SNAKE_CASE					)
									model.to(SCREAMING_SNAKE_CASE					)
									model.eval()
									lowercase__					:   List[Any]													=					model(SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE				,       labels=SCREAMING_SNAKE_CASE					)
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)					)







					def 			snake_case					(							self :     Union[str, Any]				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     List[str]				,       SCREAMING_SNAKE_CASE :     Optional[Any]				,       SCREAMING_SNAKE_CASE :     Optional[Any]				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     Any				,       SCREAMING_SNAKE_CASE :     Any					):
									lowercase__					:   Optional[Any]													=					self.num_labels
									lowercase__					:   List[Any]													=					MraForTokenClassification(config=SCREAMING_SNAKE_CASE					)
									model.to(SCREAMING_SNAKE_CASE					)
									model.eval()
									lowercase__					:   str													=					model(SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE				,       labels=SCREAMING_SNAKE_CASE					)
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.num_labels)					)







					def 			snake_case					(							self :     List[Any]				,       SCREAMING_SNAKE_CASE :     Union[str, Any]				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     int				,       SCREAMING_SNAKE_CASE :     List[str]				,       SCREAMING_SNAKE_CASE :     str				,       SCREAMING_SNAKE_CASE :     List[str]				,       SCREAMING_SNAKE_CASE :     Optional[Any]					):
									lowercase__					:   Dict													=					self.num_choices
									lowercase__					:   List[Any]													=					MraForMultipleChoice(config=SCREAMING_SNAKE_CASE					)
									model.to(SCREAMING_SNAKE_CASE					)
									model.eval()
									lowercase__					:   int													=					input_ids.unsqueeze(1					).expand(-1				,       self.num_choices				,       -1					).contiguous()
									lowercase__					:   Any													=					token_type_ids.unsqueeze(1					).expand(-1				,       self.num_choices				,       -1					).contiguous()
									lowercase__					:   Union[str, Any]													=					input_mask.unsqueeze(1					).expand(-1				,       self.num_choices				,       -1					).contiguous()
									lowercase__					:   Optional[Any]													=					model(
									    SCREAMING_SNAKE_CASE				,       attention_mask=SCREAMING_SNAKE_CASE				,       token_type_ids=SCREAMING_SNAKE_CASE				,       labels=SCREAMING_SNAKE_CASE				,       )
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_choices)					)







					def 			snake_case					(							self :     Any					):
									lowercase__					:   Union[str, Any]													=					self.prepare_config_and_inputs()
									(
									    lowercase__							
									)					:   Dict													=					config_and_inputs
									lowercase__					:   Tuple													=					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
									return config, inputs_dict






@require_torch
class 				snake_case__(_UpperCamelCase    ,      unittest.TestCase      ):

					"""simple docstring"""






					lowercase_								=		(
					    (
					        MraModel,
					        MraForMaskedLM,
					        MraForMultipleChoice,
					        MraForQuestionAnswering,
					        MraForSequenceClassification,
					        MraForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					lowercase_								=		False
					lowercase_								=		False
					lowercase_								=		False
					lowercase_								=		False

					lowercase_								=		()







					def 			snake_case					(							self :     Tuple					):
									lowercase__					:   List[str]													=					MraModelTester(self					)
									lowercase__					:   int													=					ConfigTester(self				,       config_class=SCREAMING_SNAKE_CASE				,       hidden_size=37					)







					def 			snake_case					(							self :     Union[str, Any]					):
									self.config_tester.run_common_tests()







					def 			snake_case					(							self :     List[Any]					):
									lowercase__					:   Union[str, Any]													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     List[Any]					):
									lowercase__					:   Optional[Any]													=					self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
													lowercase__					:   List[str]													=					type
													self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     Dict					):
									lowercase__					:   List[str]													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     Dict					):
									lowercase__					:   int													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     Optional[Any]					):
									lowercase__					:   List[Any]													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     int					):
									lowercase__					:   Tuple													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     List[Any]					):
									lowercase__					:   Optional[Any]													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE					)







					@slow
					def 			snake_case					(							self :     Union[str, Any]					):
									for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowercase__					:   Union[str, Any]													=					MraModel.from_pretrained(SCREAMING_SNAKE_CASE					)
													self.assertIsNotNone(SCREAMING_SNAKE_CASE					)







					@unittest.skip(reason="MRA does not output attentions"					)
					def 			snake_case					(							self :     Dict					):
									return






@require_torch
class 				snake_case__(unittest.TestCase      ):

					"""simple docstring"""






					@slow
					def 			snake_case					(							self :     Any					):
									lowercase__					:   Optional[int]													=					MraModel.from_pretrained("uw-madison/mra-base-512-4"					)
									lowercase__					:   Union[str, Any]													=					torch.arange(256					).unsqueeze(0					)

									with torch.no_grad():
													lowercase__					:   int													=					model(SCREAMING_SNAKE_CASE					)[0]

									lowercase__					:   Dict													=					torch.Size((1, 256, 768)					)
									self.assertEqual(output.shape				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   Dict													=					torch.tensor(
									    [[[-0.0_140, 0.0_830, -0.0_381], [0.1_546, 0.1_402, 0.0_220], [0.1_162, 0.0_851, 0.0_165]]]					)

									self.assertTrue(torch.allclose(output[:, :3, :3]				,       SCREAMING_SNAKE_CASE				,       atol=1E-4					)					)







					@slow
					def 			snake_case					(							self :     Optional[Any]					):
									lowercase__					:   Optional[int]													=					MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4"					)
									lowercase__					:   Dict													=					torch.arange(256					).unsqueeze(0					)

									with torch.no_grad():
													lowercase__					:   Union[str, Any]													=					model(SCREAMING_SNAKE_CASE					)[0]

									lowercase__					:   str													=					50_265

									lowercase__					:   int													=					torch.Size((1, 256, vocab_size)					)
									self.assertEqual(output.shape				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   Dict													=					torch.tensor(
									    [[[9.2_595, -3.6_038, 11.8_819], [9.3_869, -3.2_693, 11.0_956], [11.8_524, -3.4_938, 13.1_210]]]					)

									self.assertTrue(torch.allclose(output[:, :3, :3]				,       SCREAMING_SNAKE_CASE				,       atol=1E-4					)					)







					@slow
					def 			snake_case					(							self :     List[str]					):
									lowercase__					:   Tuple													=					MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3"					)
									lowercase__					:   Union[str, Any]													=					torch.arange(4_096					).unsqueeze(0					)

									with torch.no_grad():
													lowercase__					:   List[Any]													=					model(SCREAMING_SNAKE_CASE					)[0]

									lowercase__					:   Union[str, Any]													=					50_265

									lowercase__					:   int													=					torch.Size((1, 4_096, vocab_size)					)
									self.assertEqual(output.shape				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   Tuple													=					torch.tensor(
									    [[[5.4_789, -2.3_564, 7.5_064], [7.9_067, -1.3_369, 9.9_668], [9.0_712, -1.8_106, 7.0_380]]]					)

									self.assertTrue(torch.allclose(output[:, :3, :3]				,       SCREAMING_SNAKE_CASE				,       atol=1E-4					)					)









import inspect
import unittest

import numpy as np

from transformers import BeitConfig
from transformers.testing_utils import require_flax, require_vision, slow
from transformers.utils import cached_property, is_flax_available, is_vision_available

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


if is_flax_available():
				import jax

				from transformers import FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling, FlaxBeitModel

if is_vision_available():
				from PIL import Image

				from transformers import BeitImageProcessor

class 				snake_case__(unittest.TestCase      ):

					"""simple docstring"""






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

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







					def 			snake_case					(							self :     Tuple					):
									lowercase__					:   List[Any]													=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

									lowercase__					:   Union[str, Any]													=					None
									if self.use_labels:
													lowercase__					:   List[Any]													=					ids_tensor([self.batch_size]				,       self.type_sequence_label_size					)

									lowercase__					:   int													=					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=SCREAMING_SNAKE_CASE				,       initializer_range=self.initializer_range				,       )

									return config, pixel_values, labels







					def 			snake_case					(							self :     Optional[Any]				,       SCREAMING_SNAKE_CASE :     List[str]				,       SCREAMING_SNAKE_CASE :     Dict				,       SCREAMING_SNAKE_CASE :     List[Any]					):
									lowercase__					:   Optional[Any]													=					FlaxBeitModel(config=SCREAMING_SNAKE_CASE					)
									lowercase__					:   Optional[int]													=					model(SCREAMING_SNAKE_CASE					)
									self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)					)







					def 			snake_case					(							self :     Union[str, Any]				,       SCREAMING_SNAKE_CASE :     int				,       SCREAMING_SNAKE_CASE :     Tuple				,       SCREAMING_SNAKE_CASE :     List[str]					):
									lowercase__					:   int													=					FlaxBeitForMaskedImageModeling(config=SCREAMING_SNAKE_CASE					)
									lowercase__					:   Union[str, Any]													=					model(SCREAMING_SNAKE_CASE					)
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length - 1, self.vocab_size)					)







					def 			snake_case					(							self :     Optional[int]				,       SCREAMING_SNAKE_CASE :     str				,       SCREAMING_SNAKE_CASE :     List[Any]				,       SCREAMING_SNAKE_CASE :     Any					):
									lowercase__					:   Tuple													=					self.type_sequence_label_size
									lowercase__					:   Optional[int]													=					FlaxBeitForImageClassification(config=SCREAMING_SNAKE_CASE					)
									lowercase__					:   str													=					model(SCREAMING_SNAKE_CASE					)
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.type_sequence_label_size)					)

									# test greyscale images
									lowercase__					:   int													=					1
									lowercase__					:   List[str]													=					FlaxBeitForImageClassification(SCREAMING_SNAKE_CASE					)

									lowercase__					:   int													=					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
									lowercase__					:   List[str]													=					model(SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     Any					):
									lowercase__					:   Optional[Any]													=					self.prepare_config_and_inputs()
									(
									    (
									    lowercase__
									)							,							(
									    lowercase__
									)							,							(
									    lowercase__
									)							,							
									)					:   str													=					config_and_inputs
									lowercase__					:   Dict													=					{"pixel_values": pixel_values}
									return config, inputs_dict






@require_flax
class 				snake_case__(_UpperCamelCase    ,      unittest.TestCase      ):

					"""simple docstring"""






					lowercase_								=		(
					    (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
					)







					def 			snake_case					(							self :     Any					):
									lowercase__					:   List[Any]													=					FlaxBeitModelTester(self					)
									lowercase__					:   str													=					ConfigTester(self				,       config_class=SCREAMING_SNAKE_CASE				,       has_text_modality=SCREAMING_SNAKE_CASE				,       hidden_size=37					)







					def 			snake_case					(							self :     int					):
									self.config_tester.run_common_tests()







					def 			snake_case					(							self :     int					):
									lowercase__							,							lowercase__					:   Tuple													=					self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
													lowercase__					:   List[str]													=					model_class(SCREAMING_SNAKE_CASE					)
													lowercase__					:   str													=					inspect.signature(model.__call__					)
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													lowercase__					:   Optional[int]													=					[*signature.parameters.keys()]

													lowercase__					:   str													=					["pixel_values"]
													self.assertListEqual(arg_names[:1]				,       SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     List[Any]					):
									lowercase__							,							lowercase__					:   Optional[Any]													=					self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
													with self.subTest(model_class.__name__					):
																	lowercase__					:   Optional[int]													=					self._prepare_for_class(SCREAMING_SNAKE_CASE				,       SCREAMING_SNAKE_CASE					)
																	lowercase__					:   List[str]													=					model_class(SCREAMING_SNAKE_CASE					)

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

																	with self.subTest("JIT Enabled"					):
																					lowercase__					:   Union[str, Any]													=					model_jitted(**SCREAMING_SNAKE_CASE					).to_tuple()

																	with self.subTest("JIT Disabled"					):
																					with jax.disable_jit():
																									lowercase__					:   Optional[int]													=					model_jitted(**SCREAMING_SNAKE_CASE					).to_tuple()

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







					def 			snake_case					(							self :     Optional[int]					):
									lowercase__					:   Union[str, Any]													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     int					):
									lowercase__					:   List[Any]													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE					)







					def 			snake_case					(							self :     Optional[Any]					):
									lowercase__					:   int													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE					)







					@slow
					def 			snake_case					(							self :     Optional[int]					):
									for model_class_name in self.all_model_classes:
													lowercase__					:   Any													=					model_class_name.from_pretrained("microsoft/beit-base-patch16-224"					)
													lowercase__					:   Optional[int]													=					model(np.ones((1, 3, 224, 224)					)					)
													self.assertIsNotNone(SCREAMING_SNAKE_CASE					)






def  __lowerCamelCase (       ):

				"""simple docstring"""


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

@require_vision
@require_flax
class 				snake_case__(unittest.TestCase      ):

					"""simple docstring"""






					@cached_property
					def 			snake_case					(							self :     int					):
									return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"					) if is_vision_available() else None







					@slow
					def 			snake_case					(							self :     Union[str, Any]					):
									lowercase__					:   Tuple													=					FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"					)

									lowercase__					:   int													=					self.default_image_processor
									lowercase__					:   Union[str, Any]													=					prepare_img()
									lowercase__					:   str													=					image_processor(images=SCREAMING_SNAKE_CASE				,       return_tensors="np"					).pixel_values

									# prepare bool_masked_pos
									lowercase__					:   Optional[Any]													=					np.ones((1, 196)				,       dtype=SCREAMING_SNAKE_CASE					)

									# forward pass
									lowercase__					:   Any													=					model(pixel_values=SCREAMING_SNAKE_CASE				,       bool_masked_pos=SCREAMING_SNAKE_CASE					)
									lowercase__					:   Tuple													=					outputs.logits

									# verify the logits
									lowercase__					:   List[str]													=					(1, 196, 8_192)
									self.assertEqual(logits.shape				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   Dict													=					np.array(
									    [[-3.2_437, 0.5_072, -13.9_174], [-3.2_456, 0.4_948, -13.9_401], [-3.2_033, 0.5_121, -13.8_550]]					)

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







					@slow
					def 			snake_case					(							self :     Any					):
									lowercase__					:   Union[str, Any]													=					FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"					)

									lowercase__					:   Tuple													=					self.default_image_processor
									lowercase__					:   List[Any]													=					prepare_img()
									lowercase__					:   Optional[Any]													=					image_processor(images=SCREAMING_SNAKE_CASE				,       return_tensors="np"					)

									# forward pass
									lowercase__					:   str													=					model(**SCREAMING_SNAKE_CASE					)
									lowercase__					:   Dict													=					outputs.logits

									# verify the logits
									lowercase__					:   List[str]													=					(1, 1_000)
									self.assertEqual(logits.shape				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   Tuple													=					np.array([-1.2_385, -1.0_987, -1.0_108]					)

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

									lowercase__					:   str													=					281
									self.assertEqual(logits.argmax(-1					).item()				,       SCREAMING_SNAKE_CASE					)







					@slow
					def 			snake_case					(							self :     str					):
									lowercase__					:   List[Any]													=					FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"					)

									lowercase__					:   Dict													=					self.default_image_processor
									lowercase__					:   Dict													=					prepare_img()
									lowercase__					:   List[Any]													=					image_processor(images=SCREAMING_SNAKE_CASE				,       return_tensors="np"					)

									# forward pass
									lowercase__					:   Union[str, Any]													=					model(**SCREAMING_SNAKE_CASE					)
									lowercase__					:   Optional[Any]													=					outputs.logits

									# verify the logits
									lowercase__					:   int													=					(1, 21_841)
									self.assertEqual(logits.shape				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   str													=					np.array([1.6_881, -0.2_787, 0.5_901]					)

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

									lowercase__					:   Union[str, Any]													=					2_396
									self.assertEqual(logits.argmax(-1					).item()				,       SCREAMING_SNAKE_CASE					)




