


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

class 						A__							(		nn.Module  ):
						UpperCAmelCase             =						42
						UpperCAmelCase             =						jnp.floataa




						def 							__UpperCamelCase				(       self					:      int  )							->     List[str]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =nn.Conv(
										    self.out_channels				,    kernel_size=(3, 3)				,    strides=(1, 1)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )




						def __call__(       self					:      Union[str, Any]				,    _a					:      Tuple  )							->     Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE			  =hidden_states.shape
										_SCREAMING_SNAKE_CASE			  =jax.image.resize(
										    _a				,    shape=(batch, height * 2, width * 2, channels)				,    method='''nearest'''				,    )
										_SCREAMING_SNAKE_CASE			  =self.conv(_a  )
										return hidden_states

class 						A__							(		nn.Module  ):
						UpperCAmelCase             =						42
						UpperCAmelCase             =						jnp.floataa




						def 							__UpperCamelCase				(       self					:      int  )							->     Dict:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =nn.Conv(
										    self.out_channels				,    kernel_size=(3, 3)				,    strides=(2, 2)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )




						def __call__(       self					:      int				,    _a					:      Dict  )							->     List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =self.conv(_a  )
										return hidden_states

class 						A__							(		nn.Module  ):
						UpperCAmelCase             =						42
						UpperCAmelCase             =						None
						UpperCAmelCase             =						0.0
						UpperCAmelCase             =						None
						UpperCAmelCase             =						jnp.floataa




						def 							__UpperCamelCase				(       self					:      Tuple  )							->     Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =self.in_channels if self.out_channels is None else self.out_channels

										_SCREAMING_SNAKE_CASE			  =nn.GroupNorm(num_groups=32				,    epsilon=1E-5  )
										_SCREAMING_SNAKE_CASE			  =nn.Conv(
										    _a				,    kernel_size=(3, 3)				,    strides=(1, 1)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )

										_SCREAMING_SNAKE_CASE			  =nn.Dense(_a				,    dtype=self.dtype  )

										_SCREAMING_SNAKE_CASE			  =nn.GroupNorm(num_groups=32				,    epsilon=1E-5  )
										_SCREAMING_SNAKE_CASE			  =nn.Dropout(self.dropout_prob  )
										_SCREAMING_SNAKE_CASE			  =nn.Conv(
										    _a				,    kernel_size=(3, 3)				,    strides=(1, 1)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )

										_SCREAMING_SNAKE_CASE			  =self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

										_SCREAMING_SNAKE_CASE			  =None
										if use_nin_shortcut:
														_SCREAMING_SNAKE_CASE			  =nn.Conv(
														    _a				,    kernel_size=(1, 1)				,    strides=(1, 1)				,    padding='''VALID'''				,    dtype=self.dtype				,    )




						def __call__(       self					:      str				,    _a					:      str				,    _a					:      Tuple				,    _a					:      Dict=True  )							->     Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE			  =hidden_states
										_SCREAMING_SNAKE_CASE			  =self.norma(_a  )
										_SCREAMING_SNAKE_CASE			  =nn.swish(_a  )
										_SCREAMING_SNAKE_CASE			  =self.conva(_a  )

										_SCREAMING_SNAKE_CASE			  =self.time_emb_proj(nn.swish(_a  )  )
										_SCREAMING_SNAKE_CASE			  =jnp.expand_dims(jnp.expand_dims(_a				,    1  )				,    1  )
										_SCREAMING_SNAKE_CASE			  =hidden_states + temb

										_SCREAMING_SNAKE_CASE			  =self.norma(_a  )
										_SCREAMING_SNAKE_CASE			  =nn.swish(_a  )
										_SCREAMING_SNAKE_CASE			  =self.dropout(_a				,    _a  )
										_SCREAMING_SNAKE_CASE			  =self.conva(_a  )

										if self.conv_shortcut is not None:
														_SCREAMING_SNAKE_CASE			  =self.conv_shortcut(_a  )

										return hidden_states + residual



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_speech_available, is_torch_available


snake_case_						:			Optional[Any]													=							{
    '''configuration_audio_spectrogram_transformer''': [
        '''AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''ASTConfig''',
    ]
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							snake_case_						:			Optional[Any]													=							[
							    '''AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
							    '''ASTForAudioClassification''',
							    '''ASTModel''',
							    '''ASTPreTrainedModel''',
							]

try:
							if not is_speech_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							snake_case_						:			List[Any]													=							['''ASTFeatureExtractor''']

if TYPE_CHECKING:
							from .configuration_audio_spectrogram_transformer import (
							    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    ASTConfig,
							)

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_audio_spectrogram_transformer import (
														    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
														    ASTForAudioClassification,
														    ASTModel,
														    ASTPreTrainedModel,
														)

							try:
														if not is_speech_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .feature_extraction_audio_spectrogram_transformer import ASTFeatureExtractor


else:
							import sys

							snake_case_						:			Optional[Any]													=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)