

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



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





           '''simple docstring'''
           A_						=      parent
           A_						=      vocab_size
           A_						=      batch_size
           A_						=      image_size
           A_						=      patch_size
           A_						=      num_channels
           A_						=      is_training
           A_						=      use_labels
           A_						=      hidden_size
           A_						=      num_hidden_layers
           A_						=      num_attention_heads
           A_						=      intermediate_size
           A_						=      hidden_act
           A_						=      hidden_dropout_prob
           A_						=      attention_probs_dropout_prob
           A_						=      type_sequence_label_size
           A_						=      initializer_range

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



    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_						=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

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

           A_						=      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		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      Dict:





           '''simple docstring'''
           A_						=      FlaxBeitModel(config=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__       )
           self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.hidden_size)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      str:





           '''simple docstring'''
           A_						=      FlaxBeitForMaskedImageModeling(config=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__       )
           self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length - 1, self.vocab_size)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      List[Any]:





           '''simple docstring'''
           A_						=      self.type_sequence_label_size
           A_						=      FlaxBeitForImageClassification(config=UpperCamelCase__       )
           A_						=      model(UpperCamelCase__       )
           self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.type_sequence_label_size)       )

           # test greyscale images
           A_						=      1
           A_						=      FlaxBeitForImageClassification(UpperCamelCase__       )

           A_						=      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
           A_						=      model(UpperCamelCase__       )







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





           '''simple docstring'''
           A_						=      self.prepare_config_and_inputs()
           (
               (
               A_
           )					,			(
               A_
           )					,			(
               A_
           )					,			
           )						=      config_and_inputs
           A_						=      {"""pixel_values""": pixel_values}
           return config, inputs_dict


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



    def 					snake_case_			(					self       )						->      None:





           '''simple docstring'''
           A_						=      FlaxBeitModelTester(self       )
           A_						=      ConfigTester(self		,				config_class=UpperCamelCase__		,				has_text_modality=UpperCamelCase__		,				hidden_size=37       )



    def 					snake_case_			(					self       )						->      Any:





           '''simple docstring'''
           self.config_tester.run_common_tests()



    def 					snake_case_			(					self       )						->      Dict:





           '''simple docstring'''
           A_					,			A_						=      self.model_tester.prepare_config_and_inputs_for_common()

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

                  A_						=      ["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]		,				UpperCamelCase__       )



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





           '''simple docstring'''
           A_					,			A_						=      self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  with self.subTest(model_class.__name__       ):
                         A_						=      self._prepare_for_class(UpperCamelCase__		,				UpperCamelCase__       )
                         A_						=      model_class(UpperCamelCase__       )

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

                         with self.subTest("""JIT Enabled"""       ):
                                A_						=      model_jitted(**UpperCamelCase__       ).to_tuple()

                         with self.subTest("""JIT Disabled"""       ):
                                with jax.disable_jit():
                                       A_						=      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       )						->      List[str]:





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__       )



    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__       )



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





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__       )







    @slow
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           for model_class_name in self.all_model_classes:
                  A_						=      model_class_name.from_pretrained("""microsoft/beit-base-patch16-224"""       )
                  A_						=      model(np.ones((1, 3, 224, 224)       )       )
                  self.assertIsNotNone(UpperCamelCase__       )

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







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



    @cached_property
    def 					snake_case_			(					self       )						->      int:





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



    @slow
    def 					snake_case_			(					self       )						->      Dict:





           '''simple docstring'''
           A_						=      FlaxBeitForMaskedImageModeling.from_pretrained("""microsoft/beit-base-patch16-224-pt22k"""       )

           A_						=      self.default_image_processor
           A_						=      prepare_img()
           A_						=      image_processor(images=UpperCamelCase__		,				return_tensors="""np"""       ).pixel_values

           # prepare bool_masked_pos
           A_						=      np.ones((1, 196)		,				dtype=UpperCamelCase__       )

           # forward pass
           A_						=      model(pixel_values=UpperCamelCase__		,				bool_masked_pos=UpperCamelCase__       )
           A_						=      outputs.logits

           # verify the logits
           A_						=      (1, 196, 8192)
           self.assertEqual(logits.shape		,				UpperCamelCase__       )

           A_						=      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       )						->      Dict:





           '''simple docstring'''
           A_						=      FlaxBeitForImageClassification.from_pretrained("""microsoft/beit-base-patch16-224"""       )

           A_						=      self.default_image_processor
           A_						=      prepare_img()
           A_						=      image_processor(images=UpperCamelCase__		,				return_tensors="""np"""       )

           # forward pass
           A_						=      model(**UpperCamelCase__       )
           A_						=      outputs.logits

           # verify the logits
           A_						=      (1, 1000)
           self.assertEqual(logits.shape		,				UpperCamelCase__       )

           A_						=      np.array([-1.2385, -1.0987, -1.0108]       )

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

           A_						=      281
           self.assertEqual(logits.argmax(-1       ).item()		,				UpperCamelCase__       )







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





           '''simple docstring'''
           A_						=      FlaxBeitForImageClassification.from_pretrained("""microsoft/beit-large-patch16-224-pt22k-ft22k"""       )

           A_						=      self.default_image_processor
           A_						=      prepare_img()
           A_						=      image_processor(images=UpperCamelCase__		,				return_tensors="""np"""       )

           # forward pass
           A_						=      model(**UpperCamelCase__       )
           A_						=      outputs.logits

           # verify the logits
           A_						=      (1, 21841)
           self.assertEqual(logits.shape		,				UpperCamelCase__       )

           A_						=      np.array([1.6881, -0.2787, 0.5901]       )

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

           A_						=      2396
           self.assertEqual(logits.argmax(-1       ).item()		,				UpperCamelCase__       )



