





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







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






     '''simple docstring'''





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

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







  def   UpperCAmelCase					(      self  :Union[str, Any] ):






     '''simple docstring'''





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

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

     lowercase__			       =  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   UpperCAmelCase					(      self  :str   ,							_lowercase  :int   ,							_lowercase  :Any   ,							_lowercase  :List[Any] ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Union[str, Any]   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :int   ,							_lowercase  :int ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :List[Any]   ,							_lowercase  :Dict   ,							_lowercase  :Any ):






     '''simple docstring'''





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

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

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







  def   UpperCAmelCase					(      self  :str ):






     '''simple docstring'''





     lowercase__			       =  self.prepare_config_and_inputs()
     (
         (
         lowercase__ 
     )       ,							(
         lowercase__ 
     )       ,							(
         lowercase__ 
     )       ,							 
     )			       =  config_and_inputs
     lowercase__			       =  {"pixel_values": pixel_values}
     return config, inputs_dict





@require_flax
class  lowerCAmelCase							(       lowerCAmelCase__							,							unittest.TestCase ):
  __lowerCamelCase		=      (
      (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
  )







  def   UpperCAmelCase					(      self  :Tuple ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     self.config_tester.run_common_tests()







  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





     lowercase__       ,							lowercase__			       =  self.model_tester.prepare_config_and_inputs_for_common()

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

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







  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





     lowercase__       ,							lowercase__			       =  self.model_tester.prepare_config_and_inputs_for_common()

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

           @jax.jit
           def model_jitted(_lowercase  :Union[str, Any]   ,							**_lowercase  :Optional[int] ):
              return model(pixel_values=_SCREAMING_SNAKE_CASE   ,							**_SCREAMING_SNAKE_CASE )

           with self.subTest("JIT Enabled" ):
              lowercase__			       =  model_jitted(**_SCREAMING_SNAKE_CASE ).to_tuple()

           with self.subTest("JIT Disabled" ):
              with jax.disable_jit():
                 lowercase__			       =  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   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





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







  @slow
  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





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


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







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







  @cached_property
  def   UpperCAmelCase					(      self  :List[str] ):






     '''simple docstring'''





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







  @slow
  def   UpperCAmelCase					(      self  :str ):






     '''simple docstring'''





     lowercase__			       =  FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k" )

     lowercase__			       =  self.default_image_processor
     lowercase__			       =  prepare_img()
     lowercase__			       =  image_processor(images=_SCREAMING_SNAKE_CASE   ,							return_tensors="np" ).pixel_values

     # prepare bool_masked_pos
     lowercase__			       =  np.ones((1, 1_96)   ,							dtype=_SCREAMING_SNAKE_CASE )

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

     # verify the logits
     lowercase__			       =  (1, 1_96, 81_92)
     self.assertEqual(logits.shape   ,							_SCREAMING_SNAKE_CASE )

     lowercase__			       =  np.array(
         [[-3.2437, 0.5072, -13.91_74], [-3.2456, 0.4948, -13.94_01], [-3.2033, 0.5121, -13.85_50]] )

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







  @slow
  def   UpperCAmelCase					(      self  :str ):






     '''simple docstring'''





     lowercase__			       =  FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224" )

     lowercase__			       =  self.default_image_processor
     lowercase__			       =  prepare_img()
     lowercase__			       =  image_processor(images=_SCREAMING_SNAKE_CASE   ,							return_tensors="np" )

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

     # verify the logits
     lowercase__			       =  (1, 10_00)
     self.assertEqual(logits.shape   ,							_SCREAMING_SNAKE_CASE )

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

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

     lowercase__			       =  2_81
     self.assertEqual(logits.argmax(-1 ).item()   ,							_SCREAMING_SNAKE_CASE )







  @slow
  def   UpperCAmelCase					(      self  :str ):






     '''simple docstring'''





     lowercase__			       =  FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k" )

     lowercase__			       =  self.default_image_processor
     lowercase__			       =  prepare_img()
     lowercase__			       =  image_processor(images=_SCREAMING_SNAKE_CASE   ,							return_tensors="np" )

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

     # verify the logits
     lowercase__			       =  (1, 2_18_41)
     self.assertEqual(logits.shape   ,							_SCREAMING_SNAKE_CASE )

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

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

     lowercase__			       =  23_96
     self.assertEqual(logits.argmax(-1 ).item()   ,							_SCREAMING_SNAKE_CASE )













from __future__ import annotations


def _A       (   __magic_name__  , __magic_name__	):
   lowercase__			       =  []
   create_all_state(1  , __magic_name__  , __magic_name__  , []  , __magic_name__	)
   return result


def _A       (   __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , ):
   if level == 0:
      total_list.append(current_list[:]	)
      return

   for i in range(__magic_name__  , total_number - level + 2	):
      current_list.append(__magic_name__	)
      create_all_state(i + 1  , __magic_name__  , level - 1  , __magic_name__  , __magic_name__	)
      current_list.pop()


def _A       (   __magic_name__	):
   for i in total_list:
      print(*__magic_name__	)


if __name__ == "__main__":
  _snake_case           =			4
  _snake_case           =			2
  _snake_case           =			generate_all_combinations(n, k)
  print_all_state(total_list)






