




'''simple docstring'''







from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import (
       TF_MODEL_FOR_PRETRAINING_MAPPING,
       TFMobileBertForMaskedLM,
       TFMobileBertForMultipleChoice,
       TFMobileBertForNextSentencePrediction,
       TFMobileBertForPreTraining,
       TFMobileBertForQuestionAnswering,
       TFMobileBertForSequenceClassification,
       TFMobileBertForTokenClassification,
       TFMobileBertModel,
   )


@require_tf
class     _a						(			__a							,							__a							,							unittest.TestCase				):
       __a       :  Union[str, Any]            =      (
           (
               TFMobileBertModel,
               TFMobileBertForMaskedLM,
               TFMobileBertForNextSentencePrediction,
               TFMobileBertForPreTraining,
               TFMobileBertForQuestionAnswering,
               TFMobileBertForSequenceClassification,
               TFMobileBertForTokenClassification,
               TFMobileBertForMultipleChoice,
           )
           if is_tf_available()
           else ()
       )
       __a       :  Tuple            =      (
           {
               """feature-extraction""": TFMobileBertModel,
               """fill-mask""": TFMobileBertForMaskedLM,
               """question-answering""": TFMobileBertForQuestionAnswering,
               """text-classification""": TFMobileBertForSequenceClassification,
               """token-classification""": TFMobileBertForTokenClassification,
               """zero-shot""": TFMobileBertForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       __a       :  Tuple            =      False
       __a       :  List[str]            =      False
       def 					A   (			self		:			Dict     ,       lowercase		:			Any     ,       lowercase		:			Dict     ,       lowercase		:			List[str]=False       ):



           '''simple docstring'''

           UpperCAmelCase     =	super()._prepare_for_class(lowercase     ,       lowercase     ,       return_labels=lowercase       )

           if return_labels:
               if model_class in get_values(lowercase       ):
                   UpperCAmelCase     =	tf.zeros(self.model_tester.batch_size     ,       dtype=tf.intaa       )

           return inputs_dict


       class     _a						(			__a				):
              def __init__(			self		:			List[str]     ,       lowercase		:			List[str]     ,       lowercase		:			Union[str, Any]=13     ,       lowercase		:			int=7     ,       lowercase		:			List[str]=True     ,       lowercase		:			Any=True     ,       lowercase		:			Optional[Any]=True     ,       lowercase		:			Any=True     ,       lowercase		:			int=99     ,       lowercase		:			Optional[Any]=32     ,       lowercase		:			Tuple=32     ,       lowercase		:			int=2     ,       lowercase		:			List[Any]=4     ,       lowercase		:			str=37     ,       lowercase		:			Any="gelu"     ,       lowercase		:			int=0.1     ,       lowercase		:			int=0.1     ,       lowercase		:			int=512     ,       lowercase		:			Union[str, Any]=16     ,       lowercase		:			Dict=2     ,       lowercase		:			Dict=0.02     ,       lowercase		:			Any=3     ,       lowercase		:			Dict=4     ,       lowercase		:			str=None     ,       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	parent
                  UpperCAmelCase     =	batch_size
                  UpperCAmelCase     =	seq_length
                  UpperCAmelCase     =	is_training
                  UpperCAmelCase     =	use_input_mask
                  UpperCAmelCase     =	use_token_type_ids
                  UpperCAmelCase     =	use_labels
                  UpperCAmelCase     =	vocab_size
                  UpperCAmelCase     =	hidden_size
                  UpperCAmelCase     =	num_hidden_layers
                  UpperCAmelCase     =	num_attention_heads
                  UpperCAmelCase     =	intermediate_size
                  UpperCAmelCase     =	hidden_act
                  UpperCAmelCase     =	hidden_dropout_prob
                  UpperCAmelCase     =	attention_probs_dropout_prob
                  UpperCAmelCase     =	max_position_embeddings
                  UpperCAmelCase     =	type_vocab_size
                  UpperCAmelCase     =	type_sequence_label_size
                  UpperCAmelCase     =	initializer_range
                  UpperCAmelCase     =	num_labels
                  UpperCAmelCase     =	num_choices
                  UpperCAmelCase     =	scope
                  UpperCAmelCase     =	embedding_size
              def 					A   (			self		:			str       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	ids_tensor([self.batch_size, self.seq_length]     ,       self.vocab_size       )

                  UpperCAmelCase     =	None
                  if self.use_input_mask:
                      UpperCAmelCase     =	random_attention_mask([self.batch_size, self.seq_length]       )

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

                  UpperCAmelCase     =	None
                  UpperCAmelCase     =	None
                  UpperCAmelCase     =	None
                  if self.use_labels:
                      UpperCAmelCase     =	ids_tensor([self.batch_size]     ,       self.type_sequence_label_size       )
                      UpperCAmelCase     =	ids_tensor([self.batch_size, self.seq_length]     ,       self.num_labels       )
                      UpperCAmelCase     =	ids_tensor([self.batch_size]     ,       self.num_choices       )

                  UpperCAmelCase     =	MobileBertConfig(
                      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     ,       initializer_range=self.initializer_range     ,       embedding_size=self.embedding_size     ,       )

                  return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
              def 					A   (			self		:			Any     ,       lowercase		:			Optional[int]     ,       lowercase		:			List[Any]     ,       lowercase		:			List[Any]     ,       lowercase		:			Tuple     ,       lowercase		:			Dict     ,       lowercase		:			List[str]     ,       lowercase		:			int       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	TFMobileBertModel(config=lowercase       )
                  UpperCAmelCase     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
                  UpperCAmelCase     =	model(lowercase       )

                  UpperCAmelCase     =	[input_ids, input_mask]
                  UpperCAmelCase     =	model(lowercase       )

                  UpperCAmelCase     =	model(lowercase       )

                  self.parent.assertEqual(
                      result.last_hidden_state.shape     ,       (self.batch_size, self.seq_length, self.hidden_size)       )
                  self.parent.assertEqual(result.pooler_output.shape     ,       (self.batch_size, self.hidden_size)       )
              def 					A   (			self		:			List[Any]     ,       lowercase		:			Optional[Any]     ,       lowercase		:			Tuple     ,       lowercase		:			List[str]     ,       lowercase		:			Optional[int]     ,       lowercase		:			Optional[Any]     ,       lowercase		:			Any     ,       lowercase		:			Optional[Any]       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	TFMobileBertForMaskedLM(config=lowercase       )
                  UpperCAmelCase     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
                  UpperCAmelCase     =	model(lowercase       )
                  self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, self.seq_length, self.vocab_size)       )
              def 					A   (			self		:			Optional[Any]     ,       lowercase		:			List[Any]     ,       lowercase		:			Any     ,       lowercase		:			str     ,       lowercase		:			List[Any]     ,       lowercase		:			List[Any]     ,       lowercase		:			str     ,       lowercase		:			Tuple       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	TFMobileBertForNextSentencePrediction(config=lowercase       )
                  UpperCAmelCase     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
                  UpperCAmelCase     =	model(lowercase       )
                  self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, 2)       )
              def 					A   (			self		:			Union[str, Any]     ,       lowercase		:			str     ,       lowercase		:			Tuple     ,       lowercase		:			List[str]     ,       lowercase		:			Any     ,       lowercase		:			str     ,       lowercase		:			Optional[int]     ,       lowercase		:			Union[str, Any]       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	TFMobileBertForPreTraining(config=lowercase       )
                  UpperCAmelCase     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
                  UpperCAmelCase     =	model(lowercase       )
                  self.parent.assertEqual(
                      result.prediction_logits.shape     ,       (self.batch_size, self.seq_length, self.vocab_size)       )
                  self.parent.assertEqual(result.seq_relationship_logits.shape     ,       (self.batch_size, 2)       )
              def 					A   (			self		:			Optional[Any]     ,       lowercase		:			Optional[int]     ,       lowercase		:			Tuple     ,       lowercase		:			Dict     ,       lowercase		:			str     ,       lowercase		:			Optional[int]     ,       lowercase		:			Union[str, Any]     ,       lowercase		:			Any       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	self.num_labels
                  UpperCAmelCase     =	TFMobileBertForSequenceClassification(config=lowercase       )
                  UpperCAmelCase     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
                  UpperCAmelCase     =	model(lowercase       )
                  self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, self.num_labels)       )
              def 					A   (			self		:			Union[str, Any]     ,       lowercase		:			Optional[Any]     ,       lowercase		:			str     ,       lowercase		:			Union[str, Any]     ,       lowercase		:			Union[str, Any]     ,       lowercase		:			Any     ,       lowercase		:			List[str]     ,       lowercase		:			Optional[int]       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	self.num_choices
                  UpperCAmelCase     =	TFMobileBertForMultipleChoice(config=lowercase       )
                  UpperCAmelCase     =	tf.tile(tf.expand_dims(lowercase     ,       1       )     ,       (1, self.num_choices, 1)       )
                  UpperCAmelCase     =	tf.tile(tf.expand_dims(lowercase     ,       1       )     ,       (1, self.num_choices, 1)       )
                  UpperCAmelCase     =	tf.tile(tf.expand_dims(lowercase     ,       1       )     ,       (1, self.num_choices, 1)       )
                  UpperCAmelCase     =	{
                      '''input_ids''': multiple_choice_inputs_ids,
                      '''attention_mask''': multiple_choice_input_mask,
                      '''token_type_ids''': multiple_choice_token_type_ids,
                  }
                  UpperCAmelCase     =	model(lowercase       )
                  self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, self.num_choices)       )
              def 					A   (			self		:			List[Any]     ,       lowercase		:			Optional[Any]     ,       lowercase		:			int     ,       lowercase		:			int     ,       lowercase		:			Dict     ,       lowercase		:			Any     ,       lowercase		:			Any     ,       lowercase		:			Any       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	self.num_labels
                  UpperCAmelCase     =	TFMobileBertForTokenClassification(config=lowercase       )
                  UpperCAmelCase     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
                  UpperCAmelCase     =	model(lowercase       )
                  self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, self.seq_length, self.num_labels)       )
              def 					A   (			self		:			List[Any]     ,       lowercase		:			Optional[Any]     ,       lowercase		:			Optional[Any]     ,       lowercase		:			Dict     ,       lowercase		:			List[Any]     ,       lowercase		:			Optional[Any]     ,       lowercase		:			Tuple     ,       lowercase		:			str       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	TFMobileBertForQuestionAnswering(config=lowercase       )
                  UpperCAmelCase     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
                  UpperCAmelCase     =	model(lowercase       )
                  self.parent.assertEqual(result.start_logits.shape     ,       (self.batch_size, self.seq_length)       )
                  self.parent.assertEqual(result.end_logits.shape     ,       (self.batch_size, self.seq_length)       )







              def 					A   (			self		:			str       ):



                  '''simple docstring'''

                  UpperCAmelCase     =	self.prepare_config_and_inputs()
                  (
                      (
                      UpperCAmelCase
                  ) ,  (
                      UpperCAmelCase
                  ) ,  (
                      UpperCAmelCase
                  ) ,  (
                      UpperCAmelCase
                  ) ,  (
                      UpperCAmelCase
                  ) ,  (
                      UpperCAmelCase
                  ) ,  (
                      UpperCAmelCase
                  ) ,  
                  )     =	config_and_inputs
                  UpperCAmelCase     =	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
                  return config, inputs_dict






       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	TFMobileBertModelTest.TFMobileBertModelTester(self       )
           UpperCAmelCase     =	ConfigTester(self     ,       config_class=lowercase     ,       hidden_size=37       )
       def 					A   (			self		:			str       ):



           '''simple docstring'''

           self.config_tester.run_common_tests()
       def 					A   (			self		:			Dict       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_model(*lowercase       )
       def 					A   (			self		:			str       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_masked_lm(*lowercase       )
       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_multiple_choice(*lowercase       )
       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*lowercase       )
       def 					A   (			self		:			str       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_pretraining(*lowercase       )
       def 					A   (			self		:			Dict       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_question_answering(*lowercase       )
       def 					A   (			self		:			Optional[Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_sequence_classification(*lowercase       )
       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_token_classification(*lowercase       )







       @slow
       def 					A   (			self		:			Dict       ):



           '''simple docstring'''

           for model_name in ["google/mobilebert-uncased"]:
               UpperCAmelCase     =	TFMobileBertModel.from_pretrained(lowercase       )
               self.assertIsNotNone(lowercase       )






@require_tf
class     _a						(			unittest.TestCase				):
       @slow
       def 					A   (			self		:			Optional[int]       ):



           '''simple docstring'''

           UpperCAmelCase     =	TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased'''       )
           UpperCAmelCase     =	tf.constant([[0, 1, 2, 3, 4, 5]]       )
           UpperCAmelCase     =	model(lowercase       )[0]

           UpperCAmelCase     =	[1, 6, 30_522]
           self.assertEqual(output.shape     ,       lowercase       )

           UpperCAmelCase     =	tf.constant(
               [
                   [
                       [-4.591_9547, -9.24_8295, -9.64_5256],
                       [-6.730_6175, -6.44_0284, -6.605_2837],
                       [-7.274_3506, -6.784_7915, -6.02_4673],
                   ]
               ]       )
           tf.debugging.assert_near(output[:, :3, :3]     ,       lowercase     ,       atol=1E-4       )










'''simple docstring'''







import unittest

import numpy as np

from transformers import RobertaPreLayerNormConfig, 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():
   import jax.numpy as jnp

   from transformers.models.roberta_prelayernorm.modeling_flax_roberta_prelayernorm import (
       FlaxRobertaPreLayerNormForCausalLM,
       FlaxRobertaPreLayerNormForMaskedLM,
       FlaxRobertaPreLayerNormForMultipleChoice,
       FlaxRobertaPreLayerNormForQuestionAnswering,
       FlaxRobertaPreLayerNormForSequenceClassification,
       FlaxRobertaPreLayerNormForTokenClassification,
       FlaxRobertaPreLayerNormModel,
   )


class     _a						(			unittest.TestCase				):
       def __init__(			self		:			List[Any]     ,       lowercase		:			Dict     ,       lowercase		:			List[str]=13     ,       lowercase		:			str=7     ,       lowercase		:			List[str]=True     ,       lowercase		:			List[str]=True     ,       lowercase		:			Optional[Any]=True     ,       lowercase		:			Optional[Any]=True     ,       lowercase		:			Any=99     ,       lowercase		:			Any=32     ,       lowercase		:			Any=5     ,       lowercase		:			Tuple=4     ,       lowercase		:			List[Any]=37     ,       lowercase		:			List[Any]="gelu"     ,       lowercase		:			int=0.1     ,       lowercase		:			Any=0.1     ,       lowercase		:			Optional[int]=512     ,       lowercase		:			List[str]=16     ,       lowercase		:			Union[str, Any]=2     ,       lowercase		:			int=0.02     ,       lowercase		:			int=4     ,       ):



           '''simple docstring'''

           UpperCAmelCase     =	parent
           UpperCAmelCase     =	batch_size
           UpperCAmelCase     =	seq_length
           UpperCAmelCase     =	is_training
           UpperCAmelCase     =	use_attention_mask
           UpperCAmelCase     =	use_token_type_ids
           UpperCAmelCase     =	use_labels
           UpperCAmelCase     =	vocab_size
           UpperCAmelCase     =	hidden_size
           UpperCAmelCase     =	num_hidden_layers
           UpperCAmelCase     =	num_attention_heads
           UpperCAmelCase     =	intermediate_size
           UpperCAmelCase     =	hidden_act
           UpperCAmelCase     =	hidden_dropout_prob
           UpperCAmelCase     =	attention_probs_dropout_prob
           UpperCAmelCase     =	max_position_embeddings
           UpperCAmelCase     =	type_vocab_size
           UpperCAmelCase     =	type_sequence_label_size
           UpperCAmelCase     =	initializer_range
           UpperCAmelCase     =	num_choices
       def 					A   (			self		:			Any       ):



           '''simple docstring'''

           UpperCAmelCase     =	ids_tensor([self.batch_size, self.seq_length]     ,       self.vocab_size       )

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

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

           UpperCAmelCase     =	RobertaPreLayerNormConfig(
               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=lowercase     ,       initializer_range=self.initializer_range     ,       )

           return config, input_ids, token_type_ids, attention_mask
       def 					A   (			self		:			List[Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.prepare_config_and_inputs()
           UpperCAmelCase ,  UpperCAmelCase ,  UpperCAmelCase ,  UpperCAmelCase     =	config_and_inputs
           UpperCAmelCase     =	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
           return config, inputs_dict







       def 					A   (			self		:			Optional[Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	self.prepare_config_and_inputs()
           UpperCAmelCase ,  UpperCAmelCase ,  UpperCAmelCase ,  UpperCAmelCase     =	config_and_inputs

           UpperCAmelCase     =	True
           UpperCAmelCase     =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
           UpperCAmelCase     =	ids_tensor([self.batch_size, self.seq_length]     ,       vocab_size=2       )

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






@require_flax
# Copied from tests.models.roberta.test_modelling_flax_roberta.FlaxRobertaPreLayerNormModelTest with ROBERTA->ROBERTA_PRELAYERNORM,Roberta->RobertaPreLayerNorm,roberta-base->andreasmadsen/efficient_mlm_m0.40
class     _a						(			__a							,							unittest.TestCase				):
       __a       :  Any            =      True

       __a       :  str            =      (
           (
               FlaxRobertaPreLayerNormModel,
               FlaxRobertaPreLayerNormForCausalLM,
               FlaxRobertaPreLayerNormForMaskedLM,
               FlaxRobertaPreLayerNormForSequenceClassification,
               FlaxRobertaPreLayerNormForTokenClassification,
               FlaxRobertaPreLayerNormForMultipleChoice,
               FlaxRobertaPreLayerNormForQuestionAnswering,
           )
           if is_flax_available()
           else ()
       )
       def 					A   (			self		:			Any       ):



           '''simple docstring'''

           UpperCAmelCase     =	FlaxRobertaPreLayerNormModelTester(self       )







       @slow
       def 					A   (			self		:			Optional[int]       ):



           '''simple docstring'''

           for model_class_name in self.all_model_classes:
               UpperCAmelCase     =	model_class_name.from_pretrained('''andreasmadsen/efficient_mlm_m0.40'''     ,       from_pt=lowercase       )
               UpperCAmelCase     =	model(np.ones((1, 1)       )       )
               self.assertIsNotNone(lowercase       )






@require_flax
class     _a						(			unittest.TestCase				):
       @slow
       def 					A   (			self		:			Union[str, Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	FlaxRobertaPreLayerNormForMaskedLM.from_pretrained('''andreasmadsen/efficient_mlm_m0.40'''     ,       from_pt=lowercase       )

           UpperCAmelCase     =	np.array([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]     ,       dtype=jnp.intaa       )
           UpperCAmelCase     =	model(lowercase       )[0]
           UpperCAmelCase     =	[1, 11, 50_265]
           self.assertEqual(list(output.shape       )     ,       lowercase       )
           # compare the actual values for a slice.
           UpperCAmelCase     =	np.array(
               [[[40.4880, 18.0199, -5.2367], [-1.8877, -4.0885, 10.7085], [-2.2613, -5.6110, 7.2665]]]     ,       dtype=np.floataa       )
           self.assertTrue(np.allclose(output[:, :3, :3]     ,       lowercase     ,       atol=1E-4       )       )







       @slow
       def 					A   (			self		:			Optional[Any]       ):



           '''simple docstring'''

           UpperCAmelCase     =	FlaxRobertaPreLayerNormModel.from_pretrained('''andreasmadsen/efficient_mlm_m0.40'''     ,       from_pt=lowercase       )

           UpperCAmelCase     =	np.array([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]     ,       dtype=jnp.intaa       )
           UpperCAmelCase     =	model(lowercase       )[0]
           # compare the actual values for a slice.
           UpperCAmelCase     =	np.array(
               [[[0.0208, -0.0356, 0.0237], [-0.1569, -0.0411, -0.2626], [0.1879, 0.0125, -0.0089]]]     ,       dtype=np.floataa       )
           self.assertTrue(np.allclose(output[:, :3, :3]     ,       lowercase     ,       atol=1E-4       )       )




