


'''simple docstring'''






import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

  if is_vision_available():
    from transformers import MaskaFormerImageProcessor

if is_vision_available():
  from PIL import Image




class 	a__       :





       '''simple docstring'''






       def __init__(	self      ,					lowerCamelCase_      ,					lowerCamelCase_=2      ,					lowerCamelCase_=True      ,					lowerCamelCase_=False      ,					lowerCamelCase_=10      ,					lowerCamelCase_=3      ,					lowerCamelCase_=32 * 8      ,					lowerCamelCase_=32 * 8      ,					lowerCamelCase_=4      ,					lowerCamelCase_=64      ,					) ->							List[Any]:
              lowerCAmelCase__      =			parent
              lowerCAmelCase__      =			batch_size
              lowerCAmelCase__      =			is_training
              lowerCAmelCase__      =			use_auxiliary_loss
              lowerCAmelCase__      =			num_queries
              lowerCAmelCase__      =			num_channels
              lowerCAmelCase__      =			min_size
              lowerCAmelCase__      =			max_size
              lowerCAmelCase__      =			num_labels
              lowerCAmelCase__      =			hidden_dim
              lowerCAmelCase__      =			hidden_dim







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[int]:
              lowerCAmelCase__      =			floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]     ).to(
                  lowerCamelCase_     )

              lowerCAmelCase__      =			torch.ones([self.batch_size, self.min_size, self.max_size]      ,					device=lowerCamelCase_     )

              lowerCAmelCase__      =			(
                  torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]      ,					device=lowerCamelCase_     ) > 0.5
              ).float()
              lowerCAmelCase__      =			(torch.rand((self.batch_size, self.num_labels)      ,					device=lowerCamelCase_     ) > 0.5).long()

              lowerCAmelCase__      =			self.get_config()
              return config, pixel_values, pixel_mask, mask_labels, class_labels







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
              lowerCAmelCase__      =			MaskaFormerConfig(
                  hidden_size=self.hidden_dim      ,					)
              lowerCAmelCase__      =			self.num_queries
              lowerCAmelCase__      =			self.num_labels

              lowerCAmelCase__      =			[1, 1, 1, 1]
              lowerCAmelCase__      =			self.num_channels

              lowerCAmelCase__      =			64
              lowerCAmelCase__      =			1_28
              lowerCAmelCase__      =			self.hidden_dim
              lowerCAmelCase__      =			self.hidden_dim
              lowerCAmelCase__      =			self.hidden_dim
              return config







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
              lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__      =			self.prepare_config_and_inputs()
              lowerCAmelCase__      =			{'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
              return config, inputs_dict







       def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_     ) ->							Optional[int]:
              lowerCAmelCase__      =			output.encoder_hidden_states
              lowerCAmelCase__      =			output.pixel_decoder_hidden_states
              lowerCAmelCase__      =			output.transformer_decoder_hidden_states

              self.parent.assertTrue(len(lowerCamelCase_     )      ,					len(config.backbone_config.depths     )     )
              self.parent.assertTrue(len(lowerCamelCase_     )      ,					len(config.backbone_config.depths     )     )
              self.parent.assertTrue(len(lowerCamelCase_     )      ,					config.decoder_layers     )







       def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_=False     ) ->							Tuple:
              with torch.no_grad():
                     lowerCAmelCase__      =			MaskaFormerModel(config=lowerCamelCase_     )
                     model.to(lowerCamelCase_     )
                     model.eval()

                     lowerCAmelCase__      =			model(pixel_values=lowerCamelCase_      ,					pixel_mask=lowerCamelCase_     )
                     lowerCAmelCase__      =			model(lowerCamelCase_      ,					output_hidden_states=lowerCamelCase_     )

              self.parent.assertEqual(
                  output.transformer_decoder_last_hidden_state.shape      ,					(self.batch_size, self.num_queries, self.hidden_dim)      ,					)
              # let's ensure the other two hidden state exists
              self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None     )
              self.parent.assertTrue(output.encoder_last_hidden_state is not None     )

              if output_hidden_states:
                     self.check_output_hidden_state(lowerCamelCase_      ,					lowerCamelCase_     )







       def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_     ) ->							Any:
              lowerCAmelCase__      =			MaskaFormerForUniversalSegmentation(config=lowerCamelCase_     )
              model.to(lowerCamelCase_     )
              model.eval()

              def comm_check_on_output(lowerCamelCase_     ):
                     # let's still check that all the required stuff is there
                     self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None     )
                     self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None     )
                     self.parent.assertTrue(result.encoder_last_hidden_state is not None     )
                     # okay, now we need to check the logits shape
                     # due to the encoder compression, masks have a //4 spatial size
                     self.parent.assertEqual(
                         result.masks_queries_logits.shape      ,					(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)      ,					)
                     # + 1 for null class
                     self.parent.assertEqual(
                         result.class_queries_logits.shape      ,					(self.batch_size, self.num_queries, self.num_labels + 1)     )

              with torch.no_grad():
                     lowerCAmelCase__      =			model(pixel_values=lowerCamelCase_      ,					pixel_mask=lowerCamelCase_     )
                     lowerCAmelCase__      =			model(lowerCamelCase_     )

                     comm_check_on_output(lowerCamelCase_     )

                     lowerCAmelCase__      =			model(
                         pixel_values=lowerCamelCase_      ,					pixel_mask=lowerCamelCase_      ,					mask_labels=lowerCamelCase_      ,					class_labels=lowerCamelCase_     )

              comm_check_on_output(lowerCamelCase_     )

              self.parent.assertTrue(result.loss is not None     )
              self.parent.assertEqual(result.loss.shape      ,					torch.Size([1]     )     )




