


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 						A__							(		unittest.TestCase  ):




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =parent
										_SCREAMING_SNAKE_CASE			  =vocab_size
										_SCREAMING_SNAKE_CASE			  =batch_size
										_SCREAMING_SNAKE_CASE			  =image_size
										_SCREAMING_SNAKE_CASE			  =patch_size
										_SCREAMING_SNAKE_CASE			  =num_channels
										_SCREAMING_SNAKE_CASE			  =is_training
										_SCREAMING_SNAKE_CASE			  =use_labels
										_SCREAMING_SNAKE_CASE			  =hidden_size
										_SCREAMING_SNAKE_CASE			  =num_hidden_layers
										_SCREAMING_SNAKE_CASE			  =num_attention_heads
										_SCREAMING_SNAKE_CASE			  =intermediate_size
										_SCREAMING_SNAKE_CASE			  =hidden_act
										_SCREAMING_SNAKE_CASE			  =hidden_dropout_prob
										_SCREAMING_SNAKE_CASE			  =attention_probs_dropout_prob
										_SCREAMING_SNAKE_CASE			  =type_sequence_label_size
										_SCREAMING_SNAKE_CASE			  =initializer_range

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




						def 							__UpperCamelCase				(       self					:      str  )							->     Dict:


										"""simple docstring"""





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

										_SCREAMING_SNAKE_CASE			  =None
										if self.use_labels:
														_SCREAMING_SNAKE_CASE			  =ids_tensor([self.batch_size]				,    self.type_sequence_label_size  )

										_SCREAMING_SNAKE_CASE			  =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=_a				,    initializer_range=self.initializer_range				,    )

										return config, pixel_values, labels




						def 							__UpperCamelCase				(       self					:      List[str]				,    _a					:      Tuple				,    _a					:      Union[str, Any]				,    _a					:      str  )							->     Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =FlaxBeitModel(config=_a  )
										_SCREAMING_SNAKE_CASE			  =model(_a  )
										self.parent.assertEqual(result.last_hidden_state.shape				,    (self.batch_size, self.seq_length, self.hidden_size)  )




						def 							__UpperCamelCase				(       self					:      Optional[Any]				,    _a					:      str				,    _a					:      Optional[Any]				,    _a					:      List[str]  )							->     Tuple:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =FlaxBeitForMaskedImageModeling(config=_a  )
										_SCREAMING_SNAKE_CASE			  =model(_a  )
										self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.seq_length - 1, self.vocab_size)  )




						def 							__UpperCamelCase				(       self					:      int				,    _a					:      int				,    _a					:      Union[str, Any]				,    _a					:      Dict  )							->     Tuple:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =self.type_sequence_label_size
										_SCREAMING_SNAKE_CASE			  =FlaxBeitForImageClassification(config=_a  )
										_SCREAMING_SNAKE_CASE			  =model(_a  )
										self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.type_sequence_label_size)  )

										# test greyscale images
										_SCREAMING_SNAKE_CASE			  =1
										_SCREAMING_SNAKE_CASE			  =FlaxBeitForImageClassification(_a  )

										_SCREAMING_SNAKE_CASE			  =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
										_SCREAMING_SNAKE_CASE			  =model(_a  )




						def 							__UpperCamelCase				(       self					:      int  )							->     int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =self.prepare_config_and_inputs()
										(
										    (
										    _SCREAMING_SNAKE_CASE
										)  ,				(
										    _SCREAMING_SNAKE_CASE
										)  ,				(
										    _SCREAMING_SNAKE_CASE
										)  ,				
										)			  =config_and_inputs
										_SCREAMING_SNAKE_CASE			  ={'''pixel_values''': pixel_values}
										return config, inputs_dict

