



'''simple docstring'''






import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel

class 				SCREAMING_SNAKE_CASE:







	"""simple docstring"""






	def __init__(						self			:     int		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Optional[int]=13		,			__snake_case			:     str=7		,			__snake_case			:     int=True		,			__snake_case			:     List[Any]=True		,			__snake_case			:     Any=False		,			__snake_case			:     Union[str, Any]=True		,			__snake_case			:     Dict=99		,			__snake_case			:     Optional[int]=32		,			__snake_case			:     Any=5		,			__snake_case			:     Any=4		,			__snake_case			:     List[Any]=37		,			__snake_case			:     Dict="gelu"		,			__snake_case			:     Optional[int]=0.1		,			__snake_case			:     int=0.1		,			__snake_case			:     Union[str, Any]=512		,			__snake_case			:     Optional[int]=16		,			__snake_case			:     Any=2		,			__snake_case			:     Optional[Any]=0.02		,			__snake_case			:     List[str]=3		,			__snake_case			:     str=4		,			__snake_case			:     Optional[int]=None		,			)			->		Optional[int]:
								UpperCAmelCase :       Dict 							=		parent
								UpperCAmelCase :       List[str] 							=		batch_size
								UpperCAmelCase :       List[str] 							=		seq_length
								UpperCAmelCase :       Optional[Any] 							=		is_training
								UpperCAmelCase :       int 							=		use_input_mask
								UpperCAmelCase :       Any 							=		use_token_type_ids
								UpperCAmelCase :       Optional[Any] 							=		use_labels
								UpperCAmelCase :       Dict 							=		vocab_size
								UpperCAmelCase :       List[str] 							=		hidden_size
								UpperCAmelCase :       Any 							=		num_hidden_layers
								UpperCAmelCase :       Union[str, Any] 							=		num_attention_heads
								UpperCAmelCase :       int 							=		intermediate_size
								UpperCAmelCase :       Tuple 							=		hidden_act
								UpperCAmelCase :       str 							=		hidden_dropout_prob
								UpperCAmelCase :       int 							=		attention_probs_dropout_prob
								UpperCAmelCase :       Optional[int] 							=		max_position_embeddings
								UpperCAmelCase :       Dict 							=		type_vocab_size
								UpperCAmelCase :       Optional[int] 							=		type_sequence_label_size
								UpperCAmelCase :       Any 							=		initializer_range
								UpperCAmelCase :       int 							=		num_labels
								UpperCAmelCase :       Tuple 							=		num_choices
								UpperCAmelCase :       List[str] 							=		scope







	def      A	(						self			:     Any		)			->		str:
								UpperCAmelCase :       Tuple 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size		)

								UpperCAmelCase :       Optional[int] 							=		None
								if self.use_input_mask:
															UpperCAmelCase :       str 							=		random_attention_mask([self.batch_size, self.seq_length]		)

								UpperCAmelCase :       Optional[int] 							=		None
								if self.use_token_type_ids:
															UpperCAmelCase :       str 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.type_vocab_size		)

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

								UpperCAmelCase :       Dict 							=		self.get_config()

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







	def      A	(						self			:     Optional[Any]		)			->		Optional[int]:
								return OpenLlamaConfig(
								    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		,			is_decoder=__snake_case		,			initializer_range=self.initializer_range		,			use_stable_embedding=__snake_case		,			)







	def      A	(						self			:     Dict		,			__snake_case			:     Any		,			__snake_case			:     int		,			__snake_case			:     Dict		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Dict		,			__snake_case			:     Any		,			__snake_case			:     int		)			->		int:
								UpperCAmelCase :       Optional[Any] 							=		OpenLlamaModel(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Optional[int] 							=		model(__snake_case		,			attention_mask=__snake_case		)
								UpperCAmelCase :       Dict 							=		model(__snake_case		)
								self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)		)







	def      A	(						self			:     Union[str, Any]		,			__snake_case			:     int		,			__snake_case			:     int		,			__snake_case			:     Dict		,			__snake_case			:     Optional[int]		,			__snake_case			:     List[Any]		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Dict		,			__snake_case			:     Optional[int]		,			__snake_case			:     str		,			)			->		Tuple:
								UpperCAmelCase :       Optional[int] 							=		True
								UpperCAmelCase :       Optional[Any] 							=		OpenLlamaModel(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Union[str, Any] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			)
								UpperCAmelCase :       List[str] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			)
								UpperCAmelCase :       Optional[int] 							=		model(__snake_case		,			attention_mask=__snake_case		)
								self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)		)







	def      A	(						self			:     Dict		,			__snake_case			:     int		,			__snake_case			:     Optional[Any]		,			__snake_case			:     str		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Dict		,			__snake_case			:     int		,			__snake_case			:     Dict		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Dict		,			)			->		Optional[int]:
								UpperCAmelCase :       int 							=		OpenLlamaForCausalLM(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Dict 							=		model(__snake_case		,			attention_mask=__snake_case		,			labels=__snake_case		)
								self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length, self.vocab_size)		)







	def      A	(						self			:     Optional[Any]		,			__snake_case			:     Any		,			__snake_case			:     List[str]		,			__snake_case			:     List[Any]		,			__snake_case			:     str		,			__snake_case			:     List[str]		,			__snake_case			:     Dict		,			__snake_case			:     Optional[int]		,			__snake_case			:     Any		,			__snake_case			:     Optional[int]		,			)			->		List[Any]:
								UpperCAmelCase :       List[str] 							=		True
								UpperCAmelCase :       Union[str, Any] 							=		True
								UpperCAmelCase :       Union[str, Any] 							=		OpenLlamaForCausalLM(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()

								# first forward pass
								UpperCAmelCase :       Any 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			use_cache=__snake_case		,			)
								UpperCAmelCase :       Optional[int] 							=		outputs.past_key_values

								# create hypothetical multiple next token and extent to next_input_ids
								UpperCAmelCase :       str 							=		ids_tensor((self.batch_size, 3)		,			config.vocab_size		)
								UpperCAmelCase :       Optional[Any] 							=		ids_tensor((self.batch_size, 3)		,			vocab_size=2		)

								# append to next input_ids and
								UpperCAmelCase :       List[str] 							=		torch.cat([input_ids, next_tokens]		,			dim=-1		)
								UpperCAmelCase :       Tuple 							=		torch.cat([input_mask, next_mask]		,			dim=-1		)

								UpperCAmelCase :       Dict 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			output_hidden_states=__snake_case		,			)['''hidden_states'''][0]
								UpperCAmelCase :       Union[str, Any] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			past_key_values=__snake_case		,			output_hidden_states=__snake_case		,			)['''hidden_states'''][0]

								# select random slice
								UpperCAmelCase :       Union[str, Any] 							=		ids_tensor((1,)		,			output_from_past.shape[-1]		).item()
								UpperCAmelCase :       List[Any] 							=		output_from_no_past[:, -3:, random_slice_idx].detach()
								UpperCAmelCase :       Tuple 							=		output_from_past[:, :, random_slice_idx].detach()

								self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

								# test that outputs are equal for slice
								self.parent.assertTrue(torch.allclose(__snake_case		,			__snake_case		,			atol=1E-3		)		)







	def      A	(						self			:     Dict		)			->		Optional[Any]:
								UpperCAmelCase :       Optional[int] 							=		self.prepare_config_and_inputs()
								(
								    (
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		
								) :       List[str] 							=		config_and_inputs
								UpperCAmelCase :       Optional[int] 							=		{'''input_ids''': input_ids, '''attention_mask''': input_mask}
								return config, inputs_dict

@require_torch
class 				SCREAMING_SNAKE_CASE(							A__							,		A__							,		A__							,		unittest.TestCase			):







	"""simple docstring"""






	lowerCamelCase__			 =       (
	    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
	)
	lowerCamelCase__			 =       (OpenLlamaForCausalLM,) if is_torch_available() else ()
	lowerCamelCase__			 =       (
	    {
	        """feature-extraction""": OpenLlamaModel,
	        """text-classification""": OpenLlamaForSequenceClassification,
	        """text-generation""": OpenLlamaForCausalLM,
	        """zero-shot""": OpenLlamaForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	lowerCamelCase__			 =       False
	lowerCamelCase__			 =       False







	def      A	(						self			:     Union[str, Any]		)			->		int:
								UpperCAmelCase :       Optional[int] 							=		OpenLlamaModelTester(self		)
								UpperCAmelCase :       Any 							=		ConfigTester(self		,			config_class=__snake_case		,			hidden_size=37		)







	def      A	(						self			:     Dict		)			->		Union[str, Any]:
								self.config_tester.run_common_tests()







	def      A	(						self			:     Dict		)			->		str:
								UpperCAmelCase :       str 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__snake_case		)







	def      A	(						self			:     Union[str, Any]		)			->		Optional[int]:
								UpperCAmelCase :       int 							=		self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
															UpperCAmelCase :       Union[str, Any] 							=		type
															self.model_tester.create_and_check_model(*__snake_case		)







	def      A	(						self			:     Dict		)			->		Union[str, Any]:
								UpperCAmelCase							,		UpperCAmelCase :       Union[str, Any] 							=		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase :       Tuple 							=		3
								UpperCAmelCase :       Optional[int] 							=		input_dict['''input_ids''']
								UpperCAmelCase :       Optional[Any] 							=		input_ids.ne(1		).to(__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		ids_tensor([self.model_tester.batch_size]		,			self.model_tester.type_sequence_label_size		)
								UpperCAmelCase :       Any 							=		OpenLlamaForSequenceClassification(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Any 							=		model(__snake_case		,			attention_mask=__snake_case		,			labels=__snake_case		)
								self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels)		)







	def      A	(						self			:     List[str]		)			->		str:
								UpperCAmelCase							,		UpperCAmelCase :       List[Any] 							=		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase :       List[str] 							=		3
								UpperCAmelCase :       Dict 							=		'''single_label_classification'''
								UpperCAmelCase :       int 							=		input_dict['''input_ids''']
								UpperCAmelCase :       List[Any] 							=		input_ids.ne(1		).to(__snake_case		)
								UpperCAmelCase :       str 							=		ids_tensor([self.model_tester.batch_size]		,			self.model_tester.type_sequence_label_size		)
								UpperCAmelCase :       Optional[Any] 							=		OpenLlamaForSequenceClassification(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       List[Any] 							=		model(__snake_case		,			attention_mask=__snake_case		,			labels=__snake_case		)
								self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels)		)







	def      A	(						self			:     Union[str, Any]		)			->		Optional[int]:
								UpperCAmelCase							,		UpperCAmelCase :       Tuple 							=		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase :       str 							=		3
								UpperCAmelCase :       List[str] 							=		'''multi_label_classification'''
								UpperCAmelCase :       Union[str, Any] 							=		input_dict['''input_ids''']
								UpperCAmelCase :       Optional[int] 							=		input_ids.ne(1		).to(__snake_case		)
								UpperCAmelCase :       Union[str, Any] 							=		ids_tensor(
								    [self.model_tester.batch_size, config.num_labels]		,			self.model_tester.type_sequence_label_size		).to(torch.float		)
								UpperCAmelCase :       Any 							=		OpenLlamaForSequenceClassification(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Tuple 							=		model(__snake_case		,			attention_mask=__snake_case		,			labels=__snake_case		)
								self.assertEqual(result.logits.shape		,			(self.model_tester.batch_size, self.model_tester.num_labels)		)







	@unittest.skip('''Open-Llama buffers include complex numbers, which breaks this test'''		)
	def      A	(						self			:     List[str]		)			->		Optional[int]:
								pass







	@parameterized.expand([('''linear''',), ('''dynamic''',)]		)
	def      A	(						self			:     Any		,			__snake_case			:     Any		)			->		List[Any]:
								UpperCAmelCase							,		UpperCAmelCase :       str 							=		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase :       Optional[Any] 							=		ids_tensor([1, 10]		,			config.vocab_size		)
								UpperCAmelCase :       Tuple 							=		ids_tensor([1, int(config.max_position_embeddings * 1.5		)]		,			config.vocab_size		)

								set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
								UpperCAmelCase :       List[str] 							=		OpenLlamaModel(__snake_case		)
								original_model.to(__snake_case		)
								original_model.eval()
								UpperCAmelCase :       int 							=		original_model(__snake_case		).last_hidden_state
								UpperCAmelCase :       Optional[int] 							=		original_model(__snake_case		).last_hidden_state

								set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
								UpperCAmelCase :       str 							=		{'''type''': scaling_type, '''factor''': 10.0}
								UpperCAmelCase :       List[str] 							=		OpenLlamaModel(__snake_case		)
								scaled_model.to(__snake_case		)
								scaled_model.eval()
								UpperCAmelCase :       List[str] 							=		scaled_model(__snake_case		).last_hidden_state
								UpperCAmelCase :       List[str] 							=		scaled_model(__snake_case		).last_hidden_state

								# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
								# maximum sequence length, so the outputs for the short input should match.
								if scaling_type == "dynamic":
															self.assertTrue(torch.allclose(__snake_case		,			__snake_case		,			atol=1E-5		)		)
								else:
															self.assertFalse(torch.allclose(__snake_case		,			__snake_case		,			atol=1E-5		)		)

								# The output should be different for long inputs
								self.assertFalse(torch.allclose(__snake_case		,			__snake_case		,			atol=1E-5		)		)










'''simple docstring'''






import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
						import os

						from flax.core.frozen_dict import unfreeze
						from flax.traverse_util import flatten_dict

						from transformers import FlaxBertModel

						UpperCamelCase__:   Tuple								=			"0.12"  # assumed parallelism: 8

@require_flax
@is_staging_test
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	@classmethod
	def      A	(						cls			:     Union[str, Any]		)			->		int:
								UpperCAmelCase :       Optional[Any] 							=		TOKEN
								HfFolder.save_token(__snake_case		)







	@classmethod
	def      A	(						cls			:     List[str]		)			->		Tuple:
								try:
															delete_repo(token=cls._token		,			repo_id='''test-model-flax'''		)
								except HTTPError:
															pass

								try:
															delete_repo(token=cls._token		,			repo_id='''valid_org/test-model-flax-org'''		)
								except HTTPError:
															pass







	def      A	(						self			:     int		)			->		Tuple:
								UpperCAmelCase :       List[Any] 							=		BertConfig(
								    vocab_size=99		,			hidden_size=32		,			num_hidden_layers=5		,			num_attention_heads=4		,			intermediate_size=37		)
								UpperCAmelCase :       Dict 							=		FlaxBertModel(__snake_case		)
								model.push_to_hub('''test-model-flax'''		,			use_auth_token=self._token		)

								UpperCAmelCase :       Tuple 							=		FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax"""		)

								UpperCAmelCase :       List[Any] 							=		flatten_dict(unfreeze(model.params		)		)
								UpperCAmelCase :       Tuple 							=		flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															UpperCAmelCase :       Union[str, Any] 							=		(base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(__snake_case		,			1E-3		,			msg=F"""{key} not identical"""		)

								# Reset repo
								delete_repo(token=self._token		,			repo_id='''test-model-flax'''		)

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(__snake_case		,			repo_id='''test-model-flax'''		,			push_to_hub=__snake_case		,			use_auth_token=self._token		)

								UpperCAmelCase :       str 							=		FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax"""		)

								UpperCAmelCase :       Any 							=		flatten_dict(unfreeze(model.params		)		)
								UpperCAmelCase :       str 							=		flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															UpperCAmelCase :       Optional[Any] 							=		(base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(__snake_case		,			1E-3		,			msg=F"""{key} not identical"""		)







	def      A	(						self			:     Optional[Any]		)			->		Union[str, Any]:
								UpperCAmelCase :       Dict 							=		BertConfig(
								    vocab_size=99		,			hidden_size=32		,			num_hidden_layers=5		,			num_attention_heads=4		,			intermediate_size=37		)
								UpperCAmelCase :       Optional[Any] 							=		FlaxBertModel(__snake_case		)
								model.push_to_hub('''valid_org/test-model-flax-org'''		,			use_auth_token=self._token		)

								UpperCAmelCase :       Union[str, Any] 							=		FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''		)

								UpperCAmelCase :       List[Any] 							=		flatten_dict(unfreeze(model.params		)		)
								UpperCAmelCase :       int 							=		flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															UpperCAmelCase :       Any 							=		(base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(__snake_case		,			1E-3		,			msg=F"""{key} not identical"""		)

								# Reset repo
								delete_repo(token=self._token		,			repo_id='''valid_org/test-model-flax-org'''		)

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(
															    __snake_case		,			repo_id='''valid_org/test-model-flax-org'''		,			push_to_hub=__snake_case		,			use_auth_token=self._token		)

								UpperCAmelCase :       str 							=		FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''		)

								UpperCAmelCase :       Any 							=		flatten_dict(unfreeze(model.params		)		)
								UpperCAmelCase :       Optional[Any] 							=		flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															UpperCAmelCase :       int 							=		(base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(__snake_case		,			1E-3		,			msg=F"""{key} not identical"""		)


def  snake_case_     (					_lowerCAmelCase  :						Optional[Any]   ,      _lowerCAmelCase  :						Any  )      ->   Union[str, Any]:
							UpperCAmelCase :       str 							=		True
							UpperCAmelCase :       int 							=		flatten_dict(modela.params  )
							UpperCAmelCase :       Dict 							=		flatten_dict(modela.params  )
							for key in flat_params_a.keys():
														if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]  )  ) > 1e-4:
																					UpperCAmelCase :       Dict 							=		False

							return models_are_equal




