






'''simple docstring'''







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





class 						SCREAMING_SNAKE_CASE       (       nn.Module							):
			snake_case__	=					42
			snake_case__	=					jnp.floataa






			def       SCREAMING_SNAKE_CASE   (			self   :      str       )      ->      Dict:
							a_		:	List[Any]									=   nn.Conv(
							    self.out_channels	,    kernel_size=(3, 3)	,    strides=(1, 1)	,    padding=((1, 1), (1, 1))	,    dtype=self.dtype	,    )







			def __call__(			self   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      int       )      ->      List[Any]:
							a_					,      a_					,      a_					,      a_		:	Optional[Any]									=   hidden_states.shape
							a_		:	Any									=   jax.image.resize(
							    __SCREAMING_SNAKE_CASE	,    shape=(batch, height * 2, width * 2, channels)	,    method='''nearest'''	,    )
							a_		:	Any									=   self.conv(__SCREAMING_SNAKE_CASE       )
							return hidden_states





class 						SCREAMING_SNAKE_CASE       (       nn.Module							):
			snake_case__	=					42
			snake_case__	=					jnp.floataa






			def       SCREAMING_SNAKE_CASE   (			self   :      Tuple       )      ->      Optional[int]:
							a_		:	Optional[int]									=   nn.Conv(
							    self.out_channels	,    kernel_size=(3, 3)	,    strides=(2, 2)	,    padding=((1, 1), (1, 1))	,    dtype=self.dtype	,    )







			def __call__(			self   :      Any	,    __SCREAMING_SNAKE_CASE   :      Union[str, Any]       )      ->      Union[str, Any]:
							# pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
							# hidden_states = jnp.pad(hidden_states, pad_width=pad)
							a_		:	Union[str, Any]									=   self.conv(__SCREAMING_SNAKE_CASE       )
							return hidden_states





class 						SCREAMING_SNAKE_CASE       (       nn.Module							):
			snake_case__	=					42
			snake_case__	=					None
			snake_case__	=					0.0
			snake_case__	=					None
			snake_case__	=					jnp.floataa






			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[int]       )      ->      List[str]:
							a_		:	Optional[int]									=   self.in_channels if self.out_channels is None else self.out_channels

							a_		:	Tuple									=   nn.GroupNorm(num_groups=32	,    epsilon=1e-5       )
							a_		:	Optional[Any]									=   nn.Conv(
							    __SCREAMING_SNAKE_CASE	,    kernel_size=(3, 3)	,    strides=(1, 1)	,    padding=((1, 1), (1, 1))	,    dtype=self.dtype	,    )

							a_		:	Optional[Any]									=   nn.Dense(__SCREAMING_SNAKE_CASE	,    dtype=self.dtype       )

							a_		:	Tuple									=   nn.GroupNorm(num_groups=32	,    epsilon=1e-5       )
							a_		:	Union[str, Any]									=   nn.Dropout(self.dropout_prob       )
							a_		:	Optional[Any]									=   nn.Conv(
							    __SCREAMING_SNAKE_CASE	,    kernel_size=(3, 3)	,    strides=(1, 1)	,    padding=((1, 1), (1, 1))	,    dtype=self.dtype	,    )

							a_		:	int									=   self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

							a_		:	Any									=   None
							if use_nin_shortcut:
											a_		:	Optional[Any]									=   nn.Conv(
											    __SCREAMING_SNAKE_CASE	,    kernel_size=(1, 1)	,    strides=(1, 1)	,    padding='''VALID'''	,    dtype=self.dtype	,    )







			def __call__(			self   :      Dict	,    __SCREAMING_SNAKE_CASE   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      Union[str, Any]	,    __SCREAMING_SNAKE_CASE   :      int=True       )      ->      str:
							a_		:	Union[str, Any]									=   hidden_states
							a_		:	int									=   self.norma(__SCREAMING_SNAKE_CASE       )
							a_		:	Optional[Any]									=   nn.swish(__SCREAMING_SNAKE_CASE       )
							a_		:	Union[str, Any]									=   self.conva(__SCREAMING_SNAKE_CASE       )

							a_		:	List[Any]									=   self.time_emb_proj(nn.swish(__SCREAMING_SNAKE_CASE       )       )
							a_		:	List[str]									=   jnp.expand_dims(jnp.expand_dims(__SCREAMING_SNAKE_CASE	,    1       )	,    1       )
							a_		:	int									=   hidden_states + temb

							a_		:	Tuple									=   self.norma(__SCREAMING_SNAKE_CASE       )
							a_		:	List[str]									=   nn.swish(__SCREAMING_SNAKE_CASE       )
							a_		:	Dict									=   self.dropout(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
							a_		:	List[Any]									=   self.conva(__SCREAMING_SNAKE_CASE       )

							if self.conv_shortcut is not None:
											a_		:	int									=   self.conv_shortcut(__SCREAMING_SNAKE_CASE       )

							return hidden_states + residual












'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__lowerCAmelCase									=   {'configuration_sew': ['SEW_PRETRAINED_CONFIG_ARCHIVE_MAP', 'SEWConfig']}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase									=   [
		    'SEW_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'SEWForCTC',
		    'SEWForSequenceClassification',
		    'SEWModel',
		    'SEWPreTrainedModel',
		]

if TYPE_CHECKING:
		from .configuration_sew import SEW_PRETRAINED_CONFIG_ARCHIVE_MAP, SEWConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_sew import (
				    SEW_PRETRAINED_MODEL_ARCHIVE_LIST,
				    SEWForCTC,
				    SEWForSequenceClassification,
				    SEWModel,
				    SEWPreTrainedModel,
				)


else:
		import sys

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




