





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

						_SCREAMING_SNAKE_CASE      	=		"""0.12"""  # assumed parallelism: 8
@require_flax
@is_staging_test
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






				@classmethod
				def UpperCAmelCase_				(      cls      :      Any	)							->    Any:







					"""simple docstring"""


					snake_case_		:			List[str]											=	TOKEN
					HfFolder.save_token(lowerCamelCase_	)






				@classmethod
				def UpperCAmelCase_				(      cls      :      Optional[int]	)							->    Tuple:







					"""simple docstring"""


					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 UpperCAmelCase_				(      self      :      Optional[Any]	)							->    List[Any]:







					"""simple docstring"""


					snake_case_		:			Tuple											=	BertConfig(
					    vocab_size=99		,   hidden_size=32		,   num_hidden_layers=5		,   num_attention_heads=4		,   intermediate_size=37	)
					snake_case_		:			Optional[int]											=	FlaxBertModel(lowerCamelCase_	)
					model.push_to_hub('test-model-flax'		,   use_auth_token=self._token	)

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

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

					for key in base_params.keys():
						snake_case_		:			Dict											=	(base_params[key] - new_params[key]).sum().item()
						self.assertLessEqual(lowerCamelCase_		,   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(lowerCamelCase_		,   repo_id='test-model-flax'		,   push_to_hub=lowerCamelCase_		,   use_auth_token=self._token	)

					snake_case_		:			List[Any]											=	FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax"""	)

					snake_case_		:			str											=	flatten_dict(unfreeze(model.params	)	)
					snake_case_		:			Dict											=	flatten_dict(unfreeze(new_model.params	)	)

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






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







					"""simple docstring"""


					snake_case_		:			Dict											=	BertConfig(
					    vocab_size=99		,   hidden_size=32		,   num_hidden_layers=5		,   num_attention_heads=4		,   intermediate_size=37	)
					snake_case_		:			List[Any]											=	FlaxBertModel(lowerCamelCase_	)
					model.push_to_hub('valid_org/test-model-flax-org'		,   use_auth_token=self._token	)

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

					snake_case_		:			int											=	flatten_dict(unfreeze(model.params	)	)
					snake_case_		:			str											=	flatten_dict(unfreeze(new_model.params	)	)

					for key in base_params.keys():
						snake_case_		:			Tuple											=	(base_params[key] - new_params[key]).sum().item()
						self.assertLessEqual(lowerCamelCase_		,   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(
						    lowerCamelCase_		,   repo_id='valid_org/test-model-flax-org'		,   push_to_hub=lowerCamelCase_		,   use_auth_token=self._token	)

					snake_case_		:			Any											=	FlaxBertModel.from_pretrained('valid_org/test-model-flax-org'	)

					snake_case_		:			str											=	flatten_dict(unfreeze(model.params	)	)
					snake_case_		:			Union[str, Any]											=	flatten_dict(unfreeze(new_model.params	)	)

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


def 							SCREAMING_SNAKE_CASE__						(    __a   , __a   ):
	snake_case_		:			Dict											=	True
	snake_case_		:			Any											=	flatten_dict(modela.params   )
	snake_case_		:			List[Any]											=	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:
			snake_case_		:			List[Any]											=	False

	return models_are_equal


@require_flax
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






				def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    Optional[Any]:







					"""simple docstring"""


					snake_case_		:			Optional[int]											=	BertConfig.from_pretrained('hf-internal-testing/tiny-bert-flax-only'	)
					snake_case_		:			Optional[Any]											=	FlaxBertModel(lowerCamelCase_	)

					snake_case_		:			Dict											=	'bert'
					with tempfile.TemporaryDirectory() as tmp_dir:
						model.save_pretrained(os.path.join(lowerCamelCase_		,   lowerCamelCase_	)	)

						with self.assertRaises(lowerCamelCase_	):
							snake_case_		:			List[str]											=	FlaxBertModel.from_pretrained(lowerCamelCase_	)

						snake_case_		:			str											=	FlaxBertModel.from_pretrained(lowerCamelCase_		,   subfolder=lowerCamelCase_	)

					self.assertTrue(check_models_equal(lowerCamelCase_		,   lowerCamelCase_	)	)






				def UpperCAmelCase_				(      self      :      Optional[int]	)							->    List[Any]:







					"""simple docstring"""


					snake_case_		:			str											=	BertConfig.from_pretrained('hf-internal-testing/tiny-bert-flax-only'	)
					snake_case_		:			Tuple											=	FlaxBertModel(lowerCamelCase_	)

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

						with self.assertRaises(lowerCamelCase_	):
							snake_case_		:			str											=	FlaxBertModel.from_pretrained(lowerCamelCase_	)

						snake_case_		:			Tuple											=	FlaxBertModel.from_pretrained(lowerCamelCase_		,   subfolder=lowerCamelCase_	)

					self.assertTrue(check_models_equal(lowerCamelCase_		,   lowerCamelCase_	)	)






				def UpperCAmelCase_				(      self      :      List[str]	)							->    Optional[Any]:







					"""simple docstring"""


					snake_case_		:			Dict											=	'bert'
					snake_case_		:			Any											=	'hf-internal-testing/tiny-random-bert-subfolder'

					with self.assertRaises(lowerCamelCase_	):
						snake_case_		:			List[str]											=	FlaxBertModel.from_pretrained(lowerCamelCase_	)

					snake_case_		:			List[str]											=	FlaxBertModel.from_pretrained(lowerCamelCase_		,   subfolder=lowerCamelCase_	)

					self.assertIsNotNone(lowerCamelCase_	)






				def UpperCAmelCase_				(      self      :      Dict	)							->    List[Any]:







					"""simple docstring"""


					snake_case_		:			str											=	'bert'
					snake_case_		:			str											=	'hf-internal-testing/tiny-random-bert-sharded-subfolder'
					with self.assertRaises(lowerCamelCase_	):
						snake_case_		:			Optional[Any]											=	FlaxBertModel.from_pretrained(lowerCamelCase_	)

					snake_case_		:			Dict											=	FlaxBertModel.from_pretrained(lowerCamelCase_		,   subfolder=lowerCamelCase_	)

					self.assertIsNotNone(lowerCamelCase_	)








def 							SCREAMING_SNAKE_CASE__						(    __a   ):
	if not isinstance(__a   , __a   ):
		raise ValueError('Input must be an integer'   )
	if input_num <= 0:
		raise ValueError('Input must be positive'   )
	return sum(
	    divisor for divisor in range(1   , input_num // 2 + 1   ) if input_num % divisor == 0   )


if __name__ == "__main__":
						import doctest

						doctest.testmod()

