





import json
import os
from functools import lru_cache
from typing import Dict, List, Optional, Tuple, Union

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...tokenization_utils_base import BatchEncoding, EncodedInput
from ...utils import PaddingStrategy, logging


_UpperCAmelCase             =							logging.get_logger(__name__)


_UpperCAmelCase             =							{'vocab_file': 'vocab.json', 'merges_file': 'merges.txt'}

# See all LED models at https://huggingface.co/models?filter=LED
_UpperCAmelCase             =							{
    'vocab_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/vocab.json',
    },
    'merges_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/merges.txt',
    },
    'tokenizer_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/tokenizer.json',
    },
}

_UpperCAmelCase             =							{
    'allenai/led-base-16384': 1_6384,
}






@lru_cache()
# Copied from transformers.models.bart.tokenization_bart.bytes_to_unicode
def 	_SCREAMING_SNAKE_CASE		(      )      ->   Optional[Any]:
	__lowerCAmelCase					:	int      =       (
	    list(range(ord("""!"""   )   ,   ord("""~"""   ) + 1   )   ) + list(range(ord("""¡"""   )   ,   ord("""¬"""   ) + 1   )   ) + list(range(ord("""®"""   )   ,   ord("""ÿ"""   ) + 1   )   )
	)
	__lowerCAmelCase					:	Optional[int]      =       bs[:]
	__lowerCAmelCase					:	List[Any]      =       0
	for b in range(2**8   ):
		if b not in bs:
			bs.append(SCREAMING_SNAKE_CASE   )
			cs.append(2**8 + n   )
			n += 1
	__lowerCAmelCase					:	Tuple      =       [chr(SCREAMING_SNAKE_CASE   ) for n in cs]
	return dict(zip(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )   )






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:str   )      ->   Union[str, Any]:
	__lowerCAmelCase					:	List[Any]      =       set()
	__lowerCAmelCase					:	Optional[Any]      =       word[0]
	for char in word[1:]:
		pairs.add((prev_char, char)   )
		__lowerCAmelCase					:	str      =       char
	return pairs







