



'''simple docstring'''






from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)

@flax.struct.dataclass
class 				SCREAMING_SNAKE_CASE(							A__			):







	"""simple docstring"""






	lowerCamelCase__			 =       42

@flax_register_to_config
class 				SCREAMING_SNAKE_CASE(							nn.Module							,		A__							,		A__			):







	"""simple docstring"""






	lowerCamelCase__			 =       32
	lowerCamelCase__			 =       4
	lowerCamelCase__			 =       4
	lowerCamelCase__			 =       (
	    "CrossAttnDownBlock2D",
	    "CrossAttnDownBlock2D",
	    "CrossAttnDownBlock2D",
	    "DownBlock2D",
	)
	lowerCamelCase__			 =       ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
	lowerCamelCase__			 =       False
	lowerCamelCase__			 =       (320, 640, 1_280, 1_280)
	lowerCamelCase__			 =       2
	lowerCamelCase__			 =       8
	lowerCamelCase__			 =       None
	lowerCamelCase__			 =       1_280
	lowerCamelCase__			 =       0.0
	lowerCamelCase__			 =       False
	lowerCamelCase__			 =       jnp.floataa
	lowerCamelCase__			 =       True
	lowerCamelCase__			 =       0
	lowerCamelCase__			 =       False







	def      A	(						self			:     List[Any]		,			__snake_case			:     jax.random.KeyArray		)			->		FrozenDict:
								# init input tensors
								UpperCAmelCase :       List[Any] 							=		(1, self.in_channels, self.sample_size, self.sample_size)
								UpperCAmelCase :       int 							=		jnp.zeros(__snake_case		,			dtype=jnp.floataa		)
								UpperCAmelCase :       List[Any] 							=		jnp.ones((1,)		,			dtype=jnp.intaa		)
								UpperCAmelCase :       List[Any] 							=		jnp.zeros((1, 1, self.cross_attention_dim)		,			dtype=jnp.floataa		)

								UpperCAmelCase							,		UpperCAmelCase :       Optional[Any] 							=		jax.random.split(__snake_case		)
								UpperCAmelCase :       Dict 							=		{'''params''': params_rng, '''dropout''': dropout_rng}

								return self.init(__snake_case		,			__snake_case		,			__snake_case		,			__snake_case		)["params"]







	def      A	(						self			:     Tuple		)			->		int:
								UpperCAmelCase :       List[str] 							=		self.block_out_channels
								UpperCAmelCase :       Optional[Any] 							=		block_out_channels[0] * 4

								if self.num_attention_heads is not None:
															raise ValueError(
															    '''At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'''		)

								# If `num_attention_heads` is not defined (which is the case for most models)
								# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
								# The reason for this behavior is to correct for incorrectly named variables that were introduced
								# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
								# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
								# which is why we correct for the naming here.
								UpperCAmelCase :       Union[str, Any] 							=		self.num_attention_heads or self.attention_head_dim

								# input
								UpperCAmelCase :       Union[str, Any] 							=		nn.Conv(
								    block_out_channels[0]		,			kernel_size=(3, 3)		,			strides=(1, 1)		,			padding=((1, 1), (1, 1))		,			dtype=self.dtype		,			)

								# time
								UpperCAmelCase :       Optional[Any] 							=		FlaxTimesteps(
								    block_out_channels[0]		,			flip_sin_to_cos=self.flip_sin_to_cos		,			freq_shift=self.config.freq_shift		)
								UpperCAmelCase :       Optional[Any] 							=		FlaxTimestepEmbedding(__snake_case		,			dtype=self.dtype		)

								UpperCAmelCase :       Tuple 							=		self.only_cross_attention
								if isinstance(__snake_case		,			__snake_case		):
															UpperCAmelCase :       Optional[int] 							=		(only_cross_attention,) * len(self.down_block_types		)

								if isinstance(__snake_case		,			__snake_case		):
															UpperCAmelCase :       Optional[int] 							=		(num_attention_heads,) * len(self.down_block_types		)

								# down
								UpperCAmelCase :       Optional[Any] 							=		[]
								UpperCAmelCase :       str 							=		block_out_channels[0]
								for i, down_block_type in enumerate(self.down_block_types		):
															UpperCAmelCase :       Dict 							=		output_channel
															UpperCAmelCase :       Any 							=		block_out_channels[i]
															UpperCAmelCase :       Optional[int] 							=		i == len(__snake_case		) - 1

															if down_block_type == "CrossAttnDownBlock2D":
																						UpperCAmelCase :       Tuple 							=		FlaxCrossAttnDownBlockaD(
																						    in_channels=__snake_case		,			out_channels=__snake_case		,			dropout=self.dropout		,			num_layers=self.layers_per_block		,			num_attention_heads=num_attention_heads[i]		,			add_downsample=not is_final_block		,			use_linear_projection=self.use_linear_projection		,			only_cross_attention=only_cross_attention[i]		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)
															else:
																						UpperCAmelCase :       Any 							=		FlaxDownBlockaD(
																						    in_channels=__snake_case		,			out_channels=__snake_case		,			dropout=self.dropout		,			num_layers=self.layers_per_block		,			add_downsample=not is_final_block		,			dtype=self.dtype		,			)

															down_blocks.append(__snake_case		)
								UpperCAmelCase :       List[Any] 							=		down_blocks

								# mid
								UpperCAmelCase :       Tuple 							=		FlaxUNetMidBlockaDCrossAttn(
								    in_channels=block_out_channels[-1]		,			dropout=self.dropout		,			num_attention_heads=num_attention_heads[-1]		,			use_linear_projection=self.use_linear_projection		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)

								# up
								UpperCAmelCase :       Dict 							=		[]
								UpperCAmelCase :       List[Any] 							=		list(reversed(__snake_case		)		)
								UpperCAmelCase :       Optional[int] 							=		list(reversed(__snake_case		)		)
								UpperCAmelCase :       List[Any] 							=		list(reversed(__snake_case		)		)
								UpperCAmelCase :       List[Any] 							=		reversed_block_out_channels[0]
								for i, up_block_type in enumerate(self.up_block_types		):
															UpperCAmelCase :       int 							=		output_channel
															UpperCAmelCase :       Any 							=		reversed_block_out_channels[i]
															UpperCAmelCase :       List[Any] 							=		reversed_block_out_channels[min(i + 1		,			len(__snake_case		) - 1		)]

															UpperCAmelCase :       Tuple 							=		i == len(__snake_case		) - 1

															if up_block_type == "CrossAttnUpBlock2D":
																						UpperCAmelCase :       List[str] 							=		FlaxCrossAttnUpBlockaD(
																						    in_channels=__snake_case		,			out_channels=__snake_case		,			prev_output_channel=__snake_case		,			num_layers=self.layers_per_block + 1		,			num_attention_heads=reversed_num_attention_heads[i]		,			add_upsample=not is_final_block		,			dropout=self.dropout		,			use_linear_projection=self.use_linear_projection		,			only_cross_attention=only_cross_attention[i]		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)
															else:
																						UpperCAmelCase :       List[str] 							=		FlaxUpBlockaD(
																						    in_channels=__snake_case		,			out_channels=__snake_case		,			prev_output_channel=__snake_case		,			num_layers=self.layers_per_block + 1		,			add_upsample=not is_final_block		,			dropout=self.dropout		,			dtype=self.dtype		,			)

															up_blocks.append(__snake_case		)
															UpperCAmelCase :       int 							=		output_channel
								UpperCAmelCase :       str 							=		up_blocks

								# out
								UpperCAmelCase :       str 							=		nn.GroupNorm(num_groups=32		,			epsilon=1E-5		)
								UpperCAmelCase :       Union[str, Any] 							=		nn.Conv(
								    self.out_channels		,			kernel_size=(3, 3)		,			strides=(1, 1)		,			padding=((1, 1), (1, 1))		,			dtype=self.dtype		,			)







	def __call__(						self			:     Tuple		,			__snake_case			:     Tuple		,			__snake_case			:     List[str]		,			__snake_case			:     str		,			__snake_case			:     int=None		,			__snake_case			:     Optional[int]=None		,			__snake_case			:     bool = True		,			__snake_case			:     bool = False		,			)			->		Union[FlaxUNetaDConditionOutput, Tuple]:
								# 1. time
								if not isinstance(__snake_case		,			jnp.ndarray		):
															UpperCAmelCase :       List[str] 							=		jnp.array([timesteps]		,			dtype=jnp.intaa		)
								elif isinstance(__snake_case		,			jnp.ndarray		) and len(timesteps.shape		) == 0:
															UpperCAmelCase :       Union[str, Any] 							=		timesteps.astype(dtype=jnp.floataa		)
															UpperCAmelCase :       Any 							=		jnp.expand_dims(__snake_case		,			0		)

								UpperCAmelCase :       List[Any] 							=		self.time_proj(__snake_case		)
								UpperCAmelCase :       List[str] 							=		self.time_embedding(__snake_case		)

								# 2. pre-process
								UpperCAmelCase :       List[str] 							=		jnp.transpose(__snake_case		,			(0, 2, 3, 1)		)
								UpperCAmelCase :       List[str] 							=		self.conv_in(__snake_case		)

								# 3. down
								UpperCAmelCase :       List[Any] 							=		(sample,)
								for down_block in self.down_blocks:
															if isinstance(__snake_case		,			__snake_case		):
																						UpperCAmelCase							,		UpperCAmelCase :       List[Any] 							=		down_block(__snake_case		,			__snake_case		,			__snake_case		,			deterministic=not train		)
															else:
																						UpperCAmelCase							,		UpperCAmelCase :       Optional[Any] 							=		down_block(__snake_case		,			__snake_case		,			deterministic=not train		)
															down_block_res_samples += res_samples

								if down_block_additional_residuals is not None:
															UpperCAmelCase :       List[str] 							=		()

															for down_block_res_sample, down_block_additional_residual in zip(
															    __snake_case		,			__snake_case		):
																						down_block_res_sample += down_block_additional_residual
																						new_down_block_res_samples += (down_block_res_sample,)

															UpperCAmelCase :       str 							=		new_down_block_res_samples

								# 4. mid
								UpperCAmelCase :       Any 							=		self.mid_block(__snake_case		,			__snake_case		,			__snake_case		,			deterministic=not train		)

								if mid_block_additional_residual is not None:
															sample += mid_block_additional_residual

								# 5. up
								for up_block in self.up_blocks:
															UpperCAmelCase :       Tuple 							=		down_block_res_samples[-(self.layers_per_block + 1) :]
															UpperCAmelCase :       int 							=		down_block_res_samples[: -(self.layers_per_block + 1)]
															if isinstance(__snake_case		,			__snake_case		):
																						UpperCAmelCase :       Tuple 							=		up_block(
																						    __snake_case		,			temb=__snake_case		,			encoder_hidden_states=__snake_case		,			res_hidden_states_tuple=__snake_case		,			deterministic=not train		,			)
															else:
																						UpperCAmelCase :       List[Any] 							=		up_block(__snake_case		,			temb=__snake_case		,			res_hidden_states_tuple=__snake_case		,			deterministic=not train		)

        # 6. post-process
								UpperCAmelCase :       Optional[Any] 							=		self.conv_norm_out(__snake_case		)
								UpperCAmelCase :       str 							=		nn.silu(__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		self.conv_out(__snake_case		)
								UpperCAmelCase :       Dict 							=		jnp.transpose(__snake_case		,			(0, 3, 1, 2)		)

								if not return_dict:
															return (sample,)

								return FlaxUNetaDConditionOutput(sample=__snake_case		)











"""simple docstring"""






from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
			from transformers import (
			    TFAutoModel,
			    TFAutoModelForCausalLM,
			    TFAutoModelForMaskedLM,
			    TFAutoModelForPreTraining,
			    TFAutoModelForQuestionAnswering,
			    TFAutoModelForSeqaSeqLM,
			    TFAutoModelForSequenceClassification,
			    TFAutoModelForTableQuestionAnswering,
			    TFAutoModelForTokenClassification,
			    TFAutoModelWithLMHead,
			    TFBertForMaskedLM,
			    TFBertForPreTraining,
			    TFBertForQuestionAnswering,
			    TFBertForSequenceClassification,
			    TFBertModel,
			    TFFunnelBaseModel,
			    TFFunnelModel,
			    TFGPTaLMHeadModel,
			    TFRobertaForMaskedLM,
			    TFTaForConditionalGeneration,
			    TFTapasForQuestionAnswering,
			)
			from transformers.models.auto.modeling_tf_auto import (
			    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
			    TF_MODEL_FOR_MASKED_LM_MAPPING,
			    TF_MODEL_FOR_PRETRAINING_MAPPING,
			    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
			    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
			    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
			    TF_MODEL_MAPPING,
			)
			from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
			from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
			from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
			from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST






class lowerCAmelCase__	(     lowercase						):






			'''simple docstring'''



			lowerCamelCase__      =			"""new-model"""


if is_tf_available():






			class lowerCAmelCase__	(     lowercase						):






						'''simple docstring'''



						lowerCamelCase__      =			NewModelConfig






@require_tf
class lowerCAmelCase__	(     unittest.TestCase						):






			'''simple docstring'''



			@slow
			def 					A_    (  self  ):
							_lowerCamelCase     :			List[str]    =    'bert-base-cased'
							_lowerCamelCase     :			Union[str, Any]    =    AutoConfig.from_pretrained(lowercase  )
							self.assertIsNotNone(lowercase  )
							self.assertIsInstance(lowercase  ,      lowercase  )

							_lowerCamelCase     :			Union[str, Any]    =    TFAutoModel.from_pretrained(lowercase  )
							self.assertIsNotNone(lowercase  )
							self.assertIsInstance(lowercase  ,      lowercase  )
			@slow
			def 					A_    (  self  ):
							_lowerCamelCase     :			List[str]    =    'bert-base-cased'
							_lowerCamelCase     :			int    =    AutoConfig.from_pretrained(lowercase  )
							self.assertIsNotNone(lowercase  )
							self.assertIsInstance(lowercase  ,      lowercase  )

							_lowerCamelCase     :			int    =    TFAutoModelForPreTraining.from_pretrained(lowercase  )
							self.assertIsNotNone(lowercase  )
							self.assertIsInstance(lowercase  ,      lowercase  )
			@slow
			def 					A_    (  self  ):
							for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											_lowerCamelCase     :			Dict    =    AutoConfig.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )

											_lowerCamelCase     :			int    =    TFAutoModelForCausalLM.from_pretrained(lowercase  )
											_lowerCamelCase, _lowerCamelCase     :			str    =    TFAutoModelForCausalLM.from_pretrained(lowercase  ,      output_loading_info=lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )
			@slow
			def 					A_    (  self  ):
							for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											_lowerCamelCase     :			List[Any]    =    AutoConfig.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )

											_lowerCamelCase     :			str    =    TFAutoModelWithLMHead.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )
			@slow
			def 					A_    (  self  ):
							for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											_lowerCamelCase     :			Tuple    =    AutoConfig.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )

											_lowerCamelCase     :			Optional[int]    =    TFAutoModelForMaskedLM.from_pretrained(lowercase  )
											_lowerCamelCase, _lowerCamelCase     :			Tuple    =    TFAutoModelForMaskedLM.from_pretrained(lowercase  ,      output_loading_info=lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )
			@slow
			def 					A_    (  self  ):
							for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											_lowerCamelCase     :			Optional[int]    =    AutoConfig.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )

											_lowerCamelCase     :			Tuple    =    TFAutoModelForSeqaSeqLM.from_pretrained(lowercase  )
											_lowerCamelCase, _lowerCamelCase     :			Tuple    =    TFAutoModelForSeqaSeqLM.from_pretrained(lowercase  ,      output_loading_info=lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )
			@slow
			def 					A_    (  self  ):
							# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							for model_name in ["bert-base-uncased"]:
											_lowerCamelCase     :			str    =    AutoConfig.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )

											_lowerCamelCase     :			Union[str, Any]    =    TFAutoModelForSequenceClassification.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )
			@slow
			def 					A_    (  self  ):
							# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							for model_name in ["bert-base-uncased"]:
											_lowerCamelCase     :			Optional[Any]    =    AutoConfig.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )

											_lowerCamelCase     :			List[str]    =    TFAutoModelForQuestionAnswering.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )
			@slow
			@require_tensorflow_probability
			def 					A_    (  self  ):
							for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
											_lowerCamelCase     :			Dict    =    AutoConfig.from_pretrained(lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )

											_lowerCamelCase     :			Dict    =    TFAutoModelForTableQuestionAnswering.from_pretrained(lowercase  )
											_lowerCamelCase, _lowerCamelCase     :			List[Any]    =    TFAutoModelForTableQuestionAnswering.from_pretrained(
											    lowercase  ,      output_loading_info=lowercase  )
											self.assertIsNotNone(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )
			def 					A_    (  self  ):
							_lowerCamelCase     :			int    =    TFAutoModelWithLMHead.from_pretrained(lowercase  )
							self.assertIsInstance(lowercase  ,      lowercase  )
							self.assertEqual(model.num_parameters()  ,      14410  )
							self.assertEqual(model.num_parameters(only_trainable=lowercase  )  ,      14410  )
			def 					A_    (  self  ):
							_lowerCamelCase     :			Any    =    TFAutoModelWithLMHead.from_pretrained(lowercase  )
							self.assertIsInstance(lowercase  ,      lowercase  )
							self.assertEqual(model.num_parameters()  ,      14410  )
							self.assertEqual(model.num_parameters(only_trainable=lowercase  )  ,      14410  )
			def 					A_    (  self  ):
							# For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
							_lowerCamelCase     :			List[str]    =    TFAutoModel.from_pretrained('sgugger/funnel-random-tiny'  )
							self.assertIsInstance(lowercase  ,      lowercase  )

							_lowerCamelCase     :			Optional[int]    =    copy.deepcopy(model.config  )
							_lowerCamelCase     :			Dict    =    ['FunnelBaseModel']
							_lowerCamelCase     :			List[Any]    =    TFAutoModel.from_config(lowercase  )
							self.assertIsInstance(lowercase  ,      lowercase  )

							with tempfile.TemporaryDirectory() as tmp_dir:
											model.save_pretrained(lowercase  )
											_lowerCamelCase     :			Tuple    =    TFAutoModel.from_pretrained(lowercase  )
											self.assertIsInstance(lowercase  ,      lowercase  )
			def 					A_    (  self  ):
							try:
											AutoConfig.register('new-model'  ,      lowercase  )

											_lowerCamelCase     :			Tuple    =    [
											    TFAutoModel,
											    TFAutoModelForCausalLM,
											    TFAutoModelForMaskedLM,
											    TFAutoModelForPreTraining,
											    TFAutoModelForQuestionAnswering,
											    TFAutoModelForSequenceClassification,
											    TFAutoModelForTokenClassification,
											]

											for auto_class in auto_classes:
															with self.subTest(auto_class.__name__  ):
																			# Wrong config class will raise an error
																			with self.assertRaises(lowercase  ):
																							auto_class.register(lowercase  ,      lowercase  )
																			auto_class.register(lowercase  ,      lowercase  )
																			# Trying to register something existing in the Transformers library will raise an error
																			with self.assertRaises(lowercase  ):
																							auto_class.register(lowercase  ,      lowercase  )

																			# Now that the config is registered, it can be used as any other config with the auto-API
																			_lowerCamelCase     :			Optional[Any]    =    BertModelTester(self  ).get_config()
																			_lowerCamelCase     :			Dict    =    NewModelConfig(**tiny_config.to_dict()  )
																			_lowerCamelCase     :			int    =    auto_class.from_config(lowercase  )
																			self.assertIsInstance(lowercase  ,      lowercase  )

																			with tempfile.TemporaryDirectory() as tmp_dir:
																							model.save_pretrained(lowercase  )
																							_lowerCamelCase     :			List[Any]    =    auto_class.from_pretrained(lowercase  )
																							self.assertIsInstance(lowercase  ,      lowercase  )

							finally:
											if "new-model" in CONFIG_MAPPING._extra_content:
															del CONFIG_MAPPING._extra_content["new-model"]
											for mapping in (
											    TF_MODEL_MAPPING,
											    TF_MODEL_FOR_PRETRAINING_MAPPING,
											    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
											    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
											    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
											    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
											    TF_MODEL_FOR_MASKED_LM_MAPPING,
											):
															if NewModelConfig in mapping._extra_content:
																			del mapping._extra_content[NewModelConfig]
			def 					A_    (  self  ):
							with self.assertRaisesRegex(
							    lowercase  ,      'bert-base is not a local folder and is not a valid model identifier'  ):
											_lowerCamelCase     :			Union[str, Any]    =    TFAutoModel.from_pretrained('bert-base'  )
			def 					A_    (  self  ):
							with self.assertRaisesRegex(
							    lowercase  ,      r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'  ):
											_lowerCamelCase     :			str    =    TFAutoModel.from_pretrained(lowercase  ,      revision='aaaaaa'  )
			def 					A_    (  self  ):
							with self.assertRaisesRegex(
							    lowercase  ,      'hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin'  ,      ):
											_lowerCamelCase     :			Tuple    =    TFAutoModel.from_pretrained('hf-internal-testing/config-no-model'  )
			def 					A_    (  self  ):
							with self.assertRaisesRegex(lowercase  ,      'Use `from_pt=True` to load this model'  ):
											_lowerCamelCase     :			Tuple    =    TFAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'  )






			def 					A_    (  self  ):
							# Make sure we have cached the model.
							_lowerCamelCase     :			Optional[int]    =    TFAutoModel.from_pretrained('hf-internal-testing/tiny-random-bert'  )
							with RequestCounter() as counter:
											_lowerCamelCase     :			Optional[int]    =    TFAutoModel.from_pretrained('hf-internal-testing/tiny-random-bert'  )
											self.assertEqual(counter.get_request_count  ,      0  )
											self.assertEqual(counter.head_request_count  ,      1  )
											self.assertEqual(counter.other_request_count  ,      0  )

							# With a sharded checkpoint
							_lowerCamelCase     :			int    =    TFAutoModel.from_pretrained('ArthurZ/tiny-random-bert-sharded'  )
							with RequestCounter() as counter:
											_lowerCamelCase     :			List[Any]    =    TFAutoModel.from_pretrained('ArthurZ/tiny-random-bert-sharded'  )
											self.assertEqual(counter.get_request_count  ,      0  )
											self.assertEqual(counter.head_request_count  ,      1  )
											self.assertEqual(counter.other_request_count  ,      0  )