
"""simple docstring"""





from __future__ import annotations

import os
import tempfile
import unittest

from transformers import ConvBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TFConvBertForMaskedLM,
        TFConvBertForMultipleChoice,
        TFConvBertForQuestionAnswering,
        TFConvBertForSequenceClassification,
        TFConvBertForTokenClassification,
        TFConvBertModel,
    )


class A_	:




  '''simple docstring'''







  def __init__(		self				,   lowercase_				,   lowercase_=13				,   lowercase_=7				,   lowercase_=True				,   lowercase_=True				,   lowercase_=True				,   lowercase_=True				,   lowercase_=99				,   lowercase_=32				,   lowercase_=2				,   lowercase_=4				,   lowercase_=37				,   lowercase_="gelu"				,   lowercase_=0.1				,   lowercase_=0.1				,   lowercase_=512				,   lowercase_=16				,   lowercase_=2				,   lowercase_=0.02				,   lowercase_=3				,   lowercase_=4				,   lowercase_=None				,   ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    parent
     UpperCAmelCase_					:					Dict			     =    13
     UpperCAmelCase_					:					List[str]			     =    7
     UpperCAmelCase_					:					Dict			     =    True
     UpperCAmelCase_					:					Optional[Any]			     =    True
     UpperCAmelCase_					:					Union[str, Any]			     =    True
     UpperCAmelCase_					:					Dict			     =    True
     UpperCAmelCase_					:					Union[str, Any]			     =    99
     UpperCAmelCase_					:					Tuple			     =    384
     UpperCAmelCase_					:					Optional[Any]			     =    2
     UpperCAmelCase_					:					Any			     =    4
     UpperCAmelCase_					:					int			     =    37
     UpperCAmelCase_					:					Dict			     =    "gelu"
     UpperCAmelCase_					:					Optional[int]			     =    0.1
     UpperCAmelCase_					:					Tuple			     =    0.1
     UpperCAmelCase_					:					Optional[Any]			     =    512
     UpperCAmelCase_					:					Dict			     =    16
     UpperCAmelCase_					:					int			     =    2
     UpperCAmelCase_					:					Union[str, Any]			     =    0.02
     UpperCAmelCase_					:					int			     =    3
     UpperCAmelCase_					:					Optional[Any]			     =    4
     UpperCAmelCase_					:					Any			     =    128
     UpperCAmelCase_					:					List[str]			     =    2
     UpperCAmelCase_					:					List[Any]			     =    9
     UpperCAmelCase_					:					List[Any]			     =    1
     UpperCAmelCase_					:					List[Any]			     =    None







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Any			     =    ids_tensor([self.batch_size, self.seq_length]				,   self.vocab_size  )

     UpperCAmelCase_					:					List[str]			     =    None
     if self.use_input_mask:
        UpperCAmelCase_					:					Tuple			     =    random_attention_mask([self.batch_size, self.seq_length]  )

     UpperCAmelCase_					:					List[Any]			     =    None
     if self.use_token_type_ids:
        UpperCAmelCase_					:					List[str]			     =    ids_tensor([self.batch_size, self.seq_length]				,   self.type_vocab_size  )

     UpperCAmelCase_					:					Optional[Any]			     =    None
     UpperCAmelCase_					:					Dict			     =    None
     UpperCAmelCase_					:					Tuple			     =    None
     if self.use_labels:
        UpperCAmelCase_					:					List[str]			     =    ids_tensor([self.batch_size]				,   self.type_sequence_label_size  )
        UpperCAmelCase_					:					List[str]			     =    ids_tensor([self.batch_size, self.seq_length]				,   self.num_labels  )
        UpperCAmelCase_					:					Optional[Any]			     =    ids_tensor([self.batch_size]				,   self.num_choices  )

     UpperCAmelCase_					:					str			     =    ConvBertConfig(
         vocab_size=self.vocab_size				,   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				,   max_position_embeddings=self.max_position_embeddings				,   type_vocab_size=self.type_vocab_size				,   initializer_range=self.initializer_range				,   return_dict=lowercase_				,   )

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    TFConvBertModel(config=lowercase_  )
     UpperCAmelCase_					:					int			     =    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

     UpperCAmelCase_					:					List[str]			     =    [input_ids, input_mask]
     UpperCAmelCase_					:					Union[str, Any]			     =    model(lowercase_  )

     UpperCAmelCase_					:					str			     =    model(lowercase_  )

     self.parent.assertEqual(result.last_hidden_state.shape				,   (self.batch_size, self.seq_length, self.hidden_size)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[Any]			     =    TFConvBertForMaskedLM(config=lowercase_  )
     UpperCAmelCase_					:					Tuple			     =    {
         "input_ids": input_ids,
         "attention_mask": input_mask,
         "token_type_ids": token_type_ids,
     }
     UpperCAmelCase_					:					Tuple			     =    model(lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.seq_length, self.vocab_size)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    self.num_labels
     UpperCAmelCase_					:					List[str]			     =    TFConvBertForSequenceClassification(config=lowercase_  )
     UpperCAmelCase_					:					List[str]			     =    {
         "input_ids": input_ids,
         "attention_mask": input_mask,
         "token_type_ids": token_type_ids,
     }

     UpperCAmelCase_					:					List[Any]			     =    model(lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.num_labels)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					str			     =    self.num_choices
     UpperCAmelCase_					:					Tuple			     =    TFConvBertForMultipleChoice(config=lowercase_  )
     UpperCAmelCase_					:					str			     =    tf.tile(tf.expand_dims(lowercase_				,   1  )				,   (1, self.num_choices, 1)  )
     UpperCAmelCase_					:					Tuple			     =    tf.tile(tf.expand_dims(lowercase_				,   1  )				,   (1, self.num_choices, 1)  )
     UpperCAmelCase_					:					Optional[Any]			     =    tf.tile(tf.expand_dims(lowercase_				,   1  )				,   (1, self.num_choices, 1)  )
     UpperCAmelCase_					:					Any			     =    {
         "input_ids": multiple_choice_inputs_ids,
         "attention_mask": multiple_choice_input_mask,
         "token_type_ids": multiple_choice_token_type_ids,
     }
     UpperCAmelCase_					:					Union[str, Any]			     =    model(lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.num_choices)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Dict			     =    self.num_labels
     UpperCAmelCase_					:					int			     =    TFConvBertForTokenClassification(config=lowercase_  )
     UpperCAmelCase_					:					Optional[Any]			     =    {
         "input_ids": input_ids,
         "attention_mask": input_mask,
         "token_type_ids": token_type_ids,
     }
     UpperCAmelCase_					:					Any			     =    model(lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.seq_length, self.num_labels)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    TFConvBertForQuestionAnswering(config=lowercase_  )
     UpperCAmelCase_					:					str			     =    {
         "input_ids": input_ids,
         "attention_mask": input_mask,
         "token_type_ids": token_type_ids,
     }

     UpperCAmelCase_					:					Optional[int]			     =    model(lowercase_  )
     self.parent.assertEqual(result.start_logits.shape				,   (self.batch_size, self.seq_length)  )
     self.parent.assertEqual(result.end_logits.shape				,   (self.batch_size, self.seq_length)  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Dict			     =    self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   
     )					:					Optional[int]			     =    config_and_inputs
     UpperCAmelCase_					:					Tuple			     =    {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
     return config, inputs_dict




@require_tf
class A_	(lowercase__				,lowercase__				,unittest.TestCase      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						Union[str, Any]       =		(
      (
          TFConvBertModel,
          TFConvBertForMaskedLM,
          TFConvBertForQuestionAnswering,
          TFConvBertForSequenceClassification,
          TFConvBertForTokenClassification,
          TFConvBertForMultipleChoice,
      )
      if is_tf_available()
      else ()
  )
  SCREAMING_SNAKE_CASE__			:						Optional[int]       =		(
      {
          """feature-extraction""": TFConvBertModel,
          """fill-mask""": TFConvBertForMaskedLM,
          """question-answering""": TFConvBertForQuestionAnswering,
          """text-classification""": TFConvBertForSequenceClassification,
          """token-classification""": TFConvBertForTokenClassification,
          """zero-shot""": TFConvBertForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  SCREAMING_SNAKE_CASE__			:						Optional[int]       =		False
  SCREAMING_SNAKE_CASE__			:						List[str]       =		False
  SCREAMING_SNAKE_CASE__			:						Dict       =		False







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					str			     =    TFConvBertModelTester(self  )
     UpperCAmelCase_					:					int			     =    ConfigTester(self				,   config_class=lowercase_				,   hidden_size=37  )







  def        UpperCamelCase__					(		self  ):

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







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*lowercase_  )







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					Optional[Any]			     =    self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase_					:					Optional[int]			     =    True
     UpperCAmelCase_					:					Tuple			     =    True

     if hasattr(lowercase_				,   "use_cache"  ):
        UpperCAmelCase_					:					Optional[int]			     =    True

     UpperCAmelCase_					:					Any			     =    getattr(self.model_tester				,   "encoder_seq_length"				,   self.model_tester.seq_length  )
     UpperCAmelCase_					:					List[Any]			     =    getattr(self.model_tester				,   "key_length"				,   lowercase_  )

     for model_class in self.all_model_classes:
        UpperCAmelCase_					:					Dict			     =    self._prepare_for_class(lowercase_				,   lowercase_  )
        UpperCAmelCase_					:					Dict			     =    model_class(lowercase_  )
        UpperCAmelCase_					:					Dict			     =    len(model(lowercase_  )  )

        with tempfile.TemporaryDirectory() as tmpdirname:
           model.save_pretrained(lowercase_				,   saved_model=lowercase_  )
           UpperCAmelCase_					:					Union[str, Any]			     =    os.path.join(lowercase_				,   "saved_model"				,   "1"  )
           UpperCAmelCase_					:					str			     =    tf.keras.models.load_model(lowercase_  )
           UpperCAmelCase_					:					Optional[int]			     =    model(lowercase_  )

           if self.is_encoder_decoder:
              UpperCAmelCase_					:					List[Any]			     =    outputs["encoder_hidden_states"]
              UpperCAmelCase_					:					Union[str, Any]			     =    outputs["encoder_attentions"]
           else:
              UpperCAmelCase_					:					List[Any]			     =    outputs["hidden_states"]
              UpperCAmelCase_					:					int			     =    outputs["attentions"]

           self.assertEqual(len(lowercase_  )				,   lowercase_  )

           UpperCAmelCase_					:					str			     =    getattr(
               self.model_tester				,   "expected_num_hidden_layers"				,   self.model_tester.num_hidden_layers + 1  )

           self.assertEqual(len(lowercase_  )				,   lowercase_  )
           self.assertListEqual(
               list(output_hidden_states[0].shape[-2:]  )				,   [self.model_tester.seq_length, self.model_tester.hidden_size]				,   )

           self.assertEqual(len(lowercase_  )				,   self.model_tester.num_hidden_layers  )
           self.assertListEqual(
               list(output_attentions[0].shape[-3:]  )				,   [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]				,   )







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    TFConvBertModel.from_pretrained("YituTech/conv-bert-base"  )
     self.assertIsNotNone(lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					str			     =    self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase_					:					Any			     =    True
     UpperCAmelCase_					:					List[Any]			     =    getattr(self.model_tester				,   "decoder_seq_length"				,   self.model_tester.seq_length  )
     UpperCAmelCase_					:					Tuple			     =    getattr(self.model_tester				,   "encoder_seq_length"				,   self.model_tester.seq_length  )
     UpperCAmelCase_					:					Optional[Any]			     =    getattr(self.model_tester				,   "key_length"				,   lowercase_  )
     UpperCAmelCase_					:					List[Any]			     =    getattr(self.model_tester				,   "key_length"				,   lowercase_  )

     def check_decoder_attentions_output(lowercase_  ):
        UpperCAmelCase_					:					Tuple			     =    len(lowercase_  )
        self.assertEqual(out_len % 2				,   0  )
        UpperCAmelCase_					:					Any			     =    outputs.decoder_attentions
        self.assertEqual(len(lowercase_  )				,   self.model_tester.num_hidden_layers  )
        self.assertListEqual(
            list(decoder_attentions[0].shape[-3:]  )				,   [self.model_tester.num_attention_heads / 2, decoder_seq_length, decoder_key_length]				,   )

     def check_encoder_attentions_output(lowercase_  ):
        UpperCAmelCase_					:					Union[str, Any]			     =    [
            t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
        ]
        self.assertEqual(len(lowercase_  )				,   self.model_tester.num_hidden_layers  )
        self.assertListEqual(
            list(attentions[0].shape[-3:]  )				,   [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]				,   )

     for model_class in self.all_model_classes:
        UpperCAmelCase_					:					str			     =    True
        UpperCAmelCase_					:					Any			     =    False
        UpperCAmelCase_					:					List[str]			     =    model_class(lowercase_  )
        UpperCAmelCase_					:					int			     =    model(self._prepare_for_class(lowercase_				,   lowercase_  )  )
        UpperCAmelCase_					:					Any			     =    len(lowercase_  )
        self.assertEqual(config.output_hidden_states				,   lowercase_  )
        check_encoder_attentions_output(lowercase_  )

        if self.is_encoder_decoder:
           UpperCAmelCase_					:					List[Any]			     =    model_class(lowercase_  )
           UpperCAmelCase_					:					Optional[Any]			     =    model(self._prepare_for_class(lowercase_				,   lowercase_  )  )
           self.assertEqual(config.output_hidden_states				,   lowercase_  )
           check_decoder_attentions_output(lowercase_  )

        # Check that output attentions can also be changed via the config
        del inputs_dict["output_attentions"]
        UpperCAmelCase_					:					Tuple			     =    True
        UpperCAmelCase_					:					List[Any]			     =    model_class(lowercase_  )
        UpperCAmelCase_					:					Optional[Any]			     =    model(self._prepare_for_class(lowercase_				,   lowercase_  )  )
        self.assertEqual(config.output_hidden_states				,   lowercase_  )
        check_encoder_attentions_output(lowercase_  )

        # Check attention is always last and order is fine
        UpperCAmelCase_					:					Union[str, Any]			     =    True
        UpperCAmelCase_					:					List[Any]			     =    True
        UpperCAmelCase_					:					Optional[Any]			     =    model_class(lowercase_  )
        UpperCAmelCase_					:					List[str]			     =    model(self._prepare_for_class(lowercase_				,   lowercase_  )  )

        self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)				,   len(lowercase_  )  )
        self.assertEqual(model.config.output_hidden_states				,   lowercase_  )
        check_encoder_attentions_output(lowercase_  )




@require_tf
class A_	(unittest.TestCase      ):




  '''simple docstring'''







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    TFConvBertModel.from_pretrained("YituTech/conv-bert-base"  )
     UpperCAmelCase_					:					Tuple			     =    tf.constant([[0, 1, 2, 3, 4, 5]]  )
     UpperCAmelCase_					:					str			     =    model(lowercase_  )[0]

     UpperCAmelCase_					:					Any			     =    [1, 6, 768]
     self.assertEqual(output.shape				,   lowercase_  )

     UpperCAmelCase_					:					str			     =    tf.constant(
         [
             [
                 [-0.03_47_54_93, -0.4_68_60_34, -0.30_63_88_32],
                 [0.22_63_72_48, -0.26_98_86_46, -0.7_42_34_24],
                 [0.10_32_48_68, -0.45_01_35_08, -0.58_28_07_84],
             ]
         ]  )
     tf.debugging.assert_near(output[:, :3, :3]				,   lowercase_				,   atol=1E-4  )




"""simple docstring"""





import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
    import jax

    from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel


class A_	(unittest.TestCase      ):




  '''simple docstring'''







  def __init__(		self				,   lowercase_				,   lowercase_=13				,   lowercase_=30				,   lowercase_=2				,   lowercase_=3				,   lowercase_=True				,   lowercase_=True				,   lowercase_=32				,   lowercase_=5				,   lowercase_=4				,   lowercase_=37				,   lowercase_="gelu"				,   lowercase_=0.1				,   lowercase_=0.1				,   lowercase_=10				,   lowercase_=0.02				,   ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    parent
     UpperCAmelCase_					:					List[str]			     =    batch_size
     UpperCAmelCase_					:					Union[str, Any]			     =    image_size
     UpperCAmelCase_					:					List[str]			     =    patch_size
     UpperCAmelCase_					:					Union[str, Any]			     =    num_channels
     UpperCAmelCase_					:					Optional[int]			     =    is_training
     UpperCAmelCase_					:					Dict			     =    use_labels
     UpperCAmelCase_					:					Any			     =    hidden_size
     UpperCAmelCase_					:					Optional[Any]			     =    num_hidden_layers
     UpperCAmelCase_					:					Optional[Any]			     =    num_attention_heads
     UpperCAmelCase_					:					Dict			     =    intermediate_size
     UpperCAmelCase_					:					Optional[Any]			     =    hidden_act
     UpperCAmelCase_					:					Optional[Any]			     =    hidden_dropout_prob
     UpperCAmelCase_					:					Tuple			     =    attention_probs_dropout_prob
     UpperCAmelCase_					:					Dict			     =    type_sequence_label_size
     UpperCAmelCase_					:					Optional[Any]			     =    initializer_range

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







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[Any]			     =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

     UpperCAmelCase_					:					Dict			     =    ViTConfig(
         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=lowercase_				,   initializer_range=self.initializer_range				,   )

     return config, pixel_values







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    FlaxViTModel(config=lowercase_  )
     UpperCAmelCase_					:					int			     =    model(lowercase_  )
     # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
     UpperCAmelCase_					:					Optional[Any]			     =    (self.image_size, self.image_size)
     UpperCAmelCase_					:					List[Any]			     =    (self.patch_size, self.patch_size)
     UpperCAmelCase_					:					str			     =    (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
     self.parent.assertEqual(result.last_hidden_state.shape				,   (self.batch_size, num_patches + 1, self.hidden_size)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    self.type_sequence_label_size
     UpperCAmelCase_					:					Tuple			     =    FlaxViTForImageClassification(config=lowercase_  )
     UpperCAmelCase_					:					str			     =    model(lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.type_sequence_label_size)  )

     # test greyscale images
     UpperCAmelCase_					:					Any			     =    1
     UpperCAmelCase_					:					Optional[int]			     =    FlaxViTForImageClassification(lowercase_  )

     UpperCAmelCase_					:					List[Any]			     =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
     UpperCAmelCase_					:					List[Any]			     =    model(lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   
     )					:					Tuple			     =    config_and_inputs
     UpperCAmelCase_					:					Union[str, Any]			     =    {"pixel_values": pixel_values}
     return config, inputs_dict




@require_flax
class A_	(lowercase__				,unittest.TestCase      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						Any       =		(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    FlaxViTModelTester(self  )
     UpperCAmelCase_					:					Dict			     =    ConfigTester(self				,   config_class=lowercase_				,   has_text_modality=lowercase_				,   hidden_size=37  )







  def        UpperCamelCase__					(		self  ):

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







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[Any]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					List[str]			     =    self.model_tester.prepare_config_and_inputs_for_common()

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

        UpperCAmelCase_					:					List[str]			     =    ["pixel_values"]
        self.assertListEqual(arg_names[:1]				,   lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					Tuple			     =    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__  ):
           UpperCAmelCase_					:					Union[str, Any]			     =    self._prepare_for_class(lowercase_				,   lowercase_  )
           UpperCAmelCase_					:					Tuple			     =    model_class(lowercase_  )

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

           with self.subTest("JIT Enabled"  ):
              UpperCAmelCase_					:					Union[str, Any]			     =    model_jitted(**lowercase_  ).to_tuple()

           with self.subTest("JIT Disabled"  ):
              with jax.disable_jit():
                 UpperCAmelCase_					:					Tuple			     =    model_jitted(**lowercase_  ).to_tuple()

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







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     for model_class_name in self.all_model_classes:
        UpperCAmelCase_					:					Union[str, Any]			     =    model_class_name.from_pretrained("google/vit-base-patch16-224"  )
        UpperCAmelCase_					:					List[str]			     =    model(np.ones((1, 3, 224, 224)  )  )
        self.assertIsNotNone(lowercase_  )