class        snake_case_    (		__lowercase						):
		A_      							=   VOCAB_FILES_NAMES
		A_      							=   PRETRAINED_VOCAB_FILES_MAP
		A_      							=   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		A_      							=   ['input_ids', 'attention_mask']






		def __init__(							self			:					Dict       ,  _snake_case			:					List[str]       ,  _snake_case			:					Dict       ,  _snake_case			:					Tuple="replace"       ,  _snake_case			:					List[Any]="<s>"       ,  _snake_case			:					Optional[int]="</s>"       ,  _snake_case			:					Tuple="</s>"       ,  _snake_case			:					Any="<s>"       ,  _snake_case			:					str="<unk>"       ,  _snake_case			:					Dict="<pad>"       ,  _snake_case			:					Dict="<mask>"       ,  _snake_case			:					str=False       ,  **_snake_case			:					str       ,  )->Dict:






			'''simple docstring'''

			__lowerCAmelCase					:	Optional[int]      =       AddedToken(_snake_case       ,  lstrip=_snake_case       ,  rstrip=_snake_case      ) if isinstance(_snake_case       ,  _snake_case      ) else bos_token
			__lowerCAmelCase					:	List[str]      =       AddedToken(_snake_case       ,  lstrip=_snake_case       ,  rstrip=_snake_case      ) if isinstance(_snake_case       ,  _snake_case      ) else eos_token
			__lowerCAmelCase					:	List[str]      =       AddedToken(_snake_case       ,  lstrip=_snake_case       ,  rstrip=_snake_case      ) if isinstance(_snake_case       ,  _snake_case      ) else sep_token
			__lowerCAmelCase					:	int      =       AddedToken(_snake_case       ,  lstrip=_snake_case       ,  rstrip=_snake_case      ) if isinstance(_snake_case       ,  _snake_case      ) else cls_token
			__lowerCAmelCase					:	Union[str, Any]      =       AddedToken(_snake_case       ,  lstrip=_snake_case       ,  rstrip=_snake_case      ) if isinstance(_snake_case       ,  _snake_case      ) else unk_token
			__lowerCAmelCase					:	Union[str, Any]      =       AddedToken(_snake_case       ,  lstrip=_snake_case       ,  rstrip=_snake_case      ) if isinstance(_snake_case       ,  _snake_case      ) else pad_token

			# Mask token behave like a normal word, i.e. include the space before it
			__lowerCAmelCase					:	Any      =       AddedToken(_snake_case       ,  lstrip=_snake_case       ,  rstrip=_snake_case      ) if isinstance(_snake_case       ,  _snake_case      ) else mask_token

			super().__init__(
			    errors=_snake_case       ,  bos_token=_snake_case       ,  eos_token=_snake_case       ,  unk_token=_snake_case       ,  sep_token=_snake_case       ,  cls_token=_snake_case       ,  pad_token=_snake_case       ,  mask_token=_snake_case       ,  add_prefix_space=_snake_case       ,  **_snake_case       ,  )

			with open(_snake_case       ,  encoding="""utf-8"""      ) as vocab_handle:
				__lowerCAmelCase					:	Tuple      =       json.load(_snake_case      )
			__lowerCAmelCase					:	Tuple      =       {v: k for k, v in self.encoder.items()}
			__lowerCAmelCase					:	List[str]      =       errors  # how to handle errors in decoding
			__lowerCAmelCase					:	Dict      =       bytes_to_unicode()
			__lowerCAmelCase					:	List[str]      =       {v: k for k, v in self.byte_encoder.items()}
			with open(_snake_case       ,  encoding="""utf-8"""      ) as merges_handle:
				__lowerCAmelCase					:	Dict      =       merges_handle.read().split("""\n"""      )[1:-1]
			__lowerCAmelCase					:	Optional[int]      =       [tuple(merge.split()      ) for merge in bpe_merges]
			__lowerCAmelCase					:	List[str]      =       dict(zip(_snake_case       ,  range(len(_snake_case      )      )      )      )
			__lowerCAmelCase					:	str      =       {}
			__lowerCAmelCase					:	Tuple      =       add_prefix_space

			# Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
			__lowerCAmelCase					:	Union[str, Any]      =       re.compile(R"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""      )






		@property
		# Copied from transformers.models.bart.tokenization_bart.BartTokenizer.vocab_size
		def 				UpperCAmelCase__						(							self			:					str      )->List[str]:






			'''simple docstring'''

			return len(self.encoder      )






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






			'''simple docstring'''

			return dict(self.encoder       ,  **self.added_tokens_encoder      )






		def 				UpperCAmelCase__						(							self			:					Optional[Any]       ,  _snake_case			:					Tuple      )->List[str]:






			'''simple docstring'''

			if token in self.cache:
				return self.cache[token]
			__lowerCAmelCase					:	List[Any]      =       tuple(_snake_case      )
			__lowerCAmelCase					:	Union[str, Any]      =       get_pairs(_snake_case      )

			if not pairs:
				return token

			while True:
				__lowerCAmelCase					:	List[Any]      =       min(_snake_case       ,  key=lambda _snake_case      : self.bpe_ranks.get(_snake_case       ,  float("""inf"""      )      )      )
				if bigram not in self.bpe_ranks:
					break
				__lowerCAmelCase					:	Union[str, Any]      =       bigram
				__lowerCAmelCase					:	Tuple      =       []
				__lowerCAmelCase					:	Optional[int]      =       0
				while i < len(_snake_case      ):
					try:
						__lowerCAmelCase					:	Any      =       word.index(_snake_case       ,  _snake_case      )
					except ValueError:
						new_word.extend(word[i:]      )
						break
					else:
						new_word.extend(word[i:j]      )
						__lowerCAmelCase					:	Any      =       j

					if word[i] == first and i < len(_snake_case      ) - 1 and word[i + 1] == second:
						new_word.append(first + second      )
						i += 2
					else:
						new_word.append(word[i]      )
						i += 1
				__lowerCAmelCase					:	Optional[int]      =       tuple(_snake_case      )
				__lowerCAmelCase					:	Optional[int]      =       new_word
				if len(_snake_case      ) == 1:
					break
				else:
					__lowerCAmelCase					:	Optional[int]      =       get_pairs(_snake_case      )
			__lowerCAmelCase					:	Dict      =       """ """.join(_snake_case      )
			__lowerCAmelCase					:	Union[str, Any]      =       word
			return word






		def 				UpperCAmelCase__						(							self			:					List[Any]       ,  _snake_case			:					Optional[int]      )->List[str]:






			'''simple docstring'''

			__lowerCAmelCase					:	List[Any]      =       []
			for token in re.findall(self.pat       ,  _snake_case      ):
				__lowerCAmelCase					:	Dict      =       """""".join(
				    self.byte_encoder[b] for b in token.encode("""utf-8"""      )      )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
				bpe_tokens.extend(bpe_token for bpe_token in self.bpe(_snake_case      ).split(""" """      )      )
			return bpe_tokens






		def 				UpperCAmelCase__						(							self			:					Tuple       ,  _snake_case			:					List[str]      )->Tuple:






			'''simple docstring'''

			return self.encoder.get(_snake_case       ,  self.encoder.get(self.unk_token      )      )






		def 				UpperCAmelCase__						(							self			:					List[str]       ,  _snake_case			:					List[str]      )->Union[str, Any]:






			'''simple docstring'''

			return self.decoder.get(_snake_case      )






		def 				UpperCAmelCase__						(							self			:					List[str]       ,  _snake_case			:					List[Any]      )->Optional[int]:






			'''simple docstring'''

			__lowerCAmelCase					:	Dict      =       """""".join(_snake_case      )
			__lowerCAmelCase					:	int      =       bytearray([self.byte_decoder[c] for c in text]      ).decode("""utf-8"""       ,  errors=self.errors      )
			return text






		def 				UpperCAmelCase__						(							self			:					List[Any]       ,  _snake_case			:					str       ,  _snake_case			:					Optional[str] = None      )->Tuple[str]:






			'''simple docstring'''

			if not os.path.isdir(_snake_case      ):
				logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''      )
				return
			__lowerCAmelCase					:	int      =       os.path.join(
			    _snake_case       ,  (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]      )
			__lowerCAmelCase					:	List[str]      =       os.path.join(
			    _snake_case       ,  (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]      )

			with open(_snake_case       ,  """w"""       ,  encoding="""utf-8"""      ) as f:
				f.write(json.dumps(self.encoder       ,  indent=2       ,  sort_keys=_snake_case       ,  ensure_ascii=_snake_case      ) + """\n"""      )

			__lowerCAmelCase					:	List[str]      =       0
			with open(_snake_case       ,  """w"""       ,  encoding="""utf-8"""      ) as writer:
				writer.write("""#version: 0.2\n"""      )
				for bpe_tokens, token_index in sorted(self.bpe_ranks.items()       ,  key=lambda _snake_case      : kv[1]      ):
					if index != token_index:
						logger.warning(
						    F'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
						    """ Please check that the tokenizer is not corrupted!"""      )
						__lowerCAmelCase					:	Dict      =       token_index
					writer.write(""" """.join(_snake_case      ) + """\n"""      )
					index += 1

			return vocab_file, merge_file






		def 				UpperCAmelCase__						(							self			:					Tuple       ,  _snake_case			:					List[int]       ,  _snake_case			:					Optional[List[int]] = None      )->List[int]:






			'''simple docstring'''

			if token_ids_a is None:
				return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
			__lowerCAmelCase					:	List[str]      =       [self.cls_token_id]
			__lowerCAmelCase					:	Tuple      =       [self.sep_token_id]
			return cls + token_ids_a + sep + sep + token_ids_a + sep






		def 				UpperCAmelCase__						(							self			:					Optional[int]       ,  _snake_case			:					List[int]       ,  _snake_case			:					Optional[List[int]] = None       ,  _snake_case			:					bool = False      )->List[int]:






			'''simple docstring'''

			if already_has_special_tokens:
				return super().get_special_tokens_mask(
				    token_ids_a=_snake_case       ,  token_ids_a=_snake_case       ,  already_has_special_tokens=_snake_case      )

			if token_ids_a is None:
				return [1] + ([0] * len(_snake_case      )) + [1]
			return [1] + ([0] * len(_snake_case      )) + [1, 1] + ([0] * len(_snake_case      )) + [1]






		def 				UpperCAmelCase__						(							self			:					Dict       ,  _snake_case			:					List[int]       ,  _snake_case			:					Optional[List[int]] = None      )->List[int]:






			'''simple docstring'''

			__lowerCAmelCase					:	Any      =       [self.sep_token_id]
			__lowerCAmelCase					:	Union[str, Any]      =       [self.cls_token_id]

			if token_ids_a is None:
				return len(cls + token_ids_a + sep      ) * [0]
			return len(cls + token_ids_a + sep + sep + token_ids_a + sep      ) * [0]






		def 				UpperCAmelCase__						(							self			:					Dict       ,  _snake_case			:					List[str]       ,  _snake_case			:					Optional[Any]=False       ,  **_snake_case			:					int      )->Optional[int]:






			'''simple docstring'''

			__lowerCAmelCase					:	Union[str, Any]      =       kwargs.pop("""add_prefix_space"""       ,  self.add_prefix_space      )
			if (is_split_into_words or add_prefix_space) and (len(_snake_case      ) > 0 and not text[0].isspace()):
				__lowerCAmelCase					:	Union[str, Any]      =       """ """ + text
			return (text, kwargs)






		def 				UpperCAmelCase__						(							self			:					Union[str, Any]       ,  _snake_case			:					Union[Dict[str, EncodedInput], BatchEncoding]       ,  _snake_case			:					Optional[int] = None       ,  _snake_case			:					PaddingStrategy = PaddingStrategy.DO_NOT_PAD       ,  _snake_case			:					Optional[int] = None       ,  _snake_case			:					Optional[bool] = None       ,  )->dict:






			'''simple docstring'''

			__lowerCAmelCase					:	Tuple      =       super()._pad(
			    encoded_inputs=_snake_case       ,  max_length=_snake_case       ,  padding_strategy=_snake_case       ,  pad_to_multiple_of=_snake_case       ,  return_attention_mask=_snake_case       ,  )

			# Load from model defaults
			if return_attention_mask is None:
				__lowerCAmelCase					:	Union[str, Any]      =       """attention_mask""" in self.model_input_names

			if return_attention_mask and "global_attention_mask" in encoded_inputs:
				__lowerCAmelCase					:	Optional[int]      =       encoded_inputs[self.model_input_names[0]]
				# `global_attention_mask` need to have the same length as other (sequential) inputs.
				__lowerCAmelCase					:	Dict      =       len(encoded_inputs["""global_attention_mask"""]      ) != len(_snake_case      )

				if needs_to_be_padded:
					__lowerCAmelCase					:	Optional[int]      =       len(_snake_case      ) - len(encoded_inputs["""global_attention_mask"""]      )

					if self.padding_side == "right":
						# Use `-1` since `0` in `global_attention_mask` means `local attention` instead of `not to attend`
						__lowerCAmelCase					:	Any      =       (
						    encoded_inputs["""global_attention_mask"""] + [-1] * difference
						)
					elif self.padding_side == "left":
						__lowerCAmelCase					:	Union[str, Any]      =       [-1] * difference + encoded_inputs[
						    """global_attention_mask"""
						]
					else:
						raise ValueError("""Invalid padding strategy:""" + str(self.padding_side      )      )

			return encoded_inputs






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






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






   '''simple docstring'''

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

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






  def 				UpperCAmelCase__						(							self			:					Any      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	List[Any]      =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

   __lowerCAmelCase					:	Tuple      =       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=_snake_case       ,  initializer_range=self.initializer_range       ,  )

   return config, pixel_values






  def 				UpperCAmelCase__						(							self			:					List[Any]       ,  _snake_case			:					List[str]       ,  _snake_case			:					List[Any]      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	str      =       FlaxViTModel(config=_snake_case      )
   __lowerCAmelCase					:	int      =       model(_snake_case      )
   # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
   __lowerCAmelCase					:	Dict      =       (self.image_size, self.image_size)
   __lowerCAmelCase					:	Any      =       (self.patch_size, self.patch_size)
   __lowerCAmelCase					:	Any      =       (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			:					Dict       ,  _snake_case			:					str       ,  _snake_case			:					List[Any]      )->Union[str, Any]:






   '''simple docstring'''

   __lowerCAmelCase					:	Tuple      =       self.type_sequence_label_size
   __lowerCAmelCase					:	Tuple      =       FlaxViTForImageClassification(config=_snake_case      )
   __lowerCAmelCase					:	List[str]      =       model(_snake_case      )
   self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.type_sequence_label_size)      )

   # test greyscale images
   __lowerCAmelCase					:	str      =       1
   __lowerCAmelCase					:	Any      =       FlaxViTForImageClassification(_snake_case      )

   __lowerCAmelCase					:	int      =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
   __lowerCAmelCase					:	Dict      =       model(_snake_case      )






  def 				UpperCAmelCase__						(							self			:					str      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	Any      =       self.prepare_config_and_inputs()
   (
       (
       __lowerCAmelCase
   ) ,					(
       __lowerCAmelCase
   ) ,					
   )					:	Union[str, Any]      =       config_and_inputs
   __lowerCAmelCase					:	str      =       {"""pixel_values""": pixel_values}
   return config, inputs_dict






@require_flax
class        snake_case_    (		__lowercase					,unittest.TestCase						):
  A_      							=   (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()






  def 				UpperCAmelCase__						(							self			:					str      )->None:






   '''simple docstring'''

   __lowerCAmelCase					:	List[str]      =       FlaxViTModelTester(self      )
   __lowerCAmelCase					:	Optional[Any]      =       ConfigTester(self       ,  config_class=_snake_case       ,  has_text_modality=_snake_case       ,  hidden_size=37      )






  def 				UpperCAmelCase__						(							self			:					Any      )->Dict:






   '''simple docstring'''

   self.config_tester.run_common_tests()






  def 				UpperCAmelCase__						(							self			:					Any      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	Dict      =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*_snake_case      )






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






   '''simple docstring'''

   __lowerCAmelCase					:	Dict      =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*_snake_case      )






  def 				UpperCAmelCase__						(							self			:					int      )->int:






   '''simple docstring'''

   __lowerCAmelCase ,					__lowerCAmelCase					:	Any      =       self.model_tester.prepare_config_and_inputs_for_common()

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

    __lowerCAmelCase					:	str      =       ["""pixel_values"""]
    self.assertListEqual(arg_names[:1]       ,  _snake_case      )






  def 				UpperCAmelCase__						(							self			:					str      )->str:






   '''simple docstring'''

   __lowerCAmelCase ,					__lowerCAmelCase					:	Any      =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    with self.subTest(model_class.__name__      ):
     __lowerCAmelCase					:	Optional[int]      =       self._prepare_for_class(_snake_case       ,  _snake_case      )
     __lowerCAmelCase					:	List[str]      =       model_class(_snake_case      )

     @jax.jit
     def model_jitted(_snake_case			:					Dict       ,  **_snake_case			:					Union[str, Any]      ):
      return model(pixel_values=_snake_case       ,  **_snake_case      )

     with self.subTest("""JIT Enabled"""      ):
      __lowerCAmelCase					:	List[Any]      =       model_jitted(**_snake_case      ).to_tuple()

     with self.subTest("""JIT Disabled"""      ):
      with jax.disable_jit():
       __lowerCAmelCase					:	str      =       model_jitted(**_snake_case      ).to_tuple()

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






  @slow
  def 				UpperCAmelCase__						(							self			:					Dict      )->str:






   '''simple docstring'''

   for model_class_name in self.all_model_classes:
    __lowerCAmelCase					:	List[str]      =       model_class_name.from_pretrained("""google/vit-base-patch16-224"""      )
    __lowerCAmelCase					:	List[str]      =       model(np.ones((1, 3, 224, 224)      )      )
    self.assertIsNotNone(_snake_case      )