




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowerCAmelCase	:       List[Any]     =	{
    'configuration_distilbert': [
        'DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'DistilBertConfig',
        'DistilBertOnnxConfig',
    ],
    'tokenization_distilbert': ['DistilBertTokenizer'],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase	:       int     =	['DistilBertTokenizerFast']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase	:       Dict     =	[
		    'DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'DistilBertForMaskedLM',
		    'DistilBertForMultipleChoice',
		    'DistilBertForQuestionAnswering',
		    'DistilBertForSequenceClassification',
		    'DistilBertForTokenClassification',
		    'DistilBertModel',
		    'DistilBertPreTrainedModel',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase	:       Any     =	[
		    'TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'TFDistilBertForMaskedLM',
		    'TFDistilBertForMultipleChoice',
		    'TFDistilBertForQuestionAnswering',
		    'TFDistilBertForSequenceClassification',
		    'TFDistilBertForTokenClassification',
		    'TFDistilBertMainLayer',
		    'TFDistilBertModel',
		    'TFDistilBertPreTrainedModel',
		]

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase	:       Optional[int]     =	[
		    'FlaxDistilBertForMaskedLM',
		    'FlaxDistilBertForMultipleChoice',
		    'FlaxDistilBertForQuestionAnswering',
		    'FlaxDistilBertForSequenceClassification',
		    'FlaxDistilBertForTokenClassification',
		    'FlaxDistilBertModel',
		    'FlaxDistilBertPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_distilbert import (
		    DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    DistilBertConfig,
		    DistilBertOnnxConfig,
		)
		from .tokenization_distilbert import DistilBertTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_distilbert_fast import DistilBertTokenizerFast

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_distilbert import (
				    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DistilBertForMaskedLM,
				    DistilBertForMultipleChoice,
				    DistilBertForQuestionAnswering,
				    DistilBertForSequenceClassification,
				    DistilBertForTokenClassification,
				    DistilBertModel,
				    DistilBertPreTrainedModel,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_distilbert import (
				    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFDistilBertForMaskedLM,
				    TFDistilBertForMultipleChoice,
				    TFDistilBertForQuestionAnswering,
				    TFDistilBertForSequenceClassification,
				    TFDistilBertForTokenClassification,
				    TFDistilBertMainLayer,
				    TFDistilBertModel,
				    TFDistilBertPreTrainedModel,
				)

		try:
				if not is_flax_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_flax_distilbert import (
				    FlaxDistilBertForMaskedLM,
				    FlaxDistilBertForMultipleChoice,
				    FlaxDistilBertForQuestionAnswering,
				    FlaxDistilBertForSequenceClassification,
				    FlaxDistilBertForTokenClassification,
				    FlaxDistilBertModel,
				    FlaxDistilBertPreTrainedModel,
				)

else:
		import sys

		__lowerCAmelCase	:       str     =	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)












import math

import flax.linen as nn
import jax.numpy as jnp




def     __magic_name__       (				A	:      jnp.ndarray,						A	:      int,						A	:      float = 1,						A	:      float = 1,						A	:      float = 1.0E4,						A	:      bool = False,						A	:      float = 1.0,						):






			'''simple docstring'''







			assert timesteps.ndim == 1, "Timesteps should be a 1d-array"
			assert embedding_dim % 2 == 0, F"""Embedding dimension {embedding_dim} should be even"""
			a					=    float(embedding_dim // 2     )
			a					=    math.log(max_timescale / min_timescale     ) / (num_timescales - freq_shift)
			a					=    min_timescale * jnp.exp(jnp.arange(A,						dtype=jnp.floataa     ) * -log_timescale_increment     )
			a					=    jnp.expand_dims(A,						1     ) * jnp.expand_dims(A,						0     )

			# scale embeddings
			a					=    scale * emb

			if flip_sin_to_cos:
						a					=    jnp.concatenate([jnp.cos(A     ), jnp.sin(A     )],						axis=1     )
			else:
						a					=    jnp.concatenate([jnp.sin(A     ), jnp.cos(A     )],						axis=1     )
			a					=    jnp.reshape(A,						[jnp.shape(A     )[0], embedding_dim]     )
			return signal




class 					snake_case__							(nn.Module     ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_ :							int											=      32
				SCREAMING_SNAKE_CASE_ :							jnp.dtype											=      jnp.floataa





				@nn.compact
				def __call__(   self       :  Tuple      ,   __lowerCamelCase       :  Optional[Any]							)      ->  List[Any]:
							a					=    nn.Dense(self.time_embed_dim      ,   dtype=self.dtype      ,   name="linear_1"							)(__lowerCamelCase							)
							a					=    nn.silu(__lowerCamelCase							)
							a					=    nn.Dense(self.time_embed_dim      ,   dtype=self.dtype      ,   name="linear_2"							)(__lowerCamelCase							)
							return temb




class 					snake_case__							(nn.Module     ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_ :							int											=      32
				SCREAMING_SNAKE_CASE_ :							bool											=      False
				SCREAMING_SNAKE_CASE_ :							float											=      1





				@nn.compact
				def __call__(   self       :  Tuple      ,   __lowerCamelCase       :  int							)      ->  Union[str, Any]:
							return get_sinusoidal_embeddings(
							    __lowerCamelCase      ,   embedding_dim=self.dim      ,   flip_sin_to_cos=self.flip_sin_to_cos      ,   freq_shift=self.freq_shift							)






