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




				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.02							,							_UpperCamelCase=4							,							)							-> List[Any]:
							lowerCAmelCase_         =      parent
							lowerCAmelCase_         =      batch_size
							lowerCAmelCase_         =      seq_length
							lowerCAmelCase_         =      is_training
							lowerCAmelCase_         =      use_attention_mask
							lowerCAmelCase_         =      use_token_type_ids
							lowerCAmelCase_         =      use_labels
							lowerCAmelCase_         =      vocab_size
							lowerCAmelCase_         =      hidden_size
							lowerCAmelCase_         =      num_hidden_layers
							lowerCAmelCase_         =      num_attention_heads
							lowerCAmelCase_         =      intermediate_size
							lowerCAmelCase_         =      hidden_act
							lowerCAmelCase_         =      hidden_dropout_prob
							lowerCAmelCase_         =      attention_probs_dropout_prob
							lowerCAmelCase_         =      max_position_embeddings
							lowerCAmelCase_         =      type_vocab_size
							lowerCAmelCase_         =      type_sequence_label_size
							lowerCAmelCase_         =      initializer_range
							lowerCAmelCase_         =      num_choices




				def 	__a		(   self	)							-> int:
							lowerCAmelCase_         =      ids_tensor([self.batch_size, self.seq_length]							,							self.vocab_size	)

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

							lowerCAmelCase_         =      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 	__a		(   self	)							-> Optional[int]:
							lowerCAmelCase_         =      self.prepare_config_and_inputs()
							lowerCAmelCase_      ,    lowerCAmelCase_      ,    lowerCAmelCase_         =      config_and_inputs
							lowerCAmelCase_         =      {"input_ids": input_ids, "attention_mask": attention_mask}
							return config, inputs_dict





@require_flax
class       _lowerCAmelCase   (      __a   , unittest.TestCase			):
				_lowercase     						=(
				    (
				        FlaxDistilBertModel,
				        FlaxDistilBertForMaskedLM,
				        FlaxDistilBertForMultipleChoice,
				        FlaxDistilBertForQuestionAnswering,
				        FlaxDistilBertForSequenceClassification,
				        FlaxDistilBertForTokenClassification,
				        FlaxDistilBertForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)




				def 	__a		(   self	)							-> Any:
							lowerCAmelCase_         =      FlaxDistilBertModelTester(self	)




				@slow
				def 	__a		(   self	)							-> Any:
							for model_class_name in self.all_model_classes:
										lowerCAmelCase_         =      model_class_name.from_pretrained("distilbert-base-uncased"	)
										lowerCAmelCase_         =      model(np.ones((1, 1)	)	)
										self.assertIsNotNone(_UpperCamelCase	)





@require_flax
class       _lowerCAmelCase   (      unittest.TestCase			):




				@slow
				def 	__a		(   self	)							-> List[Any]:
							lowerCAmelCase_         =      FlaxDistilBertModel.from_pretrained("distilbert-base-uncased"	)
							lowerCAmelCase_         =      np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]	)
							lowerCAmelCase_         =      np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
							lowerCAmelCase_         =      model(_UpperCamelCase							,							attention_mask=_UpperCamelCase	)[0]
							lowerCAmelCase_         =      (1, 11, 768)
							self.assertEqual(output.shape							,							_UpperCamelCase	)
							lowerCAmelCase_         =      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]							,							_UpperCamelCase							,							atol=1e-4	)	)





from math import ceil






def   lowerCamelCase__							(							__lowerCAmelCase    :    Dict       ,				__lowerCAmelCase    :    int							):




			"""simple docstring"""


			lowerCAmelCase_         =      list(range(0       ,				__lowerCAmelCase							)							)

			lowerCAmelCase_         =      [item for sublist in list(device_map.values()							) for item in sublist]

			# Duplicate check
			lowerCAmelCase_         =      []
			for i in device_map_blocks:
						if device_map_blocks.count(__lowerCAmelCase							) > 1 and i not in duplicate_blocks:
									duplicate_blocks.append(__lowerCAmelCase							)
    # Missing blocks
			lowerCAmelCase_         =      [i for i in blocks if i not in device_map_blocks]
			lowerCAmelCase_         =      [i for i in device_map_blocks if i not in blocks]

			if len(__lowerCAmelCase							) != 0:
						raise ValueError(
						    "Duplicate attention blocks specified in device_map. Attention blocks must be specified to one device."
						    " These attention blocks were specified more than once: " + str(__lowerCAmelCase							)							)
			if len(__lowerCAmelCase							) != 0:
						raise ValueError(
						    "There are attention blocks for this model that are not specified in the device_map. Add these attention "
						    "blocks to a device on the device_map: " + str(__lowerCAmelCase							)							)
			if len(__lowerCAmelCase							) != 0:
						raise ValueError(
						    "The device_map contains more attention blocks than this model has. Remove these from the device_map:"
						    + str(__lowerCAmelCase							)							)






def   lowerCamelCase__							(							__lowerCAmelCase    :    List[str]       ,				__lowerCAmelCase    :    Tuple							):




			"""simple docstring"""


			lowerCAmelCase_         =      list(range(__lowerCAmelCase							)							)
			lowerCAmelCase_         =      int(ceil(n_layers / len(__lowerCAmelCase							)							)							)
			lowerCAmelCase_         =      [layers[i : i + n_blocks] for i in range(0       ,				__lowerCAmelCase       ,				__lowerCAmelCase							)]

			return dict(zip(__lowerCAmelCase       ,				__lowerCAmelCase							)							)



