






'''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						:						int  =					'0.12'  # assumed parallelism: 8


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




					@classmethod
					def 		A      (  cls  :		List[Any]      )    ->							List[str]:
								UpperCAmelCase_      :       Tuple          =       TOKEN
								HfFolder.save_token(_A      )




					@classmethod
					def 		A      (  cls  :		Dict      )    ->							str:
								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  :		Dict      )    ->							Union[str, Any]:
								UpperCAmelCase_      :       Tuple          =       BertConfig(
								    vocab_size=99				, hidden_size=32				, num_hidden_layers=5				, num_attention_heads=4				, intermediate_size=37      )
								UpperCAmelCase_      :       Dict          =       FlaxBertModel(_A      )
								model.push_to_hub('''test-model-flax'''				, use_auth_token=self._token      )

								UpperCAmelCase_      :       Optional[int]          =       FlaxBertModel.from_pretrained(F"{USER}/test-model-flax"      )

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

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

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

								UpperCAmelCase_      :       Optional[int]          =       flatten_dict(unfreeze(model.params      )      )
								UpperCAmelCase_      :       Dict          =       flatten_dict(unfreeze(new_model.params      )      )

								for key in base_params.keys():
											UpperCAmelCase_      :       Tuple          =       (base_params[key] - new_params[key]).sum().item()
											self.assertLessEqual(_A				, 1e-3				, msg=F"{key} not identical"      )





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

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

								UpperCAmelCase_      :       Optional[Any]          =       flatten_dict(unfreeze(model.params      )      )
								UpperCAmelCase_      :       Optional[int]          =       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(_A				, 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(
											    _A				, repo_id='''valid_org/test-model-flax-org'''				, push_to_hub=_A				, use_auth_token=self._token      )

								UpperCAmelCase_      :       Optional[Any]          =       FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''      )

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

								for key in base_params.keys():
											UpperCAmelCase_      :       str          =       (base_params[key] - new_params[key]).sum().item()
											self.assertLessEqual(_A				, 1e-3				, msg=F"{key} not identical"      )






def __UpperCAmelCase   (					A					:			List[str]			,       A					:			Tuple		)    ->   Optional[int]:
			UpperCAmelCase_      :       Optional[int]          =       True
			UpperCAmelCase_      :       List[Any]          =       flatten_dict(modela.params		)
			UpperCAmelCase_      :       List[str]          =       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_      :       List[Any]          =       False

			return models_are_equal







@require_flax
class  snake_case__      (					unittest.TestCase):




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

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

											with self.assertRaises(_A      ):
														UpperCAmelCase_      :       int          =       FlaxBertModel.from_pretrained(_A      )

											UpperCAmelCase_      :       Optional[Any]          =       FlaxBertModel.from_pretrained(_A				, subfolder=_A      )

								self.assertTrue(check_models_equal(_A				, _A      )      )




					def 		A      (  self  :		int      )    ->							str:
								UpperCAmelCase_      :       str          =       BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''      )
								UpperCAmelCase_      :       Dict          =       FlaxBertModel(_A      )

								UpperCAmelCase_      :       str          =       '''bert'''
								with tempfile.TemporaryDirectory() as tmp_dir:
											model.save_pretrained(os.path.join(_A				, _A      )				, max_shard_size='''10KB'''      )

											with self.assertRaises(_A      ):
														UpperCAmelCase_      :       Optional[int]          =       FlaxBertModel.from_pretrained(_A      )

											UpperCAmelCase_      :       Any          =       FlaxBertModel.from_pretrained(_A				, subfolder=_A      )

								self.assertTrue(check_models_equal(_A				, _A      )      )




					def 		A      (  self  :		List[str]      )    ->							Union[str, Any]:
								UpperCAmelCase_      :       List[Any]          =       '''bert'''
								UpperCAmelCase_      :       int          =       '''hf-internal-testing/tiny-random-bert-subfolder'''

								with self.assertRaises(_A      ):
											UpperCAmelCase_      :       Optional[Any]          =       FlaxBertModel.from_pretrained(_A      )

								UpperCAmelCase_      :       Dict          =       FlaxBertModel.from_pretrained(_A				, subfolder=_A      )

								self.assertIsNotNone(_A      )





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

								UpperCAmelCase_      :       Optional[int]          =       FlaxBertModel.from_pretrained(_A				, subfolder=_A      )

								self.assertIsNotNone(_A      )










'''simple docstring'''





_UpperCamelCase						:						List[str]  =					[sum(int(c, 10) ** 2 for c in i.__str__()) for i in range(100_000)]
def __UpperCAmelCase   (					A					:			int		)    ->   int:
			UpperCAmelCase_      :       str          =       0
			while number:
						# Increased Speed Slightly by checking every 5 digits together.
						sum_of_digits_squared += DIGITS_SQUARED[number % 1_0_0_0_0_0]
						number //= 1_0_0_0_0_0

			return sum_of_digits_squared


# There are 2 Chains made,
# One ends with 89 with the chain member 58 being the one which when declared first,
# there will be the least number of iterations for all the members to be checked.

# The other one ends with 1 and has only one element 1.

# So 58 and 1 are chosen to be declared at the starting.

# Changed dictionary to an array to quicken the solution
_UpperCamelCase						:						list[bool | None]  =					[None] * 10_000_000
_UpperCamelCase						:						List[str]  =					True
_UpperCamelCase						:						Tuple  =					False
def __UpperCAmelCase   (					A					:			int		)    ->   bool:

			if CHAINS[number - 1] is not None:
						return CHAINS[number - 1]  # type: ignore

			UpperCAmelCase_      :       Dict          =       chain(next_number(A		)		)
			UpperCAmelCase_      :       int          =       number_chain

			while number < 1_0_0_0_0_0_0_0:
						UpperCAmelCase_      :       Optional[Any]          =       number_chain
						number *= 1_0

			return number_chain
def __UpperCAmelCase   (					A					:			int = 1_0_0_0_0_0_0_0		)    ->   int:
			for i in range(1			,       A		):
						if CHAINS[i] is None:
									chain(i + 1		)

			return CHAINS[:number].count(A		)


if __name__ == "__main__":
						import doctest

						doctest.testmod()
						print(f'''{solution() = }''')


