


'''simple docstring'''



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						):




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

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




					def 					snake_case__    (   self):
								UpperCAmelCase__	:       str   =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

								UpperCAmelCase__	:       Dict   =  None
								if self.use_labels:
											UpperCAmelCase__	:       List[Any]   =  ids_tensor([self.batch_size]				,				self.type_sequence_label_size)

								UpperCAmelCase__	:       Any   =  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=__UpperCAmelCase				,				initializer_range=self.initializer_range				,				)

								return config, pixel_values, labels




					def 					snake_case__    (   self				,				_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase):
								UpperCAmelCase__	:       Optional[Any]   =  FlaxBeitModel(config=__UpperCAmelCase)
								UpperCAmelCase__	:       Dict   =  model(__UpperCAmelCase)
								self.parent.assertEqual(result.last_hidden_state.shape				,				(self.batch_size, self.seq_length, self.hidden_size))




					def 					snake_case__    (   self				,				_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase):
								UpperCAmelCase__	:       List[str]   =  FlaxBeitForMaskedImageModeling(config=__UpperCAmelCase)
								UpperCAmelCase__	:       Union[str, Any]   =  model(__UpperCAmelCase)
								self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length - 1, self.vocab_size))




					def 					snake_case__    (   self				,				_lowerCamelCase				,				_lowerCamelCase				,				_lowerCamelCase):
								UpperCAmelCase__	:       Optional[int]   =  self.type_sequence_label_size
								UpperCAmelCase__	:       List[str]   =  FlaxBeitForImageClassification(config=__UpperCAmelCase)
								UpperCAmelCase__	:       int   =  model(__UpperCAmelCase)
								self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.type_sequence_label_size))

								# test greyscale images
								UpperCAmelCase__	:       Tuple   =  1
								UpperCAmelCase__	:       Union[str, Any]   =  FlaxBeitForImageClassification(__UpperCAmelCase)

								UpperCAmelCase__	:       Tuple   =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
								UpperCAmelCase__	:       int   =  model(__UpperCAmelCase)




					def 					snake_case__    (   self):
								UpperCAmelCase__	:       Optional[Any]   =  self.prepare_config_and_inputs()
								(
								    UpperCAmelCase__							
								)	:       str   =  config_and_inputs
								UpperCAmelCase__	:       Optional[Any]   =  {'pixel_values': pixel_values}
								return config, inputs_dict


@require_flax
class 		_snake_case       (  a__							, unittest.TestCase						):
					lowerCAmelCase   :Optional[int]				=       (
					    (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
					)




					def 					snake_case__    (   self):
								UpperCAmelCase__	:       Optional[Any]   =  FlaxBeitModelTester(self)
								UpperCAmelCase__	:       Optional[Any]   =  ConfigTester(self				,				config_class=__UpperCAmelCase				,				has_text_modality=__UpperCAmelCase				,				hidden_size=37)




					def 					snake_case__    (   self):
								self.config_tester.run_common_tests()




					def 					snake_case__    (   self):
								UpperCAmelCase__	:       Union[str, Any]   =  self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											UpperCAmelCase__	:       Dict   =  model_class(__UpperCAmelCase)
											UpperCAmelCase__	:       Union[str, Any]   =  inspect.signature(model.__call__)
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											UpperCAmelCase__	:       int   =  [*signature.parameters.keys()]

											UpperCAmelCase__	:       List[Any]   =  ['pixel_values']
											self.assertListEqual(arg_names[:1]				,				__UpperCAmelCase)




					def 					snake_case__    (   self):
								UpperCAmelCase__	:       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__):
														UpperCAmelCase__	:       List[str]   =  self._prepare_for_class(__UpperCAmelCase				,				__UpperCAmelCase)
														UpperCAmelCase__	:       Dict   =  model_class(__UpperCAmelCase)

														@jax.jit
														def model_jitted(_lowerCamelCase				,				**_lowerCamelCase):
																	return model(pixel_values=__UpperCAmelCase				,				**__UpperCAmelCase)

														with self.subTest("""JIT Enabled"""):
																	UpperCAmelCase__	:       str   =  model_jitted(**__UpperCAmelCase).to_tuple()

														with self.subTest("""JIT Disabled"""):
																	with jax.disable_jit():
																				UpperCAmelCase__	:       List[Any]   =  model_jitted(**__UpperCAmelCase).to_tuple()

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




					def 					snake_case__    (   self):
								UpperCAmelCase__	:       Optional[Any]   =  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__UpperCAmelCase)




					def 					snake_case__    (   self):
								UpperCAmelCase__	:       int   =  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase)




					def 					snake_case__    (   self):
								UpperCAmelCase__	:       Optional[int]   =  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase)




					@slow
					def 					snake_case__    (   self):

								for model_class_name in self.all_model_classes:
											UpperCAmelCase__	:       str   =  model_class_name.from_pretrained("""microsoft/beit-base-patch16-224""")
											UpperCAmelCase__	:       Any   =  model(np.ones((1, 3, 224, 224)))
											self.assertIsNotNone(__UpperCAmelCase)







