


from __future__ import annotations

import unittest

from transformers import is_tf_available, is_torch_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, is_pt_tf_cross_test, slow


if is_tf_available():
						from transformers import (
						    AutoConfig,
						    BertConfig,
						    GPTaConfig,
						    TaConfig,
						    TFAutoModel,
						    TFAutoModelForCausalLM,
						    TFAutoModelForMaskedLM,
						    TFAutoModelForPreTraining,
						    TFAutoModelForQuestionAnswering,
						    TFAutoModelForSeqaSeqLM,
						    TFAutoModelForSequenceClassification,
						    TFAutoModelWithLMHead,
						    TFBertForMaskedLM,
						    TFBertForPreTraining,
						    TFBertForQuestionAnswering,
						    TFBertForSequenceClassification,
						    TFBertModel,
						    TFGPTaLMHeadModel,
						    TFRobertaForMaskedLM,
						    TFTaForConditionalGeneration,
						)
						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

if is_torch_available():
						from transformers import (
						    AutoModel,
						    AutoModelForCausalLM,
						    AutoModelForMaskedLM,
						    AutoModelForPreTraining,
						    AutoModelForQuestionAnswering,
						    AutoModelForSeqaSeqLM,
						    AutoModelForSequenceClassification,
						    AutoModelWithLMHead,
						    BertForMaskedLM,
						    BertForPreTraining,
						    BertForQuestionAnswering,
						    BertForSequenceClassification,
						    BertModel,
						    GPTaLMHeadModel,
						    RobertaForMaskedLM,
						    TaForConditionalGeneration,
						)
@is_pt_tf_cross_test
class 			_lowercase (     unittest.TestCase			):






				'''simple docstring'''



				@slow
				def 						__magic_name__(			self     :Optional[Any]      )						->     Union[str, Any]:
										# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										for model_name in ["bert-base-uncased"]:
																__SCREAMING_SNAKE_CASE		:	List[Any]						     =						AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Tuple						     =						TFAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						AutoModel.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )


				@slow
				def 						__magic_name__(			self     :List[Any]      )						->     Optional[int]:
										# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										for model_name in ["bert-base-uncased"]:
																__SCREAMING_SNAKE_CASE		:	Any						     =						AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	int						     =						TFAutoModelForPreTraining.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						AutoModelForPreTraining.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )


				@slow
				def 						__magic_name__(			self     :Union[str, Any]      )						->     List[str]:
										for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	int						     =						TFAutoModelForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																__SCREAMING_SNAKE_CASE		:	str						     =						TFAutoModelForCausalLM.from_pretrained(
																    __SCREAMING_SNAKE_CASE     ,     output_loading_info=__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						AutoModelForCausalLM.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						AutoModelForCausalLM.from_pretrained(
																    __SCREAMING_SNAKE_CASE     ,     output_loading_info=__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )


				@slow
				def 						__magic_name__(			self     :Tuple      )						->     List[Any]:
										for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__SCREAMING_SNAKE_CASE		:	List[str]						     =						AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	List[str]						     =						TFAutoModelWithLMHead.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	List[str]						     =						AutoModelWithLMHead.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )


				@slow
				def 						__magic_name__(			self     :Optional[int]      )						->     Optional[Any]:
										for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__SCREAMING_SNAKE_CASE		:	Tuple						     =						AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	List[Any]						     =						TFAutoModelForMaskedLM.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						TFAutoModelForMaskedLM.from_pretrained(
																    __SCREAMING_SNAKE_CASE     ,     output_loading_info=__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Dict						     =						AutoModelForMaskedLM.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						AutoModelForMaskedLM.from_pretrained(
																    __SCREAMING_SNAKE_CASE     ,     output_loading_info=__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )


				@slow
				def 						__magic_name__(			self     :Optional[Any]      )						->     Tuple:
										for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__SCREAMING_SNAKE_CASE		:	Tuple						     =						AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Dict						     =						TFAutoModelForSeqaSeqLM.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																__SCREAMING_SNAKE_CASE		:	Dict						     =						TFAutoModelForSeqaSeqLM.from_pretrained(
																    __SCREAMING_SNAKE_CASE     ,     output_loading_info=__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Any						     =						AutoModelForSeqaSeqLM.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																__SCREAMING_SNAKE_CASE		:	Tuple						     =						AutoModelForSeqaSeqLM.from_pretrained(
																    __SCREAMING_SNAKE_CASE     ,     output_loading_info=__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )


				@slow
				def 						__magic_name__(			self     :Optional[Any]      )						->     Tuple:
										# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										for model_name in ["bert-base-uncased"]:
																__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Dict						     =						TFAutoModelForSequenceClassification.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	List[str]						     =						AutoModelForSequenceClassification.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )


				@slow
				def 						__magic_name__(			self     :Union[str, Any]      )						->     Tuple:
										# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										for model_name in ["bert-base-uncased"]:
																__SCREAMING_SNAKE_CASE		:	List[str]						     =						AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	str						     =						TFAutoModelForQuestionAnswering.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )

																__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						AutoModelForQuestionAnswering.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
																self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )
																self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )


				def 						__magic_name__(			self     :str      )						->     Union[str, Any]:
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						TFAutoModelWithLMHead.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
										self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )
										self.assertEqual(model.num_parameters()     ,     14_410      )
										self.assertEqual(model.num_parameters(only_trainable=__SCREAMING_SNAKE_CASE      )     ,     14_410      )

										__SCREAMING_SNAKE_CASE		:	Dict						     =						AutoModelWithLMHead.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
										self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )
										self.assertEqual(model.num_parameters()     ,     14_410      )
										self.assertEqual(model.num_parameters(only_trainable=__SCREAMING_SNAKE_CASE      )     ,     14_410      )



				def 						__magic_name__(			self     :Any      )						->     Dict:
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						TFAutoModelWithLMHead.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_pt=__SCREAMING_SNAKE_CASE      )
										self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )
										self.assertEqual(model.num_parameters()     ,     14_410      )
										self.assertEqual(model.num_parameters(only_trainable=__SCREAMING_SNAKE_CASE      )     ,     14_410      )

										__SCREAMING_SNAKE_CASE		:	List[str]						     =						AutoModelWithLMHead.from_pretrained(__SCREAMING_SNAKE_CASE     ,     from_tf=__SCREAMING_SNAKE_CASE      )
										self.assertIsInstance(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE      )
										self.assertEqual(model.num_parameters()     ,     14_410      )
										self.assertEqual(model.num_parameters(only_trainable=__SCREAMING_SNAKE_CASE      )     ,     14_410      )








