






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

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





	@classmethod
	def 						_A (		cls				:					str	):
		UpperCamelCase     :Optional[int]          =	TOKEN
		HfFolder.save_token(__lowerCamelCase	)





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

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

		UpperCamelCase     :str          =	flatten_dict(unfreeze(model.params	)	)
		UpperCamelCase     :List[Any]          =	flatten_dict(unfreeze(new_model.params	)	)

		for key in base_params.keys():
			UpperCamelCase     :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	)

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

		UpperCamelCase     :int          =	flatten_dict(unfreeze(model.params	)	)
		UpperCamelCase     :Tuple          =	flatten_dict(unfreeze(new_model.params	)	)

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





	def 						_A (		self				:					int	):
		UpperCamelCase     :Optional[int]          =	BertConfig(
		    vocab_size=99 ,  hidden_size=32 ,  num_hidden_layers=5 ,  num_attention_heads=4 ,  intermediate_size=37	)
		UpperCamelCase     :Any          =	FlaxBertModel(__lowerCamelCase	)
		model.push_to_hub("""valid_org/test-model-flax-org""" ,  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     :List[str]          =	(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	)

		UpperCamelCase     :int          =	FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""	)

		UpperCamelCase     :str          =	flatten_dict(unfreeze(model.params	)	)
		UpperCamelCase     :Union[str, Any]          =	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(__lowerCamelCase ,  1E-3 ,  msg=F"""{key} not identical"""	)

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  Tuple	,			__magic_name__					:  Any					)		->			Tuple:






	"""simple docstring"""



	UpperCamelCase     :Optional[int]          =	True
	UpperCamelCase     :str          =	flatten_dict(modela.params					)
	UpperCamelCase     :Optional[int]          =	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     :Any          =	False

	return models_are_equal

@require_flax
class     _SCREAMING_SNAKE_CASE       (			unittest.TestCase      ):





	def 						_A (		self				:					Dict	):
		UpperCamelCase     :Optional[Any]          =	BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""	)
		UpperCamelCase     :Union[str, Any]          =	FlaxBertModel(__lowerCamelCase	)

		UpperCamelCase     :Dict          =	"""bert"""
		with tempfile.TemporaryDirectory() as tmp_dir:
			model.save_pretrained(os.path.join(__lowerCamelCase ,  __lowerCamelCase	)	)

			with self.assertRaises(__lowerCamelCase	):
				UpperCamelCase     :str          =	FlaxBertModel.from_pretrained(__lowerCamelCase	)

			UpperCamelCase     :Any          =	FlaxBertModel.from_pretrained(__lowerCamelCase ,  subfolder=__lowerCamelCase	)

		self.assertTrue(check_models_equal(__lowerCamelCase ,  __lowerCamelCase	)	)





	def 						_A (		self				:					Dict	):
		UpperCamelCase     :int          =	BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""	)
		UpperCamelCase     :List[Any]          =	FlaxBertModel(__lowerCamelCase	)

		UpperCamelCase     :Union[str, Any]          =	"""bert"""
		with tempfile.TemporaryDirectory() as tmp_dir:
			model.save_pretrained(os.path.join(__lowerCamelCase ,  __lowerCamelCase	) ,  max_shard_size="""10KB"""	)

			with self.assertRaises(__lowerCamelCase	):
				UpperCamelCase     :List[str]          =	FlaxBertModel.from_pretrained(__lowerCamelCase	)

			UpperCamelCase     :str          =	FlaxBertModel.from_pretrained(__lowerCamelCase ,  subfolder=__lowerCamelCase	)

		self.assertTrue(check_models_equal(__lowerCamelCase ,  __lowerCamelCase	)	)





	def 						_A (		self				:					Optional[Any]	):
		UpperCamelCase     :Union[str, Any]          =	"""bert"""
		UpperCamelCase     :Optional[Any]          =	"""hf-internal-testing/tiny-random-bert-subfolder"""

		with self.assertRaises(__lowerCamelCase	):
			UpperCamelCase     :Tuple          =	FlaxBertModel.from_pretrained(__lowerCamelCase	)

		UpperCamelCase     :Tuple          =	FlaxBertModel.from_pretrained(__lowerCamelCase ,  subfolder=__lowerCamelCase	)

		self.assertIsNotNone(__lowerCamelCase	)





	def 						_A (		self				:					Optional[Any]	):
		UpperCamelCase     :List[str]          =	"""bert"""
		UpperCamelCase     :str          =	"""hf-internal-testing/tiny-random-bert-sharded-subfolder"""
		with self.assertRaises(__lowerCamelCase	):
			UpperCamelCase     :Dict          =	FlaxBertModel.from_pretrained(__lowerCamelCase	)

		UpperCamelCase     :str          =	FlaxBertModel.from_pretrained(__lowerCamelCase ,  subfolder=__lowerCamelCase	)

		self.assertIsNotNone(__lowerCamelCase	)








def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  list[list[int]]	,			__magic_name__					:  int	,			__magic_name__					:  int	,			__magic_name__					:  list[int]					)		->			bool:






	"""simple docstring"""



	if graph[path[curr_ind - 1]][next_ver] == 0:
		return False

	# 2. Validate that next vertex is not already in path
	return not any(vertex == next_ver for vertex in path					)

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  list[list[int]]	,			__magic_name__					:  list[int]	,			__magic_name__					:  int					)		->			bool:






	"""simple docstring"""



	if curr_ind == len(__magic_name__					):
		# return whether path exists between current and starting vertices
		return graph[path[curr_ind - 1]][path[0]] == 1

	# Recursive Step
	for next_ver in range(0	,			len(__magic_name__					)					):
		if valid_connection(__magic_name__	,			__magic_name__	,			__magic_name__	,			__magic_name__					):
			# Insert current vertex  into path as next transition
			UpperCamelCase     :str          =	next_ver
			# Validate created path
			if util_hamilton_cycle(__magic_name__	,			__magic_name__	,			curr_ind + 1					):
				return True
			# Backtrack
			UpperCamelCase     :Union[str, Any]          =	-1
	return False

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  list[list[int]]	,			__magic_name__					:  int = 0					)		->			list[int]:






	"""simple docstring"""



	UpperCamelCase     :Union[str, Any]          =	[-1] * (len(__magic_name__					) + 1)
	# initialize start and end of path with starting index
	UpperCamelCase     :Any          =	start_index
	# evaluate and if we find answer return path either return empty array
	return path if util_hamilton_cycle(__magic_name__	,			__magic_name__	,			1					) else []
