




import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
     import jax
     import jax.numpy as jnp

     from transformers.generation import (
         FlaxForcedBOSTokenLogitsProcessor,
         FlaxForcedEOSTokenLogitsProcessor,
         FlaxLogitsProcessorList,
         FlaxMinLengthLogitsProcessor,
         FlaxTemperatureLogitsWarper,
         FlaxTopKLogitsWarper,
         FlaxTopPLogitsWarper,
     )

@require_flax
class   __a(					unittest.TestCase       ):






    """simple docstring"""





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )			->		int:
         UpperCAmelCase_      :							Optional[Any]       =      jnp.ones((batch_size, length)  ) / length
         return scores





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							Optional[Any]       =      None
         UpperCAmelCase_      :							str       =      20

         UpperCAmelCase_      :							List[str]       =      self._get_uniform_logits(batch_size=2       ,length=_SCREAMING_SNAKE_CASE  )

         # tweak scores to not be uniform anymore
         UpperCAmelCase_      :							Dict       =      scores.at[1, 5].set((1 / length) + 0.1  )  # peak, 1st batch
         UpperCAmelCase_      :							Tuple       =      scores.at[1, 10].set((1 / length) - 0.4  )  # valley, 1st batch

         # compute softmax
         UpperCAmelCase_      :							str       =      jax.nn.softmax(_SCREAMING_SNAKE_CASE       ,axis=-1  )

         UpperCAmelCase_      :							Tuple       =      FlaxTemperatureLogitsWarper(temperature=0.5  )
         UpperCAmelCase_      :							Optional[Any]       =      FlaxTemperatureLogitsWarper(temperature=1.3  )

         UpperCAmelCase_      :							Union[str, Any]       =      jax.nn.softmax(temp_dist_warper_sharper(_SCREAMING_SNAKE_CASE       ,scores.copy()       ,cur_len=_SCREAMING_SNAKE_CASE  )       ,axis=-1  )
         UpperCAmelCase_      :							List[Any]       =      jax.nn.softmax(temp_dist_warper_smoother(_SCREAMING_SNAKE_CASE       ,scores.copy()       ,cur_len=_SCREAMING_SNAKE_CASE  )       ,axis=-1  )

         # uniform distribution stays uniform
         self.assertTrue(jnp.allclose(probs[0, :]       ,warped_prob_sharp[0, :]       ,atol=1e-3  )  )
         self.assertTrue(jnp.allclose(probs[0, :]       ,warped_prob_smooth[0, :]       ,atol=1e-3  )  )

         # sharp peaks get higher, valleys get lower
         self.assertLess(probs[1, :].max()       ,warped_prob_sharp[1, :].max()  )
         self.assertGreater(probs[1, :].min()       ,warped_prob_sharp[1, :].min()  )

         # smooth peaks get lower, valleys get higher
         self.assertGreater(probs[1, :].max()       ,warped_prob_smooth[1, :].max()  )
         self.assertLess(probs[1, :].min()       ,warped_prob_smooth[1, :].min()  )





    def 							a__				(				self  )			->		Dict:
         UpperCAmelCase_      :							Optional[Any]       =      None
         UpperCAmelCase_      :							Union[str, Any]       =      10
         UpperCAmelCase_      :							Union[str, Any]       =      2

         # create ramp distribution
         UpperCAmelCase_      :							Dict       =      np.broadcast_to(np.arange(_SCREAMING_SNAKE_CASE  )[None, :]       ,(batch_size, vocab_size)  ).copy()
         UpperCAmelCase_      :							Dict       =      ramp_logits[1:, : vocab_size // 2] + vocab_size

         UpperCAmelCase_      :							Union[str, Any]       =      FlaxTopKLogitsWarper(3  )

         UpperCAmelCase_      :							Any       =      top_k_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )

         # check that correct tokens are filtered
         self.assertListEqual(jnp.isinf(scores[0]  ).tolist()       ,7 * [True] + 3 * [False]  )
         self.assertListEqual(jnp.isinf(scores[1]  ).tolist()       ,2 * [True] + 3 * [False] + 5 * [True]  )

         # check special case
         UpperCAmelCase_      :							Dict       =      5
         UpperCAmelCase_      :							str       =      FlaxTopKLogitsWarper(top_k=1       ,filter_value=0.0       ,min_tokens_to_keep=3  )

         UpperCAmelCase_      :							int       =      np.broadcast_to(np.arange(_SCREAMING_SNAKE_CASE  )[None, :]       ,(batch_size, length)  ).copy()
         UpperCAmelCase_      :							Optional[int]       =      top_k_warp_safety_check(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )

         # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
         self.assertListEqual((scores == 0.0).sum(axis=-1  ).tolist()       ,[2, 2]  )





    def 							a__				(				self  )			->		List[Any]:
         UpperCAmelCase_      :							List[Any]       =      None
         UpperCAmelCase_      :							Union[str, Any]       =      10
         UpperCAmelCase_      :							Optional[Any]       =      2

         # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
         UpperCAmelCase_      :							Union[str, Any]       =      np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]  )  )

         UpperCAmelCase_      :							Optional[int]       =      FlaxTopPLogitsWarper(0.8  )
         UpperCAmelCase_      :							str       =      np.exp(top_p_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )  )

         # dist should be filtered to keep min num values so that sum is >= top_p
         # exp (-inf) => 0
         UpperCAmelCase_      :							List[str]       =      np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]  )
         self.assertTrue(np.allclose(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,atol=1e-3  )  )

         # check edge cases with negative and extreme logits
         UpperCAmelCase_      :							Dict       =      np.broadcast_to(np.arange(_SCREAMING_SNAKE_CASE  )[None, :]       ,(batch_size, vocab_size)  ).copy() - (
             vocab_size // 2
         )

         # make ramp_logits more extreme
         UpperCAmelCase_      :							Any       =      ramp_logits[1] * 1_00.0

         # make sure at least 2 tokens are kept
         UpperCAmelCase_      :							Any       =      FlaxTopPLogitsWarper(0.9       ,min_tokens_to_keep=2       ,filter_value=0.0  )
         UpperCAmelCase_      :							Optional[Any]       =      top_p_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )

         # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
         self.assertListEqual((filtered_dist != 0.0).sum(axis=-1  ).tolist()       ,[3, 2]  )





    def 							a__				(				self  )			->		Any:
         UpperCAmelCase_      :							str       =      20
         UpperCAmelCase_      :							List[Any]       =      4
         UpperCAmelCase_      :							Optional[int]       =      0

         UpperCAmelCase_      :							Any       =      FlaxMinLengthLogitsProcessor(min_length=10       ,eos_token_id=_SCREAMING_SNAKE_CASE  )

         # check that min length is applied at length 5
         UpperCAmelCase_      :							List[str]       =      ids_tensor((batch_size, 20)       ,vocab_size=20  )
         UpperCAmelCase_      :							Optional[Any]       =      5
         UpperCAmelCase_      :							List[Any]       =      self._get_uniform_logits(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Optional[Any]       =      min_dist_processor(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()       ,4 * [-float('''inf'''  )]  )

         # check that min length is not applied anymore at length 15
         UpperCAmelCase_      :							Optional[Any]       =      self._get_uniform_logits(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Optional[Any]       =      15
         UpperCAmelCase_      :							Any       =      min_dist_processor(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         self.assertFalse(jnp.isinf(_SCREAMING_SNAKE_CASE  ).any()  )





    def 							a__				(				self  )			->		Tuple:
         UpperCAmelCase_      :							int       =      20
         UpperCAmelCase_      :							Any       =      4
         UpperCAmelCase_      :							List[Any]       =      0

         UpperCAmelCase_      :							Optional[Any]       =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_SCREAMING_SNAKE_CASE  )

         # check that all scores are -inf except the bos_token_id score
         UpperCAmelCase_      :							str       =      ids_tensor((batch_size, 1)       ,vocab_size=20  )
         UpperCAmelCase_      :							Dict       =      1
         UpperCAmelCase_      :							int       =      self._get_uniform_logits(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      logits_processor(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]  ).all()  )
         self.assertListEqual(scores[:, bos_token_id].tolist()       ,4 * [0]  )  # score for bos_token_id shold be zero

         # check that bos_token_id is not forced if current length is greater than 1
         UpperCAmelCase_      :							Dict       =      3
         UpperCAmelCase_      :							Tuple       =      self._get_uniform_logits(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Dict       =      logits_processor(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         self.assertFalse(jnp.isinf(_SCREAMING_SNAKE_CASE  ).any()  )





    def 							a__				(				self  )			->		Dict:
         UpperCAmelCase_      :							List[str]       =      20
         UpperCAmelCase_      :							Dict       =      4
         UpperCAmelCase_      :							List[str]       =      0
         UpperCAmelCase_      :							List[str]       =      5

         UpperCAmelCase_      :							str       =      FlaxForcedEOSTokenLogitsProcessor(max_length=_SCREAMING_SNAKE_CASE       ,eos_token_id=_SCREAMING_SNAKE_CASE  )

         # check that all scores are -inf except the eos_token_id when max_length is reached
         UpperCAmelCase_      :							List[str]       =      ids_tensor((batch_size, 4)       ,vocab_size=20  )
         UpperCAmelCase_      :							Any       =      4
         UpperCAmelCase_      :							Any       =      self._get_uniform_logits(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      logits_processor(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]  ).all()  )
         self.assertListEqual(scores[:, eos_token_id].tolist()       ,4 * [0]  )  # score for eos_token_id should be zero

         # check that eos_token_id is not forced if max_length is not reached
         UpperCAmelCase_      :							List[str]       =      3
         UpperCAmelCase_      :							str       =      self._get_uniform_logits(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[Any]       =      logits_processor(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         self.assertFalse(jnp.isinf(_SCREAMING_SNAKE_CASE  ).any()  )





    def 							a__				(				self  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Optional[Any]       =      4
         UpperCAmelCase_      :							Dict       =      10
         UpperCAmelCase_      :							List[str]       =      15
         UpperCAmelCase_      :							Optional[int]       =      2
         UpperCAmelCase_      :							Tuple       =      1
         UpperCAmelCase_      :							Union[str, Any]       =      15

         # dummy input_ids and scores
         UpperCAmelCase_      :							Union[str, Any]       =      ids_tensor((batch_size, sequence_length)       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      input_ids.copy()

         UpperCAmelCase_      :							str       =      self._get_uniform_logits(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Any       =      scores.copy()

         # instantiate all dist processors
         UpperCAmelCase_      :							List[Any]       =      FlaxTemperatureLogitsWarper(temperature=0.5  )
         UpperCAmelCase_      :							Tuple       =      FlaxTopKLogitsWarper(3  )
         UpperCAmelCase_      :							Tuple       =      FlaxTopPLogitsWarper(0.8  )

         # instantiate all logits processors
         UpperCAmelCase_      :							Optional[Any]       =      FlaxMinLengthLogitsProcessor(min_length=10       ,eos_token_id=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							str       =      FlaxForcedEOSTokenLogitsProcessor(max_length=_SCREAMING_SNAKE_CASE       ,eos_token_id=_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Optional[Any]       =      10

         # no processor list
         UpperCAmelCase_      :							List[str]       =      temp_dist_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      top_k_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      top_p_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Optional[Any]       =      min_dist_proc(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Any       =      bos_dist_proc(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Any       =      eos_dist_proc(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )

         # with processor list
         UpperCAmelCase_      :							Tuple       =      FlaxLogitsProcessorList(
             [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
         UpperCAmelCase_      :							List[str]       =      processor(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )

         # scores should be equal
         self.assertTrue(jnp.allclose(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,atol=1e-3  )  )

         # input_ids should never be changed
         self.assertListEqual(input_ids.tolist()       ,input_ids_comp.tolist()  )





    def 							a__				(				self  )			->		Dict:
         UpperCAmelCase_      :							Tuple       =      4
         UpperCAmelCase_      :							Any       =      10
         UpperCAmelCase_      :							Union[str, Any]       =      15
         UpperCAmelCase_      :							Optional[int]       =      2
         UpperCAmelCase_      :							Tuple       =      1
         UpperCAmelCase_      :							Optional[Any]       =      15

         # dummy input_ids and scores
         UpperCAmelCase_      :							List[Any]       =      ids_tensor((batch_size, sequence_length)       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							int       =      input_ids.copy()

         UpperCAmelCase_      :							Optional[int]       =      self._get_uniform_logits(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Optional[int]       =      scores.copy()

         # instantiate all dist processors
         UpperCAmelCase_      :							List[str]       =      FlaxTemperatureLogitsWarper(temperature=0.5  )
         UpperCAmelCase_      :							Dict       =      FlaxTopKLogitsWarper(3  )
         UpperCAmelCase_      :							int       =      FlaxTopPLogitsWarper(0.8  )

         # instantiate all logits processors
         UpperCAmelCase_      :							Dict       =      FlaxMinLengthLogitsProcessor(min_length=10       ,eos_token_id=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							int       =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      FlaxForcedEOSTokenLogitsProcessor(max_length=_SCREAMING_SNAKE_CASE       ,eos_token_id=_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Optional[int]       =      10

         # no processor list
         def run_no_processor_list(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ):
              UpperCAmelCase_      :							int       =      temp_dist_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Union[str, Any]       =      top_k_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							List[Any]       =      top_p_warp(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Dict       =      min_dist_proc(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Optional[Any]       =      bos_dist_proc(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Optional[Any]       =      eos_dist_proc(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
              return scores

         # with processor list
         def run_processor_list(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ):
              UpperCAmelCase_      :							List[str]       =      FlaxLogitsProcessorList(
                  [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
              UpperCAmelCase_      :							int       =      processor(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,cur_len=_SCREAMING_SNAKE_CASE  )
              return scores

         UpperCAmelCase_      :							Any       =      jax.jit(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      jax.jit(_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Optional[int]       =      jitted_run_no_processor_list(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      jitted_run_processor_list(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )

         # scores should be equal
         self.assertTrue(jnp.allclose(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,atol=1e-3  )  )

         # input_ids should never be changed
         self.assertListEqual(input_ids.tolist()       ,input_ids_comp.tolist()  )





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






    """simple docstring"""





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

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





    def 							a__				(				self  )			->		str:
         UpperCAmelCase_      :							List[Any]       =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

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

         UpperCAmelCase_      :							Optional[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 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )			->		Optional[Any]:
         UpperCAmelCase_      :							Union[str, Any]       =      FlaxBeitModel(config=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      model(_SCREAMING_SNAKE_CASE  )
         self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, self.seq_length, self.hidden_size)  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )			->		Tuple:
         UpperCAmelCase_      :							Union[str, Any]       =      FlaxBeitForMaskedImageModeling(config=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Any       =      model(_SCREAMING_SNAKE_CASE  )
         self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length - 1, self.vocab_size)  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )			->		str:
         UpperCAmelCase_      :							Dict       =      self.type_sequence_label_size
         UpperCAmelCase_      :							int       =      FlaxBeitForImageClassification(config=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      model(_SCREAMING_SNAKE_CASE  )
         self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.type_sequence_label_size)  )

         # test greyscale images
         UpperCAmelCase_      :							Any       =      1
         UpperCAmelCase_      :							List[Any]       =      FlaxBeitForImageClassification(_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							str       =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
         UpperCAmelCase_      :							Optional[Any]       =      model(_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							Optional[int]       =      self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase_
         ),   (
             UpperCAmelCase_
         ),   (
             UpperCAmelCase_
         ),   
         )      :							List[str]       =      config_and_inputs
         UpperCAmelCase_      :							int       =      {'''pixel_values''': pixel_values}
         return config, inputs_dict


@require_flax
class   __a(					_a    ,						unittest.TestCase       ):






    """simple docstring"""





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





    def 							a__				(				self  )			->		None:
         UpperCAmelCase_      :							List[Any]       =      FlaxBeitModelTester(self  )
         UpperCAmelCase_      :							List[str]       =      ConfigTester(self       ,config_class=_SCREAMING_SNAKE_CASE       ,has_text_modality=_SCREAMING_SNAKE_CASE       ,hidden_size=37  )





    def 							a__				(				self  )			->		Optional[int]:
         self.config_tester.run_common_tests()





    def 							a__				(				self  )			->		List[Any]:
         UpperCAmelCase_,   UpperCAmelCase_      :							Optional[Any]       =      self.model_tester.prepare_config_and_inputs_for_common()

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

              UpperCAmelCase_      :							Optional[int]       =      ['''pixel_values''']
              self.assertListEqual(arg_names[:1]       ,_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		Union[str, Any]:
         UpperCAmelCase_,   UpperCAmelCase_      :							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[Any]       =      self._prepare_for_class(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
                   UpperCAmelCase_      :							Union[str, Any]       =      model_class(_SCREAMING_SNAKE_CASE  )

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

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

                   with self.subTest('''JIT Disabled'''  ):
                        with jax.disable_jit():
                             UpperCAmelCase_      :							List[str]       =      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 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							Dict       =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Optional[Any]       =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							int       =      self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE  )





    @slow
    def 							a__				(				self  )			->		List[Any]:
         for model_class_name in self.all_model_classes:
              UpperCAmelCase_      :							List[Any]       =      model_class_name.from_pretrained('''microsoft/beit-base-patch16-224'''  )
              UpperCAmelCase_      :							Optional[int]       =      model(np.ones((1, 3, 224, 224)  )  )
              self.assertIsNotNone(_SCREAMING_SNAKE_CASE  )







def lowerCamelCase__       (      ):

     '''simple docstring'''


     UpperCAmelCase_      :							int       =      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
     return image



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






    """simple docstring"""





    @cached_property
    def 							a__				(				self  )			->		Dict:
         return BeitImageProcessor.from_pretrained('''microsoft/beit-base-patch16-224'''  ) if is_vision_available() else None





    @slow
    def 							a__				(				self  )			->		Optional[int]:
         UpperCAmelCase_      :							Tuple       =      FlaxBeitForMaskedImageModeling.from_pretrained('''microsoft/beit-base-patch16-224-pt22k'''  )

         UpperCAmelCase_      :							List[Any]       =      self.default_image_processor
         UpperCAmelCase_      :							Optional[Any]       =      prepare_img()
         UpperCAmelCase_      :							Optional[Any]       =      image_processor(images=_SCREAMING_SNAKE_CASE       ,return_tensors='''np'''  ).pixel_values

         # prepare bool_masked_pos
         UpperCAmelCase_      :							Union[str, Any]       =      np.ones((1, 196)       ,dtype=_SCREAMING_SNAKE_CASE  )

         # forward pass
         UpperCAmelCase_      :							Optional[int]       =      model(pixel_values=_SCREAMING_SNAKE_CASE       ,bool_masked_pos=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[str]       =      outputs.logits

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

         UpperCAmelCase_      :							List[str]       =      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]       ,_SCREAMING_SNAKE_CASE       ,atol=1e-2  )  )





    @slow
    def 							a__				(				self  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Any       =      FlaxBeitForImageClassification.from_pretrained('''microsoft/beit-base-patch16-224'''  )

         UpperCAmelCase_      :							Any       =      self.default_image_processor
         UpperCAmelCase_      :							Any       =      prepare_img()
         UpperCAmelCase_      :							Union[str, Any]       =      image_processor(images=_SCREAMING_SNAKE_CASE       ,return_tensors='''np'''  )

         # forward pass
         UpperCAmelCase_      :							Optional[Any]       =      model(**_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Optional[Any]       =      outputs.logits

         # verify the logits
         UpperCAmelCase_      :							Dict       =      (1, 1_000)
         self.assertEqual(logits.shape       ,_SCREAMING_SNAKE_CASE  )

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

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

         UpperCAmelCase_      :							Dict       =      281
         self.assertEqual(logits.argmax(-1  ).item()       ,_SCREAMING_SNAKE_CASE  )





    @slow
    def 							a__				(				self  )			->		Union[str, Any]:
         UpperCAmelCase_      :							str       =      FlaxBeitForImageClassification.from_pretrained('''microsoft/beit-large-patch16-224-pt22k-ft22k'''  )

         UpperCAmelCase_      :							Tuple       =      self.default_image_processor
         UpperCAmelCase_      :							Any       =      prepare_img()
         UpperCAmelCase_      :							Dict       =      image_processor(images=_SCREAMING_SNAKE_CASE       ,return_tensors='''np'''  )

         # forward pass
         UpperCAmelCase_      :							Dict       =      model(**_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Optional[Any]       =      outputs.logits

         # verify the logits
         UpperCAmelCase_      :							Union[str, Any]       =      (1, 21_841)
         self.assertEqual(logits.shape       ,_SCREAMING_SNAKE_CASE  )

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

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

         UpperCAmelCase_      :							Dict       =      2_396
         self.assertEqual(logits.argmax(-1  ).item()       ,_SCREAMING_SNAKE_CASE  )