@require_flax
class 						A__							(		UpperCamelCase__   ,     unittest.TestCase  ):
						UpperCAmelCase             =						(
						    (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
						)




						def 							__UpperCamelCase				(       self					:      List[Any]  )							->     None:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =FlaxBeitModelTester(self  )
										_SCREAMING_SNAKE_CASE			  =ConfigTester(self				,    config_class=_a				,    has_text_modality=_a				,    hidden_size=37  )




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


										"""simple docstring"""





										self.config_tester.run_common_tests()




						def 							__UpperCamelCase				(       self					:      Tuple  )							->     int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE			  =self.model_tester.prepare_config_and_inputs_for_common()

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

														_SCREAMING_SNAKE_CASE			  =['''pixel_values''']
														self.assertListEqual(arg_names[:1]				,    _a  )




						def 							__UpperCamelCase				(       self					:      Tuple  )							->     Tuple:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE			  =self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__  ):
																		_SCREAMING_SNAKE_CASE			  =self._prepare_for_class(_a				,    _a  )
																		_SCREAMING_SNAKE_CASE			  =model_class(_a  )

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

																		with self.subTest('''JIT Enabled'''  ):
																						_SCREAMING_SNAKE_CASE			  =model_jitted(**_a  ).to_tuple()

																		with self.subTest('''JIT Disabled'''  ):
																						with jax.disable_jit():
																										_SCREAMING_SNAKE_CASE			  =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  )




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_a  )




						def 							__UpperCamelCase				(       self					:      int  )							->     List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*_a  )




						def 							__UpperCamelCase				(       self					:      Tuple  )							->     Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*_a  )




						@slow
						def 							__UpperCamelCase				(       self					:      List[Any]  )							->     Union[str, Any]:


										"""simple docstring"""





										for model_class_name in self.all_model_classes:
														_SCREAMING_SNAKE_CASE			  =model_class_name.from_pretrained('''microsoft/beit-base-patch16-224'''  )
														_SCREAMING_SNAKE_CASE			  =model(np.ones((1, 3, 224, 224)  )  )
														self.assertIsNotNone(_a  )







def 						lowerCamelCase(   ):
				_SCREAMING_SNAKE_CASE			  =Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
				return image


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




						@cached_property
						def 							__UpperCamelCase				(       self					:      Optional[Any]  )							->     List[str]:


										"""simple docstring"""





										return BeitImageProcessor.from_pretrained('''microsoft/beit-base-patch16-224'''  ) if is_vision_available() else None




						@slow
						def 							__UpperCamelCase				(       self					:      Union[str, Any]  )							->     List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =FlaxBeitForMaskedImageModeling.from_pretrained('''microsoft/beit-base-patch16-224-pt22k'''  )

										_SCREAMING_SNAKE_CASE			  =self.default_image_processor
										_SCREAMING_SNAKE_CASE			  =prepare_img()
										_SCREAMING_SNAKE_CASE			  =image_processor(images=_a				,    return_tensors='''np'''  ).pixel_values

										# prepare bool_masked_pos
										_SCREAMING_SNAKE_CASE			  =np.ones((1, 196)				,    dtype=_a  )

										# forward pass
										_SCREAMING_SNAKE_CASE			  =model(pixel_values=_a				,    bool_masked_pos=_a  )
										_SCREAMING_SNAKE_CASE			  =outputs.logits

										# verify the logits
										_SCREAMING_SNAKE_CASE			  =(1, 196, 8192)
										self.assertEqual(logits.shape				,    _a  )

										_SCREAMING_SNAKE_CASE			  =np.array(
										    [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]  )

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




						@slow
						def 							__UpperCamelCase				(       self					:      List[Any]  )							->     Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =FlaxBeitForImageClassification.from_pretrained('''microsoft/beit-base-patch16-224'''  )

										_SCREAMING_SNAKE_CASE			  =self.default_image_processor
										_SCREAMING_SNAKE_CASE			  =prepare_img()
										_SCREAMING_SNAKE_CASE			  =image_processor(images=_a				,    return_tensors='''np'''  )

										# forward pass
										_SCREAMING_SNAKE_CASE			  =model(**_a  )
										_SCREAMING_SNAKE_CASE			  =outputs.logits

										# verify the logits
										_SCREAMING_SNAKE_CASE			  =(1, 1000)
										self.assertEqual(logits.shape				,    _a  )

										_SCREAMING_SNAKE_CASE			  =np.array([-1.23_85, -1.09_87, -1.01_08]  )

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

										_SCREAMING_SNAKE_CASE			  =281
										self.assertEqual(logits.argmax(-1  ).item()				,    _a  )




						@slow
						def 							__UpperCamelCase				(       self					:      str  )							->     Dict:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =FlaxBeitForImageClassification.from_pretrained('''microsoft/beit-large-patch16-224-pt22k-ft22k'''  )

										_SCREAMING_SNAKE_CASE			  =self.default_image_processor
										_SCREAMING_SNAKE_CASE			  =prepare_img()
										_SCREAMING_SNAKE_CASE			  =image_processor(images=_a				,    return_tensors='''np'''  )

										# forward pass
										_SCREAMING_SNAKE_CASE			  =model(**_a  )
										_SCREAMING_SNAKE_CASE			  =outputs.logits

										# verify the logits
										_SCREAMING_SNAKE_CASE			  =(1, 2_1841)
										self.assertEqual(logits.shape				,    _a  )

										_SCREAMING_SNAKE_CASE			  =np.array([1.68_81, -0.27_87, 0.59_01]  )

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

										_SCREAMING_SNAKE_CASE			  =2396
										self.assertEqual(logits.argmax(-1  ).item()				,    _a  )



