






"""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        _SCREAMING_SNAKE_CASE						(  unittest.TestCase							):





       """simple docstring"""





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

            '''simple docstring'''

            a__          =  parent
            a__          =  batch_size
            a__          =  seq_length
            a__          =  is_training
            a__          =  use_attention_mask
            a__          =  use_token_type_ids
            a__          =  use_labels
            a__          =  vocab_size
            a__          =  hidden_size
            a__          =  num_hidden_layers
            a__          =  num_attention_heads
            a__          =  intermediate_size
            a__          =  hidden_act
            a__          =  hidden_dropout_prob
            a__          =  attention_probs_dropout_prob
            a__          =  max_position_embeddings
            a__          =  type_vocab_size
            a__          =  type_sequence_label_size
            a__          =  initializer_range
            a__          =  num_choices




       def  lowercase		(   self:     List[Any]    ):

            '''simple docstring'''

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

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

            a__          =  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_=UpperCamelCase__				,		)

            return config, input_ids, attention_mask





       def  lowercase		(   self:     int    ):

            '''simple docstring'''

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







@require_flax
class        _SCREAMING_SNAKE_CASE						(  UpperCamelCase_       , unittest.TestCase							):





       """simple docstring"""





       _SCREAMING_SNAKE_CASE    =(
           (
               FlaxDistilBertModel,
               FlaxDistilBertForMaskedLM,
               FlaxDistilBertForMultipleChoice,
               FlaxDistilBertForQuestionAnswering,
               FlaxDistilBertForSequenceClassification,
               FlaxDistilBertForTokenClassification,
               FlaxDistilBertForQuestionAnswering,
           )
           if is_flax_available()
           else ()
       )




       def  lowercase		(   self:     Any    ):

            '''simple docstring'''

            a__          =  FlaxDistilBertModelTester(self    )





       @slow
       def  lowercase		(   self:     Any    ):

            '''simple docstring'''

            for model_class_name in self.all_model_classes:
                 a__          =  model_class_name.from_pretrained('''distilbert-base-uncased'''    )
                 a__          =  model(np.ones((1, 1)    )    )
                 self.assertIsNotNone(UpperCamelCase__    )







@require_flax
class        _SCREAMING_SNAKE_CASE						(  unittest.TestCase							):





       """simple docstring"""





       @slow
       def  lowercase		(   self:     int    ):

            '''simple docstring'''

            a__          =  FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased'''    )
            a__          =  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
            a__          =  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
            a__          =  model(UpperCamelCase__				,		attention_mask=UpperCamelCase__    )[0]
            a__          =  (1, 11, 768)
            self.assertEqual(output.shape				,		UpperCamelCase__    )
            a__          =  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]				,		UpperCamelCase__				,		atol=1e-4    )    )










"""simple docstring"""







__a       :	List[Any]             = frozenset(
    [
        'prompt',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
        'cross_attention_kwargs',
    ]
)

__a       :	Union[str, Any]             = frozenset(['prompt', 'negative_prompt'])

__a       :	Any             = frozenset([])

__a       :	Union[str, Any]             = frozenset(['image'])

__a       :	Dict             = frozenset(
    [
        'image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	Dict             = frozenset(['image'])

__a       :	Dict             = frozenset(
    [
        'prompt',
        'image',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
    ]
)

__a       :	Optional[Any]             = frozenset(['prompt', 'image', 'negative_prompt'])

__a       :	List[Any]             = frozenset(
    [
        # Text guided image variation with an image mask
        'prompt',
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
    ]
)

__a       :	Union[str, Any]             = frozenset(['prompt', 'image', 'mask_image', 'negative_prompt'])

__a       :	Optional[Any]             = frozenset(
    [
        # image variation with an image mask
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	int             = frozenset(['image', 'mask_image'])

__a       :	Tuple             = frozenset(
    [
        'example_image',
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	Optional[Any]             = frozenset(['example_image', 'image', 'mask_image'])

__a       :	Optional[Any]             = frozenset(['class_labels'])

__a       :	Tuple             = frozenset(['class_labels'])

__a       :	int             = frozenset(['batch_size'])

__a       :	int             = frozenset([])

__a       :	Union[str, Any]             = frozenset(['batch_size'])

__a       :	Tuple             = frozenset([])

__a       :	Dict             = frozenset(
    [
        'prompt',
        'audio_length_in_s',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
        'cross_attention_kwargs',
    ]
)

__a       :	Dict             = frozenset(['prompt', 'negative_prompt'])
__a       :	Optional[int]             = frozenset(['input_tokens'])

__a       :	str             = frozenset(['input_tokens'])


