


import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
			from transformers.models.bert.modeling_flax_bert import (
			    FlaxBertForMaskedLM,
			    FlaxBertForMultipleChoice,
			    FlaxBertForNextSentencePrediction,
			    FlaxBertForPreTraining,
			    FlaxBertForQuestionAnswering,
			    FlaxBertForSequenceClassification,
			    FlaxBertForTokenClassification,
			    FlaxBertModel,
			)

class        a	(unittest.TestCase   ):




							"""simple docstring"""


							def __init__(							self :						Dict					,			lowerCamelCase :						int					,			lowerCamelCase :						Optional[int]=13					,			lowerCamelCase :						Union[str, Any]=7					,			lowerCamelCase :						Union[str, Any]=True					,			lowerCamelCase :						str=True					,			lowerCamelCase :						List[str]=True					,			lowerCamelCase :						Optional[int]=True					,			lowerCamelCase :						int=99					,			lowerCamelCase :						int=32					,			lowerCamelCase :						Union[str, Any]=5					,			lowerCamelCase :						Tuple=4					,			lowerCamelCase :						int=37					,			lowerCamelCase :						List[Any]="gelu"					,			lowerCamelCase :						Optional[int]=0.1					,			lowerCamelCase :						Optional[int]=0.1					,			lowerCamelCase :						List[Any]=512					,			lowerCamelCase :						Optional[int]=16					,			lowerCamelCase :						Tuple=2					,			lowerCamelCase :						Any=0.02					,			lowerCamelCase :						int=4					,			)			->					Optional[Any]:
														__snake_case	:    Any											=     parent
														__snake_case	:    Any											=     batch_size
														__snake_case	:    List[Any]											=     seq_length
														__snake_case	:    Union[str, Any]											=     is_training
														__snake_case	:    Tuple											=     use_attention_mask
														__snake_case	:    Union[str, Any]											=     use_token_type_ids
														__snake_case	:    List[str]											=     use_labels
														__snake_case	:    str											=     vocab_size
														__snake_case	:    List[str]											=     hidden_size
														__snake_case	:    List[Any]											=     num_hidden_layers
														__snake_case	:    str											=     num_attention_heads
														__snake_case	:    List[str]											=     intermediate_size
														__snake_case	:    Optional[Any]											=     hidden_act
														__snake_case	:    Optional[Any]											=     hidden_dropout_prob
														__snake_case	:    Optional[int]											=     attention_probs_dropout_prob
														__snake_case	:    Optional[Any]											=     max_position_embeddings
														__snake_case	:    Tuple											=     type_vocab_size
														__snake_case	:    Dict											=     type_sequence_label_size
														__snake_case	:    List[Any]											=     initializer_range
														__snake_case	:    List[Any]											=     num_choices

							def 					__snake_case      (							self :						Dict       )			->					Union[str, Any]:
														__snake_case	:    str											=     ids_tensor([self.batch_size, self.seq_length]					,			self.vocab_size       )

														__snake_case	:    List[str]											=     None
														if self.use_attention_mask:
																					__snake_case	:    int											=     random_attention_mask([self.batch_size, self.seq_length]       )

														__snake_case	:    Any											=     None
														if self.use_token_type_ids:
																					__snake_case	:    Optional[int]											=     ids_tensor([self.batch_size, self.seq_length]					,			self.type_vocab_size       )

														__snake_case	:    Union[str, Any]											=     BertConfig(
														    vocab_size=self.vocab_size					,			hidden_size=self.hidden_size					,			num_hidden_layers=self.num_hidden_layers					,			num_attention_heads=self.num_attention_heads					,			intermediate_size=self.intermediate_size					,			hidden_act=self.hidden_act					,			hidden_dropout_prob=self.hidden_dropout_prob					,			attention_probs_dropout_prob=self.attention_probs_dropout_prob					,			max_position_embeddings=self.max_position_embeddings					,			type_vocab_size=self.type_vocab_size					,			is_decoder=lowerCamelCase					,			initializer_range=self.initializer_range					,			)

														return config, input_ids, token_type_ids, attention_mask

							def 					__snake_case      (							self :						List[str]       )			->					Tuple:
														__snake_case	:    Optional[int]											=     self.prepare_config_and_inputs()
														__snake_case	:    Union[str, Any]											=     config_and_inputs
														__snake_case	:    str											=     {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
														return config, inputs_dict






							def 					__snake_case      (							self :						Optional[int]       )			->					List[str]:
														__snake_case	:    Optional[Any]											=     self.prepare_config_and_inputs()
														__snake_case	:    str											=     config_and_inputs

														__snake_case	:    Optional[Any]											=     True
														__snake_case	:    Optional[Any]											=     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
														__snake_case	:    str											=     ids_tensor([self.batch_size, self.seq_length]					,			vocab_size=2       )

														return (
														    config,
														    input_ids,
														    attention_mask,
														    encoder_hidden_states,
														    encoder_attention_mask,
														)

@require_flax
class        a	(_lowerCAmelCase  ,       unittest.TestCase   ):




							"""simple docstring"""


							__UpperCAmelCase					:							Union[str, Any]  						=						True

							__UpperCAmelCase					:							str  						=						(
							    (
							        FlaxBertModel,
							        FlaxBertForPreTraining,
							        FlaxBertForMaskedLM,
							        FlaxBertForMultipleChoice,
							        FlaxBertForQuestionAnswering,
							        FlaxBertForNextSentencePrediction,
							        FlaxBertForSequenceClassification,
							        FlaxBertForTokenClassification,
							        FlaxBertForQuestionAnswering,
							    )
							    if is_flax_available()
							    else ()
							)

							def 					__snake_case      (							self :						Any       )			->					Union[str, Any]:
														__snake_case	:    str											=     FlaxBertModelTester(self       )






							@slow
							def 					__snake_case      (							self :						Optional[Any]       )			->					List[str]:
														# Only check this for base model, not necessary for all model classes.
														# This will also help speed-up tests.
														__snake_case	:    Dict											=     FlaxBertModel.from_pretrained("bert-base-cased"       )
														__snake_case	:    Tuple											=     model(np.ones((1, 1)       )       )
														self.assertIsNotNone(lowerCamelCase       )





import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
			from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
			import torch_xla.core.xla_model as xm







def        lowerCAmelCase_  (      __lowerCamelCase			):
							if is_torch_version("<"							,			"2.0.0"			) or not hasattr(__lowerCamelCase							,			"_dynamo"			):
														return False
							return isinstance(__lowerCamelCase							,			torch._dynamo.eval_frame.OptimizedModule			)







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase = True			):
							__snake_case	:    Optional[Any]											=     (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

							__snake_case	:    Tuple											=     is_compiled_module(__lowerCamelCase			)
							if is_compiled:
														__snake_case	:    Tuple											=     model
														__snake_case	:    List[str]											=     model._orig_mod

							if is_deepspeed_available():
														options += (DeepSpeedEngine,)

							while isinstance(__lowerCamelCase							,			__lowerCamelCase			):
														__snake_case	:    List[Any]											=     model.module

							if not keep_fpaa_wrapper:
														__snake_case	:    Tuple											=     getattr(__lowerCamelCase							,			"forward"			)
														__snake_case	:    int											=     model.__dict__.pop("_original_forward"							,			__lowerCamelCase			)
														if original_forward is not None:
																					while hasattr(__lowerCamelCase							,			"__wrapped__"			):
																												__snake_case	:    str											=     forward.__wrapped__
																												if forward == original_forward:
																																			break
																					__snake_case	:    Dict											=     forward
														if getattr(__lowerCamelCase							,			"_converted_to_transformer_engine"							,			__lowerCamelCase			):
																					convert_model(__lowerCamelCase							,			to_transformer_engine=__lowerCamelCase			)

							if is_compiled:
														__snake_case	:    Tuple											=     model
														__snake_case	:    str											=     compiled_model

							return model







def        lowerCAmelCase_  (      ):
							PartialState().wait_for_everyone()







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase			):
							if PartialState().distributed_type == DistributedType.TPU:
														xm.save(__lowerCamelCase							,			__lowerCamelCase			)
							elif PartialState().local_process_index == 0:
														torch.save(__lowerCamelCase							,			__lowerCamelCase			)