import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment







def 						lowerCamelCase(   a__	,a__=()	,a__=None	,a__="no"	,a__="29500"):
				_SCREAMING_SNAKE_CASE			  =False
				_SCREAMING_SNAKE_CASE			  =False
				if any(key.startswith('''KAGGLE''') for key in os.environ.keys()):
								_SCREAMING_SNAKE_CASE			  =True
				elif "IPython" in sys.modules:
								_SCREAMING_SNAKE_CASE			  ='''google.colab''' in str(sys.modules['''IPython'''].get_ipython())

				try:
								_SCREAMING_SNAKE_CASE			  =PrecisionType(mixed_precision.lower())
				except ValueError:
								raise ValueError(
								    f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.")

				if (in_colab or in_kaggle) and (os.environ.get('''TPU_NAME'''	,a__) is not None):
								# TPU launch
								import torch_xla.distributed.xla_multiprocessing as xmp

								if len(AcceleratorState._shared_state) > 0:
												raise ValueError(
												    '''To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside '''
												    '''your training function. Restart your notebook and make sure no cells initializes an '''
												    '''`Accelerator`.''')
								if num_processes is None:
												_SCREAMING_SNAKE_CASE			  =8

								_SCREAMING_SNAKE_CASE			  =PrepareForLaunch(a__	,distributed_type='''TPU''')
								print(f"Launching a training on {num_processes} TPU cores.")
								xmp.spawn(a__	,args=a__	,nprocs=a__	,start_method='''fork''')
				elif in_colab:
								# No need for a distributed launch otherwise as it's either CPU or one GPU.
								if torch.cuda.is_available():
												print('''Launching training on one GPU.''')
								else:
												print('''Launching training on one CPU.''')
								function(*a__)
				else:
								if num_processes is None:
												raise ValueError(
												    '''You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call.''')

								if num_processes > 1:
												# Multi-GPU launch
												from torch.multiprocessing import start_processes
												from torch.multiprocessing.spawn import ProcessRaisedException

												if len(AcceleratorState._shared_state) > 0:
																raise ValueError(
																    '''To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized '''
																    '''inside your training function. Restart your notebook and make sure no cells initializes an '''
																    '''`Accelerator`.''')

												if torch.cuda.is_initialized():
																raise ValueError(
																    '''To launch a multi-GPU training from your notebook, you need to avoid running any instruction '''
																    '''using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA '''
																    '''function.''')

												# torch.distributed will expect a few environment variable to be here. We set the ones common to each
												# process here (the other ones will be set be the launcher).
												with patch_environment(
												    world_size=a__	,master_addr='''127.0.01'''	,master_port=a__	,mixed_precision=a__):
																_SCREAMING_SNAKE_CASE			  =PrepareForLaunch(a__	,distributed_type='''MULTI_GPU''')
																print(f"Launching training on {num_processes} GPUs.")
																try:
																				start_processes(a__	,args=a__	,nprocs=a__	,start_method='''fork''')
																except ProcessRaisedException as e:
																				if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
																								raise RuntimeError(
																								    '''CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. '''
																								    '''This likely stems from an outside import causing issues once the `notebook_launcher()` is called. '''
																								    '''Please review your imports and test them when running the `notebook_launcher()` to identify '''
																								    '''which one is problematic.''') from e

								else:
												# No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
												if is_mps_available():
																_SCREAMING_SNAKE_CASE			  ='''1'''
																print('''Launching training on MPS.''')
												elif torch.cuda.is_available():
																print('''Launching training on one GPU.''')
												else:
																print('''Launching training on CPU.''')
												function(*a__)







def 						lowerCamelCase(   a__	,a__=()	,a__=2):
				from torch.multiprocessing import start_processes

				with tempfile.NamedTemporaryFile() as tmp_file:
								# torch.distributed will expect a few environment variable to be here. We set the ones common to each
								# process here (the other ones will be set be the launcher).
								with patch_environment(
								    world_size=a__	,master_addr='''127.0.01'''	,master_port='''29500'''	,accelerate_mixed_precision='''no'''	,accelerate_debug_rdv_file=tmp_file.name	,accelerate_use_cpu='''yes'''	,):
												_SCREAMING_SNAKE_CASE			  =PrepareForLaunch(a__	,debug=a__)
												start_processes(a__	,args=a__	,nprocs=a__	,start_method='''fork''')