




'''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 	lowerCAmelCase_(    unittest.TestCase					):
  '''simple docstring'''




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




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

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

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

      lowerCAmelCase__			:  List[str]											= 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=__UpperCAmelCase  ,initializer_range=self.initializer_range  ,)

      return config, input_ids, token_type_ids, attention_mask




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






  def 							UpperCAmelCase_     (   self				)   ->		Any:
      lowerCAmelCase__			:  List[Any]											= self.prepare_config_and_inputs()
      lowerCAmelCase__			:  Dict											= config_and_inputs

      lowerCAmelCase__			:  Union[str, Any]											= True
      lowerCAmelCase__			:  Optional[Any]											= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
      lowerCAmelCase__			:  Tuple											= 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 	lowerCAmelCase_(    SCREAMING_SNAKE_CASE_						,			unittest.TestCase					):
  '''simple docstring'''




  __lowercase							:				Any           =				True

  __lowercase							:				int           =				(
      (
          FlaxRobertaPreLayerNormModel,
          FlaxRobertaPreLayerNormForCausalLM,
          FlaxRobertaPreLayerNormForMaskedLM,
          FlaxRobertaPreLayerNormForSequenceClassification,
          FlaxRobertaPreLayerNormForTokenClassification,
          FlaxRobertaPreLayerNormForMultipleChoice,
          FlaxRobertaPreLayerNormForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )




  def 							UpperCAmelCase_     (   self				)   ->		Any:
      lowerCAmelCase__			:  str											= FlaxRobertaPreLayerNormModelTester(self				)






  @slow
  def 							UpperCAmelCase_     (   self				)   ->		List[str]:
      for model_class_name in self.all_model_classes:
          lowerCAmelCase__			:  int											= model_class_name.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""  ,from_pt=__UpperCAmelCase				)
          lowerCAmelCase__			:  Any											= model(np.ones((1, 1)				)				)
          self.assertIsNotNone(__UpperCAmelCase				)







@require_flax
class 	lowerCAmelCase_(    unittest.TestCase					):
  '''simple docstring'''




  @slow
  def 							UpperCAmelCase_     (   self				)   ->		str:
      lowerCAmelCase__			:  Any											= FlaxRobertaPreLayerNormForMaskedLM.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""  ,from_pt=__UpperCAmelCase				)

      lowerCAmelCase__			:  List[Any]											= np.array([[0, 3_1414, 232, 328, 740, 1140, 1_2695, 69, 4_6078, 1588, 2]]  ,dtype=jnp.intaa				)
      lowerCAmelCase__			:  Optional[int]											= model(__UpperCAmelCase				)[0]
      lowerCAmelCase__			:  Optional[Any]											= [1, 11, 5_0265]
      self.assertEqual(list(output.shape				)  ,__UpperCAmelCase				)
      # compare the actual values for a slice.
      lowerCAmelCase__			:  Dict											= np.array(
          [[[40.4880, 18.0199, -5.2_3_6_7], [-1.8_8_7_7, -4.0_8_8_5, 10.7085], [-2.2_6_1_3, -5.6_1_1_0, 7.2_6_6_5]]]  ,dtype=np.floataa				)
      self.assertTrue(np.allclose(output[:, :3, :3]  ,__UpperCAmelCase  ,atol=1E-4				)				)






  @slow
  def 							UpperCAmelCase_     (   self				)   ->		str:
      lowerCAmelCase__			:  Tuple											= FlaxRobertaPreLayerNormModel.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""  ,from_pt=__UpperCAmelCase				)

      lowerCAmelCase__			:  Union[str, Any]											= np.array([[0, 3_1414, 232, 328, 740, 1140, 1_2695, 69, 4_6078, 1588, 2]]  ,dtype=jnp.intaa				)
      lowerCAmelCase__			:  Any											= model(__UpperCAmelCase				)[0]
      # compare the actual values for a slice.
      lowerCAmelCase__			:  List[str]											= np.array(
          [[[0.0_2_0_8, -0.0_3_5_6, 0.0_2_3_7], [-0.1_5_6_9, -0.0_4_1_1, -0.2_6_2_6], [0.1_8_7_9, 0.0_1_2_5, -0.0_0_8_9]]]  ,dtype=np.floataa				)
      self.assertTrue(np.allclose(output[:, :3, :3]  ,__UpperCAmelCase  ,atol=1E-4				)				)







'''simple docstring'''





from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


_lowerCAmelCase										=     logging.get_logger(__name__)

_lowerCAmelCase										=     {
    '''facebook/convnextv2-tiny-1k-224''': '''https://huggingface.co/facebook/convnextv2-tiny-1k-224/resolve/main/config.json''',
}






class 	lowerCAmelCase_(    SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_					):
  '''simple docstring'''




  __lowercase							:				Tuple           =				'''convnextv2'''




  def __init__(   self  ,__UpperCAmelCase=3  ,__UpperCAmelCase=4  ,__UpperCAmelCase=4  ,__UpperCAmelCase=None  ,__UpperCAmelCase=None  ,__UpperCAmelCase="gelu"  ,__UpperCAmelCase=0.0_2  ,__UpperCAmelCase=1E-12  ,__UpperCAmelCase=0.0  ,__UpperCAmelCase=224  ,__UpperCAmelCase=None  ,__UpperCAmelCase=None  ,**__UpperCAmelCase  ,)   ->		Union[str, Any]:
      super().__init__(**__UpperCAmelCase				)

      lowerCAmelCase__			:  int											= num_channels
      lowerCAmelCase__			:  List[Any]											= patch_size
      lowerCAmelCase__			:  Union[str, Any]											= num_stages
      lowerCAmelCase__			:  Tuple											= [96, 192, 384, 768] if hidden_sizes is None else hidden_sizes
      lowerCAmelCase__			:  str											= [3, 3, 9, 3] if depths is None else depths
      lowerCAmelCase__			:  Optional[Any]											= hidden_act
      lowerCAmelCase__			:  str											= initializer_range
      lowerCAmelCase__			:  List[str]											= layer_norm_eps
      lowerCAmelCase__			:  Dict											= drop_path_rate
      lowerCAmelCase__			:  int											= image_size
      lowerCAmelCase__			:  int											= ["""stem"""] + [F"""stage{idx}""" for idx in range(1  ,len(self.depths				) + 1				)]
      lowerCAmelCase__				,       lowerCAmelCase__			:  Tuple											= get_aligned_output_features_output_indices(
          out_features=__UpperCAmelCase  ,out_indices=__UpperCAmelCase  ,stage_names=self.stage_names				)


