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

		__a  :     str       =  """0.12"""  # assumed parallelism: 8




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


			"""simple docstring"""






			@classmethod
			def        _SCREAMING_SNAKE_CASE			(   cls						)     ->					int:

							'''simple docstring'''


							__lowercase 			=							TOKEN
							HfFolder.save_token(lowerCAmelCase__						)





			@classmethod
			def        _SCREAMING_SNAKE_CASE			(   cls						)     ->					Dict:

							'''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        _SCREAMING_SNAKE_CASE			(   self						)     ->					Dict:

							'''simple docstring'''


							__lowercase 			=							BertConfig(
							    vocab_size=99      ,   hidden_size=32      ,   num_hidden_layers=5      ,   num_attention_heads=4      ,   intermediate_size=37						)
							__lowercase 			=							FlaxBertModel(lowerCAmelCase__						)
							model.push_to_hub('''test-model-flax'''      ,   use_auth_token=self._token						)

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

							__lowercase 			=							flatten_dict(unfreeze(model.params						)						)
							__lowercase 			=							flatten_dict(unfreeze(new_model.params						)						)

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

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

							__lowercase 			=							flatten_dict(unfreeze(model.params						)						)
							__lowercase 			=							flatten_dict(unfreeze(new_model.params						)						)

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





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					List[str]:

							'''simple docstring'''


							__lowercase 			=							BertConfig(
							    vocab_size=99      ,   hidden_size=32      ,   num_hidden_layers=5      ,   num_attention_heads=4      ,   intermediate_size=37						)
							__lowercase 			=							FlaxBertModel(lowerCAmelCase__						)
							model.push_to_hub('''valid_org/test-model-flax-org'''      ,   use_auth_token=self._token						)

							__lowercase 			=							FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''						)

							__lowercase 			=							flatten_dict(unfreeze(model.params						)						)
							__lowercase 			=							flatten_dict(unfreeze(new_model.params						)						)

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

							__lowercase 			=							FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''						)

							__lowercase 			=							flatten_dict(unfreeze(model.params						)						)
							__lowercase 			=							flatten_dict(unfreeze(new_model.params						)						)

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





def 				UpperCAmelCase						(    lowercase							,			lowercase							):







				"""simple docstring"""



				__lowercase 			=							True
				__lowercase 			=							flatten_dict(modela.params							)
				__lowercase 			=							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:
												__lowercase 			=							False

				return models_are_equal




@require_flax
class    _UpperCamelCase      (						unittest.TestCase ):


			"""simple docstring"""






			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Tuple:

							'''simple docstring'''


							__lowercase 			=							BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''						)
							__lowercase 			=							FlaxBertModel(lowerCAmelCase__						)

							__lowercase 			=							'''bert'''
							with tempfile.TemporaryDirectory() as tmp_dir:
											model.save_pretrained(os.path.join(lowerCAmelCase__      ,   lowerCAmelCase__						)						)

											with self.assertRaises(lowerCAmelCase__						):
															__lowercase 			=							FlaxBertModel.from_pretrained(lowerCAmelCase__						)

											__lowercase 			=							FlaxBertModel.from_pretrained(lowerCAmelCase__      ,   subfolder=lowerCAmelCase__						)

							self.assertTrue(check_models_equal(lowerCAmelCase__      ,   lowerCAmelCase__						)						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					List[str]:

							'''simple docstring'''


							__lowercase 			=							BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''						)
							__lowercase 			=							FlaxBertModel(lowerCAmelCase__						)

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

											with self.assertRaises(lowerCAmelCase__						):
															__lowercase 			=							FlaxBertModel.from_pretrained(lowerCAmelCase__						)

											__lowercase 			=							FlaxBertModel.from_pretrained(lowerCAmelCase__      ,   subfolder=lowerCAmelCase__						)

							self.assertTrue(check_models_equal(lowerCAmelCase__      ,   lowerCAmelCase__						)						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Any:

							'''simple docstring'''


							__lowercase 			=							'''bert'''
							__lowercase 			=							'''hf-internal-testing/tiny-random-bert-subfolder'''

							with self.assertRaises(lowerCAmelCase__						):
											__lowercase 			=							FlaxBertModel.from_pretrained(lowerCAmelCase__						)

							__lowercase 			=							FlaxBertModel.from_pretrained(lowerCAmelCase__      ,   subfolder=lowerCAmelCase__						)

							self.assertIsNotNone(lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					List[Any]:

							'''simple docstring'''


							__lowercase 			=							'''bert'''
							__lowercase 			=							'''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
							with self.assertRaises(lowerCAmelCase__						):
											__lowercase 			=							FlaxBertModel.from_pretrained(lowerCAmelCase__						)

							__lowercase 			=							FlaxBertModel.from_pretrained(lowerCAmelCase__      ,   subfolder=lowerCAmelCase__						)

							self.assertIsNotNone(lowerCAmelCase__						)
import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest




class    _UpperCamelCase      (						_UpperCAmelCase ):


			"""simple docstring"""






			__a  : Optional[Any]			 =						(KDPMaDiscreteScheduler,)
			__a  : Dict			 =						10





			def        _SCREAMING_SNAKE_CASE			(   self      ,   **lowerCAmelCase__						)     ->					int:

							'''simple docstring'''


							__lowercase 			=							{
							    '''num_train_timesteps''': 11_00,
							    '''beta_start''': 0.0001,
							    '''beta_end''': 0.02,
							    '''beta_schedule''': '''linear''',
							}

							config.update(**lowerCAmelCase__						)
							return config





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					List[str]:

							'''simple docstring'''


							for timesteps in [10, 50, 1_00, 10_00]:
											self.check_over_configs(num_train_timesteps=lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Any:

							'''simple docstring'''


							for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]      ,   [0.0002, 0.002, 0.02]						):
											self.check_over_configs(beta_start=lowerCAmelCase__      ,   beta_end=lowerCAmelCase__						)





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

							'''simple docstring'''


							for schedule in ["linear", "scaled_linear"]:
											self.check_over_configs(beta_schedule=lowerCAmelCase__						)





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

							'''simple docstring'''


							for prediction_type in ["epsilon", "v_prediction"]:
											self.check_over_configs(prediction_type=lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Tuple:

							'''simple docstring'''


							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config(prediction_type='''v_prediction'''						)
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps						)

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter * scheduler.init_noise_sigma
							__lowercase 			=							sample.to(lowerCAmelCase__						)

							for i, t in enumerate(scheduler.timesteps						):
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if torch_device in ["cpu", "mps"]:
											assert abs(result_sum.item() - 4.6934E-07						) < 1E-2
											assert abs(result_mean.item() - 6.1112E-10						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 4.693_4286_5017_0972E-07						) < 1E-2
											assert abs(result_mean.item() - 0.0002						) < 1E-3





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					str:

							'''simple docstring'''


							if torch_device == "mps":
											return
							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config()
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps						)

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter * scheduler.init_noise_sigma
							__lowercase 			=							sample.to(lowerCAmelCase__						)

							for i, t in enumerate(scheduler.timesteps						):
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if torch_device in ["cpu", "mps"]:
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Dict:

							'''simple docstring'''


							if torch_device == "mps":
											return
							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config()
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps      ,   device=lowerCAmelCase__						)

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter.to(lowerCAmelCase__						) * scheduler.init_noise_sigma

							for t in scheduler.timesteps:
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if str(lowerCAmelCase__						).startswith('''cpu'''						):
											# The following sum varies between 148 and 156 on mps. Why?
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3