def      _UpperCamelCase   (		):
			UpperCAmelCase__	:       List[str]   =  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""       )
			return image


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




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




					@slow
					def 					snake_case__    (   self):
								UpperCAmelCase__	:       Any   =  FlaxBeitForMaskedImageModeling.from_pretrained("""microsoft/beit-base-patch16-224-pt22k""")

								UpperCAmelCase__	:       Optional[int]   =  self.default_image_processor
								UpperCAmelCase__	:       List[str]   =  prepare_img()
								UpperCAmelCase__	:       Any   =  image_processor(images=__UpperCAmelCase				,				return_tensors="""np""").pixel_values

								# prepare bool_masked_pos
								UpperCAmelCase__	:       str   =  np.ones((1, 196)				,				dtype=__UpperCAmelCase)

								# forward pass
								UpperCAmelCase__	:       int   =  model(pixel_values=__UpperCAmelCase				,				bool_masked_pos=__UpperCAmelCase)
								UpperCAmelCase__	:       Dict   =  outputs.logits

								# verify the logits
								UpperCAmelCase__	:       Optional[Any]   =  (1, 196, 8192)
								self.assertEqual(logits.shape				,				__UpperCAmelCase)

								UpperCAmelCase__	:       List[Any]   =  np.array(
								    [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]])

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




					@slow
					def 					snake_case__    (   self):
								UpperCAmelCase__	:       str   =  FlaxBeitForImageClassification.from_pretrained("""microsoft/beit-base-patch16-224""")

								UpperCAmelCase__	:       int   =  self.default_image_processor
								UpperCAmelCase__	:       Any   =  prepare_img()
								UpperCAmelCase__	:       Optional[Any]   =  image_processor(images=__UpperCAmelCase				,				return_tensors="""np""")

								# forward pass
								UpperCAmelCase__	:       int   =  model(**__UpperCAmelCase)
								UpperCAmelCase__	:       Tuple   =  outputs.logits

								# verify the logits
								UpperCAmelCase__	:       str   =  (1, 1000)
								self.assertEqual(logits.shape				,				__UpperCAmelCase)

								UpperCAmelCase__	:       Tuple   =  np.array([-1.2385, -1.0987, -1.0108])

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

								UpperCAmelCase__	:       Optional[Any]   =  281
								self.assertEqual(logits.argmax(-1).item()				,				__UpperCAmelCase)




					@slow
					def 					snake_case__    (   self):
								UpperCAmelCase__	:       str   =  FlaxBeitForImageClassification.from_pretrained("""microsoft/beit-large-patch16-224-pt22k-ft22k""")

								UpperCAmelCase__	:       Dict   =  self.default_image_processor
								UpperCAmelCase__	:       Optional[Any]   =  prepare_img()
								UpperCAmelCase__	:       List[str]   =  image_processor(images=__UpperCAmelCase				,				return_tensors="""np""")

								# forward pass
								UpperCAmelCase__	:       Dict   =  model(**__UpperCAmelCase)
								UpperCAmelCase__	:       Any   =  outputs.logits

								# verify the logits
								UpperCAmelCase__	:       Tuple   =  (1, 2_1841)
								self.assertEqual(logits.shape				,				__UpperCAmelCase)

								UpperCAmelCase__	:       Any   =  np.array([1.6881, -0.2787, 0.5901])

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

								UpperCAmelCase__	:       int   =  2396
								self.assertEqual(logits.argmax(-1).item()				,				__UpperCAmelCase)