@require_flax
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	def      A	(						self			:     Tuple		)			->		Union[str, Any]:
								UpperCAmelCase :       List[Any] 							=		BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''		)
								UpperCAmelCase :       Dict 							=		FlaxBertModel(__snake_case		)

								UpperCAmelCase :       int 							=		'''bert'''
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(__snake_case		,			__snake_case		)		)

															with self.assertRaises(__snake_case		):
																						UpperCAmelCase :       Tuple 							=		FlaxBertModel.from_pretrained(__snake_case		)

															UpperCAmelCase :       str 							=		FlaxBertModel.from_pretrained(__snake_case		,			subfolder=__snake_case		)

								self.assertTrue(check_models_equal(__snake_case		,			__snake_case		)		)







	def      A	(						self			:     List[str]		)			->		Dict:
								UpperCAmelCase :       Dict 							=		BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''		)
								UpperCAmelCase :       Dict 							=		FlaxBertModel(__snake_case		)

								UpperCAmelCase :       Optional[int] 							=		'''bert'''
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(__snake_case		,			__snake_case		)		,			max_shard_size='''10KB'''		)

															with self.assertRaises(__snake_case		):
																						UpperCAmelCase :       Any 							=		FlaxBertModel.from_pretrained(__snake_case		)

															UpperCAmelCase :       Union[str, Any] 							=		FlaxBertModel.from_pretrained(__snake_case		,			subfolder=__snake_case		)

								self.assertTrue(check_models_equal(__snake_case		,			__snake_case		)		)







	def      A	(						self			:     Optional[int]		)			->		str:
								UpperCAmelCase :       Dict 							=		'''bert'''
								UpperCAmelCase :       int 							=		'''hf-internal-testing/tiny-random-bert-subfolder'''

								with self.assertRaises(__snake_case		):
															UpperCAmelCase :       Optional[Any] 							=		FlaxBertModel.from_pretrained(__snake_case		)

								UpperCAmelCase :       Tuple 							=		FlaxBertModel.from_pretrained(__snake_case		,			subfolder=__snake_case		)

								self.assertIsNotNone(__snake_case		)







	def      A	(						self			:     Dict		)			->		List[Any]:
								UpperCAmelCase :       Optional[int] 							=		'''bert'''
								UpperCAmelCase :       int 							=		'''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
								with self.assertRaises(__snake_case		):
															UpperCAmelCase :       Dict 							=		FlaxBertModel.from_pretrained(__snake_case		)

								UpperCAmelCase :       Union[str, Any] 							=		FlaxBertModel.from_pretrained(__snake_case		,			subfolder=__snake_case		)

								self.assertIsNotNone(__snake_case		)