'''simple docstring'''







import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__lowerCAmelCase									=   object()

# For specifying empty leaf dict `{}`
__lowerCAmelCase									=   object()





def 		_UpperCAmelCase (				__A				: List[str]    ,		__A				: Tuple   ):
				a_		:	List[Any]									=   tuple((re.compile(x + '''$'''   ) for x in qs)   )
				for i in range(len(__A   ) - len(__A   ) + 1   ):
								a_		:	Union[str, Any]									=   [x.match(__A   ) for x, y in zip(__A    ,		ks[i:]   )]
								if matches and all(__A   ):
												return True
				return False





def 		_UpperCAmelCase (				__A				: List[str]   ):
				def replace(__A				: int    ,		__A				: Union[str, Any]   ):
								for rule, replacement in rules:
												if _match(__A    ,		__A   ):
																return replacement
								return val

				return replace





def 		_UpperCAmelCase (				):
				return [
				    # embeddings
				    (("transformer", "wpe", "embedding"), P('''mp'''    ,		__A   )),
				    (("transformer", "wte", "embedding"), P('''mp'''    ,		__A   )),
				    # atention
				    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__A    ,		'''mp'''   )),
				    (("attention", "out_proj", "kernel"), P('''mp'''    ,		__A   )),
				    (("attention", "out_proj", "bias"), None),
				    # mlp
				    (("mlp", "c_fc", "kernel"), P(__A    ,		'''mp'''   )),
				    (("mlp", "c_fc", "bias"), P('''mp'''   )),
				    (("mlp", "c_proj", "kernel"), P('''mp'''    ,		__A   )),
				    (("mlp", "c_proj", "bias"), None),
				    # layer norms
				    ((r"ln_\d+", "bias"), None),
				    ((r"\d+", r"ln_\d+", "scale"), None),
				    (("ln_f", "bias"), None),
				    (("ln_f", "scale"), None),
				]





def 		_UpperCAmelCase (				__A				: Union[str, Any]   ):
				a_		:	Tuple									=   _get_partition_rules()
				a_		:	Tuple									=   _replacement_rules(__A   )
				a_		:	Optional[Any]									=   {k: _unmatched for k in flatten_dict(__A   )}
				a_		:	Optional[Any]									=   {k: replace(__A    ,		__A   ) for k, v in initd.items()}
				assert _unmatched not in result.values(), "Incomplete partition spec."
				return freeze(unflatten_dict(__A   )   )