"""simple docstring"""





from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin






@dataclass
class    _lowerCAmelCase  (	a			):
       """simple docstring"""




       __magic_name__	:torch.FloatTensor






class    _lowerCAmelCase  (	a    ,  a			):
       """simple docstring"""




       @register_to_config
       def __init__(       self       ,					__UpperCAmelCase = 1_6       ,					__UpperCAmelCase = 8_8       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = 1       ,					__UpperCAmelCase = 0.0       ,					__UpperCAmelCase = 3_2       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = False       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = "geglu"       ,					__UpperCAmelCase = True       ,					__UpperCAmelCase = True       ,					):



         '''simple docstring'''

         super().__init__()
         lowerCAmelCase__     :Dict											=				num_attention_heads
         lowerCAmelCase__     :Any											=				attention_head_dim
         lowerCAmelCase__     :Optional[int]											=				num_attention_heads * attention_head_dim

         lowerCAmelCase__     :Any											=				in_channels

         lowerCAmelCase__     :str											=				torch.nn.GroupNorm(num_groups=__UpperCAmelCase       ,					num_channels=__UpperCAmelCase       ,					eps=1E-6       ,					affine=__UpperCAmelCase )
         lowerCAmelCase__     :int											=				nn.Linear(__UpperCAmelCase       ,					__UpperCAmelCase )

         # 3. Define transformers blocks
         lowerCAmelCase__     :List[Any]											=				nn.ModuleList(
             [
                 BasicTransformerBlock(
                     __UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					dropout=__UpperCAmelCase       ,					cross_attention_dim=__UpperCAmelCase       ,					activation_fn=__UpperCAmelCase       ,					attention_bias=__UpperCAmelCase       ,					double_self_attention=__UpperCAmelCase       ,					norm_elementwise_affine=__UpperCAmelCase       ,					)
                 for d in range(__UpperCAmelCase )
             ] )

         lowerCAmelCase__     :List[Any]											=				nn.Linear(__UpperCAmelCase       ,					__UpperCAmelCase )






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase=1       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase = True       ,					):



         '''simple docstring'''

         lowerCAmelCase__				,							lowerCAmelCase__				,							lowerCAmelCase__				,							lowerCAmelCase__     :int											=				hidden_states.shape
         lowerCAmelCase__     :Tuple											=				batch_frames // num_frames

         lowerCAmelCase__     :str											=				hidden_states

         lowerCAmelCase__     :Union[str, Any]											=				hidden_states[None, :].reshape(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase )
         lowerCAmelCase__     :str											=				hidden_states.permute(0       ,					2       ,					1       ,					3       ,					4 )

         lowerCAmelCase__     :Optional[int]											=				self.norm(__UpperCAmelCase )
         lowerCAmelCase__     :Optional[Any]											=				hidden_states.permute(0       ,					3       ,					4       ,					2       ,					1 ).reshape(batch_size * height * width       ,					__UpperCAmelCase       ,					__UpperCAmelCase )

         lowerCAmelCase__     :List[Any]											=				self.proj_in(__UpperCAmelCase )

         # 2. Blocks
         for block in self.transformer_blocks:
           lowerCAmelCase__     :Optional[int]											=				block(
               __UpperCAmelCase       ,					encoder_hidden_states=__UpperCAmelCase       ,					timestep=__UpperCAmelCase       ,					cross_attention_kwargs=__UpperCAmelCase       ,					class_labels=__UpperCAmelCase       ,					)

         # 3. Output
         lowerCAmelCase__     :Any											=				self.proj_out(__UpperCAmelCase )
         lowerCAmelCase__     :Dict											=				(
             hidden_states[None, None, :]
             .reshape(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase )
             .permute(0       ,					3       ,					4       ,					1       ,					2 )
             .contiguous()
         )
         lowerCAmelCase__     :Optional[Any]											=				hidden_states.reshape(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase )

         lowerCAmelCase__     :List[Any]											=				hidden_states + residual

         if not return_dict:
           return (output,)

         return TransformerTemporalModelOutput(sample=__UpperCAmelCase )