'''simple docstring'''







import tensorflow as tf

from ...tf_utils import shape_list


class 					A__      (      tf.keras.layers.Layer			):



    def __init__(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=1		,				UpperCamelCase__=False		,				**UpperCamelCase__       )						->      Union[str, Any]:





           '''simple docstring'''
           super().__init__(**UpperCamelCase__       )

           A_						=      vocab_size
           A_						=      d_embed
           A_						=      d_proj

           A_						=      cutoffs + [vocab_size]
           A_						=      [0] + self.cutoffs
           A_						=      div_val

           A_						=      self.cutoffs[0]
           A_						=      len(self.cutoffs       ) - 1
           A_						=      self.shortlist_size + self.n_clusters
           A_						=      keep_order

           A_						=      []
           A_						=      []



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





           '''simple docstring'''
           if self.n_clusters > 0:
                  A_						=      self.add_weight(
                      shape=(self.n_clusters, self.d_embed)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name="""cluster_weight"""       )
                  A_						=      self.add_weight(
                      shape=(self.n_clusters,)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name="""cluster_bias"""       )

           if self.div_val == 1:
                  for i in range(len(self.cutoffs       )       ):
                         if self.d_proj != self.d_embed:
                                A_						=      self.add_weight(
                                    shape=(self.d_embed, self.d_proj)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_projs_._{i}'''		,				)
                                self.out_projs.append(UpperCamelCase__       )
                         else:
                                self.out_projs.append(UpperCamelCase__       )
                         A_						=      self.add_weight(
                             shape=(self.vocab_size, self.d_embed)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_layers_._{i}_._weight'''		,				)
                         A_						=      self.add_weight(
                             shape=(self.vocab_size,)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_layers_._{i}_._bias'''		,				)
                         self.out_layers.append((weight, bias)       )
           else:
                  for i in range(len(self.cutoffs       )       ):
                         A_					,			A_						=      self.cutoff_ends[i], self.cutoff_ends[i + 1]
                         A_						=      self.d_embed // (self.div_val**i)

                         A_						=      self.add_weight(
                             shape=(d_emb_i, self.d_proj)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_projs_._{i}'''       )
                         self.out_projs.append(UpperCamelCase__       )
                         A_						=      self.add_weight(
                             shape=(r_idx - l_idx, d_emb_i)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_layers_._{i}_._weight'''		,				)
                         A_						=      self.add_weight(
                             shape=(r_idx - l_idx,)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name=f'''out_layers_._{i}_._bias'''		,				)
                         self.out_layers.append((weight, bias)       )
           super().build(UpperCamelCase__       )



    @staticmethod
    def 					snake_case_			(					UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=None       )						->      List[Any]:





           '''simple docstring'''
           A_						=      x
           if proj is not None:
                  A_						=      tf.einsum("""ibd,ed->ibe"""		,				UpperCamelCase__		,				UpperCamelCase__       )
           return tf.einsum("""ibd,nd->ibn"""		,				UpperCamelCase__		,				UpperCamelCase__       ) + b



    @staticmethod
    def 					snake_case_			(					UpperCamelCase__		,				UpperCamelCase__       )						->      List[str]:





           '''simple docstring'''
           A_						=      shape_list(UpperCamelCase__       )
           A_						=      tf.range(lp_size[0]		,				dtype=target.dtype       )
           A_						=      tf.stack([r, target]		,				1       )
           return tf.gather_nd(UpperCamelCase__		,				UpperCamelCase__       )







    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=True		,				UpperCamelCase__=False       )						->      Optional[int]:





           '''simple docstring'''
           A_						=      0
           if self.n_clusters == 0:
                  A_						=      self._logit(UpperCamelCase__		,				self.out_layers[0][0]		,				self.out_layers[0][1]		,				self.out_projs[0]       )
                  if target is not None:
                         A_						=      tf.nn.sparse_softmax_cross_entropy_with_logits(labels=UpperCamelCase__		,				logits=UpperCamelCase__       )
                  A_						=      tf.nn.log_softmax(UpperCamelCase__		,				axis=-1       )
           else:
                  A_						=      shape_list(UpperCamelCase__       )
                  A_						=      []
                  A_						=      tf.zeros(hidden_sizes[:2]       )
                  for i in range(len(self.cutoffs       )       ):
                         A_					,			A_						=      self.cutoff_ends[i], self.cutoff_ends[i + 1]
                         if target is not None:
                                A_						=      (target >= l_idx) & (target < r_idx)
                                A_						=      tf.where(UpperCamelCase__       )
                                A_						=      tf.boolean_mask(UpperCamelCase__		,				UpperCamelCase__       ) - l_idx

                         if self.div_val == 1:
                                A_						=      self.out_layers[0][0][l_idx:r_idx]
                                A_						=      self.out_layers[0][1][l_idx:r_idx]
                         else:
                                A_						=      self.out_layers[i][0]
                                A_						=      self.out_layers[i][1]

                         if i == 0:
                                A_						=      tf.concat([cur_W, self.cluster_weight]		,				0       )
                                A_						=      tf.concat([cur_b, self.cluster_bias]		,				0       )

                                A_						=      self._logit(UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				self.out_projs[0]       )
                                A_						=      tf.nn.log_softmax(UpperCamelCase__       )
                                out.append(head_logprob[..., : self.cutoffs[0]]       )
                                if target is not None:
                                       A_						=      tf.boolean_mask(UpperCamelCase__		,				UpperCamelCase__       )
                                       A_						=      self._gather_logprob(UpperCamelCase__		,				UpperCamelCase__       )
                         else:
                                A_						=      self._logit(UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				self.out_projs[i]       )
                                A_						=      tf.nn.log_softmax(UpperCamelCase__       )
                                A_						=      self.cutoffs[0] + i - 1  # No probability for the head cluster
                                A_						=      head_logprob[..., cluster_prob_idx, None] + tail_logprob
                                out.append(UpperCamelCase__       )
                                if target is not None:
                                       A_						=      tf.boolean_mask(UpperCamelCase__		,				UpperCamelCase__       )
                                       A_						=      tf.boolean_mask(UpperCamelCase__		,				UpperCamelCase__       )
                                       A_						=      self._gather_logprob(UpperCamelCase__		,				UpperCamelCase__       )
                                       cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
                         if target is not None:
                                loss += tf.scatter_nd(UpperCamelCase__		,				-cur_logprob		,				shape_list(UpperCamelCase__       )       )
                  A_						=      tf.concat(UpperCamelCase__		,				axis=-1       )

           if target is not None:
                  if return_mean:
                         A_						=      tf.reduce_mean(UpperCamelCase__       )
                  # Add the training-time loss value to the layer using `self.add_loss()`.
                  self.add_loss(UpperCamelCase__       )

                  # Log the loss as a metric (we could log arbitrary metrics,
                  # including different metrics for training and inference.
                  self.add_metric(UpperCamelCase__		,				name=self.name		,				aggregation="""mean""" if return_mean else """"""       )

           return out
