


"""simple docstring"""





import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax.numpy as jnp

  from transformers.models.distilbert.modeling_flax_distilbert import (
      FlaxDistilBertForMaskedLM,
      FlaxDistilBertForMultipleChoice,
      FlaxDistilBertForQuestionAnswering,
      FlaxDistilBertForSequenceClassification,
      FlaxDistilBertForTokenClassification,
      FlaxDistilBertModel,
  )




class 							snake_case							(			unittest.TestCase ):

      '''simple docstring'''



      def __init__(							self       :      Tuple						,				__lowercase       :      List[str]						,				__lowercase       :      str=13						,				__lowercase       :      Union[str, Any]=7						,				__lowercase       :      Union[str, Any]=True						,				__lowercase       :      Union[str, Any]=True						,				__lowercase       :      Optional[int]=True						,				__lowercase       :      str=True						,				__lowercase       :      Union[str, Any]=99						,				__lowercase       :      Union[str, Any]=32						,				__lowercase       :      str=5						,				__lowercase       :      str=4						,				__lowercase       :      Optional[Any]=37						,				__lowercase       :      Union[str, Any]="gelu"						,				__lowercase       :      List[str]=0.1						,				__lowercase       :      Dict=0.1						,				__lowercase       :      Dict=512						,				__lowercase       :      Any=16						,				__lowercase       :      Optional[Any]=2						,				__lowercase       :      Any=0.0_2						,				__lowercase       :      List[str]=4						,				):


        '''simple docstring'''




        __UpperCAmelCase    :						List[Any]       =						parent
        __UpperCAmelCase    :						str       =						batch_size
        __UpperCAmelCase    :						Optional[int]       =						seq_length
        __UpperCAmelCase    :						str       =						is_training
        __UpperCAmelCase    :						str       =						use_attention_mask
        __UpperCAmelCase    :						int       =						use_token_type_ids
        __UpperCAmelCase    :						List[Any]       =						use_labels
        __UpperCAmelCase    :						str       =						vocab_size
        __UpperCAmelCase    :						Dict       =						hidden_size
        __UpperCAmelCase    :						Tuple       =						num_hidden_layers
        __UpperCAmelCase    :						str       =						num_attention_heads
        __UpperCAmelCase    :						Any       =						intermediate_size
        __UpperCAmelCase    :						Optional[Any]       =						hidden_act
        __UpperCAmelCase    :						List[str]       =						hidden_dropout_prob
        __UpperCAmelCase    :						Tuple       =						attention_probs_dropout_prob
        __UpperCAmelCase    :						Tuple       =						max_position_embeddings
        __UpperCAmelCase    :						Dict       =						type_vocab_size
        __UpperCAmelCase    :						List[Any]       =						type_sequence_label_size
        __UpperCAmelCase    :						Dict       =						initializer_range
        __UpperCAmelCase    :						List[Any]       =						num_choices

      def       A_					(							self       :      int ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[Any]       =						ids_tensor([self.batch_size, self.seq_length]						,				self.vocab_size )

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

        __UpperCAmelCase    :						Optional[Any]       =						DistilBertConfig(
            vocab_size=self.vocab_size						,				dim=self.hidden_size						,				n_layers=self.num_hidden_layers						,				n_heads=self.num_attention_heads						,				hidden_dim=self.intermediate_size						,				hidden_act=self.hidden_act						,				dropout=self.hidden_dropout_prob						,				attention_dropout=self.attention_probs_dropout_prob						,				max_position_embeddings=self.max_position_embeddings						,				initializer_range=self.initializer_range						,				tie_weights_=__lowercase						,				)

        return config, input_ids, attention_mask







      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[str]       =						self.prepare_config_and_inputs()
        __UpperCAmelCase       ,     __UpperCAmelCase       ,     __UpperCAmelCase    :						Dict       =						config_and_inputs
        __UpperCAmelCase    :						Tuple       =						{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
        return config, inputs_dict




@require_flax
class 							snake_case							(			lowercase__      ,  unittest.TestCase ):

      '''simple docstring'''



      _A     :       List[Any]         =       (
          (
              FlaxDistilBertModel,
              FlaxDistilBertForMaskedLM,
              FlaxDistilBertForMultipleChoice,
              FlaxDistilBertForQuestionAnswering,
              FlaxDistilBertForSequenceClassification,
              FlaxDistilBertForTokenClassification,
              FlaxDistilBertForQuestionAnswering,
          )
          if is_flax_available()
          else ()
      )

      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




        __UpperCAmelCase    :						List[Any]       =						FlaxDistilBertModelTester(self )







      @slow
      def       A_					(							self       :      int ):


        '''simple docstring'''




        for model_class_name in self.all_model_classes:
          __UpperCAmelCase    :						Any       =						model_class_name.from_pretrained('''distilbert-base-uncased''' )
          __UpperCAmelCase    :						Optional[int]       =						model(np.ones((1, 1) ) )
          self.assertIsNotNone(__lowercase )




@require_flax
class 							snake_case							(			unittest.TestCase ):

      '''simple docstring'''



      @slow
      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased''' )
        __UpperCAmelCase    :						Optional[Any]       =						np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]] )
        __UpperCAmelCase    :						Any       =						np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
        __UpperCAmelCase    :						str       =						model(__lowercase						,				attention_mask=__lowercase )[0]
        __UpperCAmelCase    :						str       =						(1, 11, 768)
        self.assertEqual(output.shape						,				__lowercase )
        __UpperCAmelCase    :						List[str]       =						np.array([[[-0.1_6_3_9, 0.3_2_9_9, 0.1_6_4_8], [-0.1_7_4_6, 0.3_2_8_9, 0.1_7_1_0], [-0.1_8_8_4, 0.3_3_5_7, 0.1_8_1_0]]] )

        self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]						,				__lowercase						,				atol=1e-4 ) )






"""simple docstring"""





def       lowerCamelCase_       (   UpperCAmelCase_ = 10_00  )						->int:





  """simple docstring"""


  return sum(2 * a * ((a - 1) // 2) for a in range(3  ,			n + 1  )  )


if __name__ == "__main__":
  print(solution())