

'''simple docstring'''
from typing import TYPE_CHECKING

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


_lowerCamelCase      					=     {"""configuration_unispeech""": ["""UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP""", """UniSpeechConfig"""]}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					_lowerCamelCase      					=     [
					    """UNISPEECH_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """UniSpeechForCTC""",
					    """UniSpeechForPreTraining""",
					    """UniSpeechForSequenceClassification""",
					    """UniSpeechModel""",
					    """UniSpeechPreTrainedModel""",
					]

if TYPE_CHECKING:
					from .configuration_unispeech import UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP, UniSpeechConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_unispeech import (
										    UNISPEECH_PRETRAINED_MODEL_ARCHIVE_LIST,
										    UniSpeechForCTC,
										    UniSpeechForPreTraining,
										    UniSpeechForSequenceClassification,
										    UniSpeechModel,
										    UniSpeechPreTrainedModel,
										)

else:
					import sys

					_lowerCamelCase      					=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)







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





class 					__lowerCAmelCase	(   nn.Module				):
  _a		     =    42
  _a		     =    jnp.floataa




  def     A__	(							self   )    ->					int:
        '''simple docstring'''






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




  def __call__(							self				,  lowerCAmelCase   )    ->					Any:
        '''simple docstring'''






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





class 					__lowerCAmelCase	(   nn.Module				):
  _a		     =    42
  _a		     =    jnp.floataa




  def     A__	(							self   )    ->					Any:
        '''simple docstring'''






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




  def __call__(							self				,  lowerCAmelCase   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			=self.conv(lowerCAmelCase   )
        return hidden_states





class 					__lowerCAmelCase	(   nn.Module				):
  _a		     =    42
  _a		     =    None
  _a		     =    0.0
  _a		     =    None
  _a		     =    jnp.floataa




  def     A__	(							self   )    ->					Optional[int]:
        '''simple docstring'''






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

        _lowercase      			=nn.GroupNorm(num_groups=32				,  epsilon=1e-5   )
        _lowercase      			=nn.Conv(
            lowerCAmelCase				,  kernel_size=(3, 3)				,  strides=(1, 1)				,  padding=((1, 1), (1, 1))				,  dtype=self.dtype				,  )

        _lowercase      			=nn.Dense(lowerCAmelCase				,  dtype=self.dtype   )

        _lowercase      			=nn.GroupNorm(num_groups=32				,  epsilon=1e-5   )
        _lowercase      			=nn.Dropout(self.dropout_prob   )
        _lowercase      			=nn.Conv(
            lowerCAmelCase				,  kernel_size=(3, 3)				,  strides=(1, 1)				,  padding=((1, 1), (1, 1))				,  dtype=self.dtype				,  )

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

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




  def __call__(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase=True   )    ->					Tuple:
        '''simple docstring'''






        _lowercase      			=hidden_states
        _lowercase      			=self.norma(lowerCAmelCase   )
        _lowercase      			=nn.swish(lowerCAmelCase   )
        _lowercase      			=self.conva(lowerCAmelCase   )

        _lowercase      			=self.time_emb_proj(nn.swish(lowerCAmelCase   )   )
        _lowercase      			=jnp.expand_dims(jnp.expand_dims(lowerCAmelCase				,  1   )				,  1   )
        _lowercase      			=hidden_states + temb

        _lowercase      			=self.norma(lowerCAmelCase   )
        _lowercase      			=nn.swish(lowerCAmelCase   )
        _lowercase      			=self.dropout(lowerCAmelCase				,  lowerCAmelCase   )
        _lowercase      			=self.conva(lowerCAmelCase   )

        if self.conv_shortcut is not None:
              _lowercase      			=self.conv_shortcut(lowerCAmelCase   )

        return hidden_states + residual