@require_torch
class 	a__       (							a__       ,       a__       ,       unittest.TestCase						):





       '''simple docstring'''






       lowercase__    :     Tuple					=   (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
       lowercase__    :     Any					=   {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

       lowercase__    :     Union[str, Any]					=   False
       lowercase__    :     str					=   False
       lowercase__    :     List[Any]					=   False
       lowercase__    :     Dict					=   False







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Any:
              lowerCAmelCase__      =			MaskaFormerModelTester(self     )
              lowerCAmelCase__      =			ConfigTester(self      ,					config_class=lowerCamelCase_      ,					has_text_modality=lowerCamelCase_     )







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
              self.config_tester.run_common_tests()







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
              lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_common()
              self.model_tester.create_and_check_maskaformer_model(lowerCamelCase_      ,					**lowerCamelCase_      ,					output_hidden_states=lowerCamelCase_     )







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
              lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*lowerCamelCase_     )







       @unittest.skip(reason='''Mask2Former does not use inputs_embeds'''     )
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
              pass







       @unittest.skip(reason='''Mask2Former does not have a get_input_embeddings method'''     )
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[int]:
              pass







       @unittest.skip(reason='''Mask2Former is not a generative model'''     )
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Any:
              pass







       @unittest.skip(reason='''Mask2Former does not use token embeddings'''     )
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Union[str, Any]:
              pass







       @require_torch_multi_gpu
       @unittest.skip(
           reason='''Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''     )
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
              pass







       @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''     )
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
              pass







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Union[str, Any]:
              lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     lowerCAmelCase__      =			model_class(lowerCamelCase_     )
                     lowerCAmelCase__      =			inspect.signature(model.forward     )
                     # signature.parameters is an OrderedDict => so arg_names order is deterministic
                     lowerCAmelCase__      =			[*signature.parameters.keys()]

                     lowerCAmelCase__      =			['''pixel_values''']
                     self.assertListEqual(arg_names[:1]      ,					lowerCamelCase_     )







       @slow
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
              for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
                     lowerCAmelCase__      =			MaskaFormerModel.from_pretrained(lowerCamelCase_     )
                     self.assertIsNotNone(lowerCamelCase_     )







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Any:
              lowerCAmelCase__      =			(self.model_tester.min_size,) * 2
              lowerCAmelCase__      =			{
                  '''pixel_values''': torch.randn((2, 3, *size)      ,					device=lowerCamelCase_     ),
                  '''mask_labels''': torch.randn((2, 10, *size)      ,					device=lowerCamelCase_     ),
                  '''class_labels''': torch.zeros(2      ,					10      ,					device=lowerCamelCase_     ).long(),
              }
              lowerCAmelCase__      =			self.model_tester.get_config()

              lowerCAmelCase__      =			MaskaFormerForUniversalSegmentation(lowerCamelCase_     ).to(lowerCamelCase_     )
              lowerCAmelCase__      =			model(**lowerCamelCase_     )
              self.assertTrue(outputs.loss is not None     )







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
              lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_common()
              self.model_tester.create_and_check_maskaformer_model(lowerCamelCase_      ,					**lowerCamelCase_      ,					output_hidden_states=lowerCamelCase_     )







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
              lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     lowerCAmelCase__      =			model_class(lowerCamelCase_     ).to(lowerCamelCase_     )
                     lowerCAmelCase__      =			model(**lowerCamelCase_      ,					output_attentions=lowerCamelCase_     )
                     self.assertTrue(outputs.attentions is not None     )







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
              if not self.model_tester.is_training:
                     return

              lowerCAmelCase__      =			self.all_model_classes[1]
              lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()

              lowerCAmelCase__      =			model_class(lowerCamelCase_     )
              model.to(lowerCamelCase_     )
              model.train()

              lowerCAmelCase__      =			model(lowerCamelCase_      ,					mask_labels=lowerCamelCase_      ,					class_labels=lowerCamelCase_     ).loss
              loss.backward()







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
              lowerCAmelCase__      =			self.all_model_classes[1]
              lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()
              lowerCAmelCase__      =			True
              lowerCAmelCase__      =			True

              lowerCAmelCase__      =			model_class(lowerCamelCase_     ).to(lowerCamelCase_     )
              model.train()

              lowerCAmelCase__      =			model(lowerCamelCase_      ,					mask_labels=lowerCamelCase_      ,					class_labels=lowerCamelCase_     )

              lowerCAmelCase__      =			outputs.encoder_hidden_states[0]
              encoder_hidden_states.retain_grad()

              lowerCAmelCase__      =			outputs.pixel_decoder_hidden_states[0]
              pixel_decoder_hidden_states.retain_grad()

              lowerCAmelCase__      =			outputs.transformer_decoder_hidden_states[0]
              transformer_decoder_hidden_states.retain_grad()

              lowerCAmelCase__      =			outputs.attentions[0]
              attentions.retain_grad()

              outputs.loss.backward(retain_graph=lowerCamelCase_     )

              self.assertIsNotNone(encoder_hidden_states.grad     )
              self.assertIsNotNone(pixel_decoder_hidden_states.grad     )
              self.assertIsNotNone(transformer_decoder_hidden_states.grad     )
              self.assertIsNotNone(attentions.grad     )

__UpperCAmelCase					       =							1e-4


def       _snake_case     (  )		->					Dict:
       lowerCAmelCase__      =			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''' )
       return image






@require_vision
@slow
class 	a__       (							unittest.TestCase						):





       '''simple docstring'''






       @cached_property
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
              return "facebook/mask2former-swin-small-coco-instance"







       @cached_property
       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
              return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints     ) if is_vision_available() else None







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Any:
              lowerCAmelCase__      =			MaskaFormerModel.from_pretrained(self.model_checkpoints     ).to(lowerCamelCase_     )
              lowerCAmelCase__      =			self.default_image_processor
              lowerCAmelCase__      =			prepare_img()
              lowerCAmelCase__      =			image_processor(lowerCamelCase_      ,					return_tensors='''pt'''     ).to(lowerCamelCase_     )
              lowerCAmelCase__      =			inputs['''pixel_values'''].shape
              # check size is divisible by 32
              self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0     )
              # check size
              self.assertEqual(lowerCamelCase_      ,					(1, 3, 3_84, 3_84)     )

              with torch.no_grad():
                     lowerCAmelCase__      =			model(**lowerCamelCase_     )

              lowerCAmelCase__      =			torch.tensor(
                  [[-0.2_790, -1.0_717, -1.1_668], [-0.5_128, -0.3_128, -0.4_987], [-0.5_832, 0.1_971, -0.0_197]]     ).to(lowerCamelCase_     )
              self.assertTrue(
                  torch.allclose(
                      outputs.encoder_last_hidden_state[0, 0, :3, :3]      ,					lowerCamelCase_      ,					atol=lowerCamelCase_     )     )

              lowerCAmelCase__      =			torch.tensor(
                  [[0.8_973, 1.1_847, 1.1_776], [1.1_934, 1.5_040, 1.5_128], [1.1_153, 1.4_486, 1.4_951]]     ).to(lowerCamelCase_     )
              self.assertTrue(
                  torch.allclose(
                      outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]      ,					lowerCamelCase_      ,					atol=lowerCamelCase_     )     )

              lowerCAmelCase__      =			torch.tensor(
                  [[2.1_152, 1.7_000, -0.8_603], [1.5_808, 1.8_004, -0.9_353], [1.6_043, 1.7_495, -0.5_999]]     ).to(lowerCamelCase_     )
              self.assertTrue(
                  torch.allclose(
                      outputs.transformer_decoder_last_hidden_state[0, :3, :3]      ,					lowerCamelCase_      ,					atol=lowerCamelCase_     )     )







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
              lowerCAmelCase__      =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints     ).to(lowerCamelCase_     ).eval()
              lowerCAmelCase__      =			self.default_image_processor
              lowerCAmelCase__      =			prepare_img()
              lowerCAmelCase__      =			image_processor(lowerCamelCase_      ,					return_tensors='''pt'''     ).to(lowerCamelCase_     )
              lowerCAmelCase__      =			inputs['''pixel_values'''].shape
              # check size is divisible by 32
              self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0     )
              # check size
              self.assertEqual(lowerCamelCase_      ,					(1, 3, 3_84, 3_84)     )

              with torch.no_grad():
                     lowerCAmelCase__      =			model(**lowerCamelCase_     )
              # masks_queries_logits
              lowerCAmelCase__      =			outputs.masks_queries_logits
              self.assertEqual(
                  masks_queries_logits.shape      ,					(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)     )
              lowerCAmelCase__      =			[
                  [-8.7_839, -9.0_056, -8.8_121],
                  [-7.4_104, -7.0_313, -6.5_401],
                  [-6.6_105, -6.3_427, -6.4_675],
              ]
              lowerCAmelCase__      =			torch.tensor(lowerCamelCase_     ).to(lowerCamelCase_     )
              self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]      ,					lowerCamelCase_      ,					atol=lowerCamelCase_     )     )
              # class_queries_logits
              lowerCAmelCase__      =			outputs.class_queries_logits
              self.assertEqual(class_queries_logits.shape      ,					(1, model.config.num_queries, model.config.num_labels + 1)     )
              lowerCAmelCase__      =			torch.tensor(
                  [
                      [1.8_324, -8.0_835, -4.1_922],
                      [0.8_450, -9.0_050, -3.6_053],
                      [0.3_045, -7.7_293, -3.0_275],
                  ]     ).to(lowerCamelCase_     )
              self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]      ,					lowerCamelCase_      ,					atol=lowerCamelCase_     )     )







       def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[int]:
              lowerCAmelCase__      =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints     ).to(lowerCamelCase_     ).eval()
              lowerCAmelCase__      =			self.default_image_processor

              lowerCAmelCase__      =			image_processor(
                  [np.zeros((3, 8_00, 13_33)     ), np.zeros((3, 8_00, 13_33)     )]      ,					segmentation_maps=[np.zeros((3_84, 3_84)     ).astype(np.floataa     ), np.zeros((3_84, 3_84)     ).astype(np.floataa     )]      ,					return_tensors='''pt'''      ,					)

              lowerCAmelCase__      =			inputs['''pixel_values'''].to(lowerCamelCase_     )
              lowerCAmelCase__      =			[el.to(lowerCamelCase_     ) for el in inputs['''mask_labels''']]
              lowerCAmelCase__      =			[el.to(lowerCamelCase_     ) for el in inputs['''class_labels''']]

              with torch.no_grad():
                     lowerCAmelCase__      =			model(**lowerCamelCase_     )

              self.assertTrue(outputs.loss is not None     )



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




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





							'''simple docstring'''






							@classmethod
							def 						__SCREAMING_SNAKE_CASE							(	cls     ) ->							Any:
														lowerCAmelCase__      =			TOKEN
														HfFolder.save_token(lowerCamelCase_     )







							@classmethod
							def 						__SCREAMING_SNAKE_CASE							(	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 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			BertConfig(
														    vocab_size=99      ,					hidden_size=32      ,					num_hidden_layers=5      ,					num_attention_heads=4      ,					intermediate_size=37     )
														lowerCAmelCase__      =			FlaxBertModel(lowerCamelCase_     )
														model.push_to_hub('''test-model-flax'''      ,					use_auth_token=self._token     )

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

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

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

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

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

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







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

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

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

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

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

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

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


def       _snake_case     (  A				,					A )		->					Optional[int]:
							lowerCAmelCase__      =			True
							lowerCAmelCase__      =			flatten_dict(modela.params )
							lowerCAmelCase__      =			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:
																					lowerCAmelCase__      =			False

							return models_are_equal






@require_flax
class 	a__       (							unittest.TestCase						):





							'''simple docstring'''






							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
														lowerCAmelCase__      =			BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''     )
														lowerCAmelCase__      =			FlaxBertModel(lowerCamelCase_     )

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

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

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

														self.assertTrue(check_models_equal(lowerCamelCase_      ,					lowerCamelCase_     )     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''     )
														lowerCAmelCase__      =			FlaxBertModel(lowerCamelCase_     )

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

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

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

														self.assertTrue(check_models_equal(lowerCamelCase_      ,					lowerCamelCase_     )     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
														lowerCAmelCase__      =			'''bert'''
														lowerCAmelCase__      =			'''hf-internal-testing/tiny-random-bert-subfolder'''

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

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

														self.assertIsNotNone(lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			'''bert'''
														lowerCAmelCase__      =			'''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
														with self.assertRaises(lowerCamelCase_     ):
																					lowerCAmelCase__      =			FlaxBertModel.from_pretrained(lowerCamelCase_     )

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

														self.assertIsNotNone(lowerCamelCase_     )