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


				'''simple docstring'''






				def __init__(		self						,      lowercase						,      lowercase=1_3						,      lowercase=7						,      lowercase=True						,      lowercase=True						,      lowercase=True						,      lowercase=True						,      lowercase=9_9						,      lowercase=3_2						,      lowercase=5						,      lowercase=4						,      lowercase=3_7						,      lowercase="gelu"						,      lowercase=0.1						,      lowercase=0.1						,      lowercase=5_1_2						,      lowercase=1_6						,      lowercase=2						,      lowercase=0.02						,      lowercase=4						,      ):



						"""simple docstring"""





						A_						:							Dict      = parent
						A_						:							List[str]      = batch_size
						A_						:							Dict      = seq_length
						A_						:							List[Any]      = is_training
						A_						:							Dict      = use_attention_mask
						A_						:							Tuple      = use_token_type_ids
						A_						:							Union[str, Any]      = use_labels
						A_						:							List[str]      = vocab_size
						A_						:							str      = hidden_size
						A_						:							int      = num_hidden_layers
						A_						:							List[Any]      = num_attention_heads
						A_						:							Optional[Any]      = intermediate_size
						A_						:							Union[str, Any]      = hidden_act
						A_						:							List[str]      = hidden_dropout_prob
						A_						:							str      = attention_probs_dropout_prob
						A_						:							Any      = max_position_embeddings
						A_						:							Any      = type_vocab_size
						A_						:							str      = type_sequence_label_size
						A_						:							str      = initializer_range
						A_						:							Any      = num_choices
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Dict      = ids_tensor([self.batch_size, self.seq_length]						,      self.vocab_size   )

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

						A_						:							Tuple      = 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     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Any      = self.prepare_config_and_inputs()
						A_  ,     A_  ,     A_						:							Optional[Any]      = config_and_inputs
						A_						:							Tuple      = {'input_ids': input_ids, 'attention_mask': attention_mask}
						return config, inputs_dict





@require_flax
class UpperCAmelCase		(   __A   ,			unittest.TestCase    ):


				'''simple docstring'''






				lowerCamelCase_        =   (
				    (
				        FlaxDistilBertModel,
				        FlaxDistilBertForMaskedLM,
				        FlaxDistilBertForMultipleChoice,
				        FlaxDistilBertForQuestionAnswering,
				        FlaxDistilBertForSequenceClassification,
				        FlaxDistilBertForTokenClassification,
				        FlaxDistilBertForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Union[str, Any]      = FlaxDistilBertModelTester(self   )







				@slow
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						for model_class_name in self.all_model_classes:
								A_						:							str      = model_class_name.from_pretrained('distilbert-base-uncased'   )
								A_						:							int      = model(np.ones((1, 1)   )   )
								self.assertIsNotNone(lowercase   )





@require_flax
class UpperCAmelCase		(   unittest.TestCase    ):


				'''simple docstring'''






				@slow
				def     lowerCAmelCase_  (		self   ):



						"""simple docstring"""





						A_						:							Union[str, Any]      = FlaxDistilBertModel.from_pretrained('distilbert-base-uncased'   )
						A_						:							Union[str, Any]      = np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]   )
						A_						:							Optional[Any]      = np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
						A_						:							List[str]      = model(lowercase						,      attention_mask=lowercase   )[0]
						A_						:							Any      = (1, 1_1, 7_6_8)
						self.assertEqual(output.shape						,      lowercase   )
						A_						:							int      = np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]   )

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







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


_UpperCAmelCase          =						logging.get_logger(__name__)

_UpperCAmelCase          =						{
    """shi-labs/dinat-mini-in1k-224""": """https://huggingface.co/shi-labs/dinat-mini-in1k-224/resolve/main/config.json""",
    # See all Dinat models at https://huggingface.co/models?filter=dinat
}





class UpperCAmelCase		(   __A   ,			__A    ):


				'''simple docstring'''






				lowerCamelCase_        =   '''dinat'''

				lowerCamelCase_        =   {
				    '''num_attention_heads''': '''num_heads''',
				    '''num_hidden_layers''': '''num_layers''',
				}
				def __init__(		self						,      lowercase=4						,      lowercase=3						,      lowercase=6_4						,      lowercase=[3, 4, 6, 5]						,      lowercase=[2, 4, 8, 1_6]						,      lowercase=7						,      lowercase=[[1, 8, 1], [1, 4, 1, 4], [1, 2, 1, 2, 1, 2], [1, 1, 1, 1, 1]]						,      lowercase=3.0						,      lowercase=True						,      lowercase=0.0						,      lowercase=0.0						,      lowercase=0.1						,      lowercase="gelu"						,      lowercase=0.02						,      lowercase=1E-5						,      lowercase=0.0						,      lowercase=None						,      lowercase=None						,      **lowercase						,      ):



						"""simple docstring"""





						super().__init__(**lowercase   )

						A_						:							List[Any]      = patch_size
						A_						:							List[str]      = num_channels
						A_						:							Any      = embed_dim
						A_						:							List[Any]      = depths
						A_						:							List[Any]      = len(lowercase   )
						A_						:							Optional[int]      = num_heads
						A_						:							Any      = kernel_size
						A_						:							Tuple      = dilations
						A_						:							Optional[int]      = mlp_ratio
						A_						:							Optional[Any]      = qkv_bias
						A_						:							Dict      = hidden_dropout_prob
						A_						:							str      = attention_probs_dropout_prob
						A_						:							Optional[Any]      = drop_path_rate
						A_						:							str      = hidden_act
						A_						:							Optional[Any]      = layer_norm_eps
						A_						:							Any      = initializer_range
						# we set the hidden_size attribute in order to make Dinat work with VisionEncoderDecoderModel
						# this indicates the channel dimension after the last stage of the model
						A_						:							Dict      = int(embed_dim * 2 ** (len(lowercase   ) - 1)   )
						A_						:							Optional[int]      = layer_scale_init_value
						A_						:							Optional[int]      = ['stem'] + [F'''stage{idx}''' for idx in range(1						,      len(lowercase   ) + 1   )]
						A_  ,     A_						:							Any      = get_aligned_output_features_output_indices(
						    out_features=lowercase						,      out_indices=lowercase						,      stage_names=self.stage_names   )






