


'''simple docstring'''






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   ,       lowerCamelCase_   ,       lowerCamelCase_=13   ,       lowerCamelCase_=7   ,       lowerCamelCase_=True   ,       lowerCamelCase_=True   ,       lowerCamelCase_=True   ,       lowerCamelCase_=True   ,       lowerCamelCase_=99   ,       lowerCamelCase_=32   ,       lowerCamelCase_=5   ,       lowerCamelCase_=4   ,       lowerCamelCase_=37   ,       lowerCamelCase_="gelu"   ,       lowerCamelCase_=0.1   ,       lowerCamelCase_=0.1   ,       lowerCamelCase_=5_12   ,       lowerCamelCase_=16   ,       lowerCamelCase_=2   ,       lowerCamelCase_=0.02   ,       lowerCamelCase_=4   ,       )			->				List[Any]:
														lowerCAmelCase__			  =      parent
														lowerCAmelCase__			  =      batch_size
														lowerCAmelCase__			  =      seq_length
														lowerCAmelCase__			  =      is_training
														lowerCAmelCase__			  =      use_attention_mask
														lowerCAmelCase__			  =      use_token_type_ids
														lowerCAmelCase__			  =      use_labels
														lowerCAmelCase__			  =      vocab_size
														lowerCAmelCase__			  =      hidden_size
														lowerCAmelCase__			  =      num_hidden_layers
														lowerCAmelCase__			  =      num_attention_heads
														lowerCAmelCase__			  =      intermediate_size
														lowerCAmelCase__			  =      hidden_act
														lowerCAmelCase__			  =      hidden_dropout_prob
														lowerCAmelCase__			  =      attention_probs_dropout_prob
														lowerCAmelCase__			  =      max_position_embeddings
														lowerCAmelCase__			  =      type_vocab_size
														lowerCAmelCase__			  =      type_sequence_label_size
														lowerCAmelCase__			  =      initializer_range
														lowerCAmelCase__			  =      num_choices







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														lowerCAmelCase__			  =      ids_tensor([self.batch_size, self.seq_length]   ,       self.vocab_size  )

														lowerCAmelCase__			  =      None
														if self.use_attention_mask:
																					lowerCAmelCase__			  =      random_attention_mask([self.batch_size, self.seq_length]  )

														lowerCAmelCase__			  =      None
														if self.use_token_type_ids:
																					lowerCAmelCase__			  =      ids_tensor([self.batch_size, self.seq_length]   ,       self.type_vocab_size  )

														lowerCAmelCase__			  =      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     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      self.prepare_config_and_inputs()
														lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__			  =      config_and_inputs
														lowerCAmelCase__			  =      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
														return config, inputs_dict







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														lowerCAmelCase__			  =      self.prepare_config_and_inputs()
														lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__			  =      config_and_inputs

														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
														lowerCAmelCase__			  =      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__							(				a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  Dict      		= True

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														lowerCAmelCase__			  =      FlaxBertModelTester(self  )







							@slow
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														# Only check this for base model, not necessary for all model classes.
														# This will also help speed-up tests.
														lowerCAmelCase__			  =      FlaxBertModel.from_pretrained('''bert-base-cased'''  )
														lowerCAmelCase__			  =      model(np.ones((1, 1)  )  )
														self.assertIsNotNone(lowerCamelCase_  )





"""simple docstring"""




import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
lowerCAmelCase				:     List[str]					=				object()

# For specifying empty leaf dict `{}`
lowerCAmelCase				:     Any					=				object()



def        a__					(							snake_case__			,					snake_case__							)		->   List[Any]:
			lowerCamelCase      =   tuple((re.compile(x + """$"""							) for x in qs)							)
			for i in range(len(snake_case__							) - len(snake_case__							) + 1							):
						lowerCamelCase      =   [x.match(snake_case__							) for x, y in zip(snake_case__			,					ks[i:]							)]
						if matches and all(snake_case__							):
									return True
			return False



def        a__					(							snake_case__							)		->   str:
			def replace(snake_case__			,					snake_case__							):
						for rule, replacement in rules:
									if _match(snake_case__			,					snake_case__							):
												return replacement
						return val

			return replace



def        a__					(							)		->   Union[str, Any]:
			return [
			    # embeddings
			    (("transformer", "wpe", "embedding"), P("""mp"""			,					snake_case__							)),
			    (("transformer", "wte", "embedding"), P("""mp"""			,					snake_case__							)),
			    # atention
			    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(snake_case__			,					"""mp"""							)),
			    (("attention", "out_proj", "kernel"), P("""mp"""			,					snake_case__							)),
			    (("attention", "out_proj", "bias"), None),
			    # mlp
			    (("mlp", "c_fc", "kernel"), P(snake_case__			,					"""mp"""							)),
			    (("mlp", "c_fc", "bias"), P("""mp"""							)),
			    (("mlp", "c_proj", "kernel"), P("""mp"""			,					snake_case__							)),
			    (("mlp", "c_proj", "bias"), None),
			    # layer norms
			    ((r"ln_\d+", "bias"), None),
			    ((r"\d+", r"ln_\d+", "scale"), None),
			    (("ln_f", "bias"), None),
			    (("ln_f", "scale"), None),
			]



def        a__					(							snake_case__							)		->   Optional[Any]:
			lowerCamelCase      =   _get_partition_rules()
			lowerCamelCase      =   _replacement_rules(snake_case__							)
			lowerCamelCase      =   {k: _unmatched for k in flatten_dict(snake_case__							)}
			lowerCamelCase      =   {k: replace(snake_case__			,					snake_case__							) for k, v in initd.items()}
			assert _unmatched not in result.values(), "Incomplete partition spec."
			return freeze(unflatten_dict(snake_case__							)							)