@contextmanager
def        lowerCAmelCase_  (      **__lowerCamelCase			):
							for key, value in kwargs.items():
														__snake_case	:    Union[str, Any]											=     str(__lowerCamelCase			)

							yield

							for key in kwargs:
														if key.upper() in os.environ:
																					del os.environ[key.upper()]







def        lowerCAmelCase_  (      __lowerCamelCase			):
							if not hasattr(__lowerCamelCase							,			"__qualname__"			) and not hasattr(__lowerCamelCase							,			"__name__"			):
														__snake_case	:    Dict											=     getattr(__lowerCamelCase							,			"__class__"							,			__lowerCamelCase			)
							if hasattr(__lowerCamelCase							,			"__qualname__"			):
														return obj.__qualname__
							if hasattr(__lowerCamelCase							,			"__name__"			):
														return obj.__name__
							return str(__lowerCamelCase			)







def        lowerCAmelCase_  (      __lowerCamelCase							,			__lowerCamelCase			):
							for key, value in source.items():
														if isinstance(__lowerCamelCase							,			__lowerCamelCase			):
																					__snake_case	:    Tuple											=     destination.setdefault(__lowerCamelCase							,			{}			)
																					merge_dicts(__lowerCamelCase							,			__lowerCamelCase			)
														else:
																					__snake_case	:    Union[str, Any]											=     value

							return destination







def        lowerCAmelCase_  (      __lowerCamelCase = None			):
							if port is None:
														__snake_case	:    List[str]											=     2_9_5_0_0
							with socket.socket(socket.AF_INET							,			socket.SOCK_STREAM			) as s:
														return s.connect_ex(("localhost", port)			) == 0

