



'''simple docstring'''



import flax.linen as nn
import jax
import jax.numpy as jnp
class 	lowerCamelCase_							(					nn.Module			):
						_lowerCAmelCase     :   int
						_lowerCAmelCase     :   jnp.dtype	     =				jnp.floataa







						def 					__lowercase    (  self   :      List[Any]				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      List[str]     					=			nn.Conv(
													    self.out_channels				,    kernel_size=(3, 3)				,    strides=(1, 1)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )







						def __call__(  self   :      Optional[Any]				,    lowerCAmelCase__   :      Dict				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					,SCREAMING_SNAKE_CASE					,SCREAMING_SNAKE_CASE					,SCREAMING_SNAKE_CASE					:      Any     					=			hidden_states.shape
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			jax.image.resize(
													    lowerCAmelCase__				,    shape=(batch, height * 2, width * 2, channels)				,    method='''nearest'''				,    )
													SCREAMING_SNAKE_CASE					:      Dict     					=			self.conv(lowerCAmelCase__				)
													return hidden_states


class 	lowerCamelCase_							(					nn.Module			):
						_lowerCAmelCase     :   int
						_lowerCAmelCase     :   jnp.dtype	     =				jnp.floataa







						def 					__lowercase    (  self   :      Optional[Any]				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      List[str]     					=			nn.Conv(
													    self.out_channels				,    kernel_size=(3, 3)				,    strides=(2, 2)				,    padding=((1, 1), (1, 1))				,    dtype=self.dtype				,    )







						def __call__(  self   :      Any				,    lowerCAmelCase__   :      Optional[int]				):

													"""simple docstring"""
													# pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
													# hidden_states = jnp.pad(hidden_states, pad_width=pad)
													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			self.conv(lowerCAmelCase__				)
													return hidden_states


class 	lowerCamelCase_							(					nn.Module			):
						_lowerCAmelCase     :   int
						_lowerCAmelCase     :   int	     =				None
						_lowerCAmelCase     :   float	     =				0.0
						_lowerCAmelCase     :   bool	     =				None
						_lowerCAmelCase     :   jnp.dtype	     =				jnp.floataa







						def 					__lowercase    (  self   :      Optional[int]				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			self.in_channels if self.out_channels is None else self.out_channels

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

													SCREAMING_SNAKE_CASE					:      List[str]     					=			nn.Dense(lowerCAmelCase__				,    dtype=self.dtype				)

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

													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

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







						def __call__(  self   :      Any				,    lowerCAmelCase__   :      Optional[Any]				,    lowerCAmelCase__   :      Optional[Any]				,    lowerCAmelCase__   :      str=True				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			hidden_states
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			self.norma(lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			nn.swish(lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			self.conva(lowerCAmelCase__				)

													SCREAMING_SNAKE_CASE					:      List[str]     					=			self.time_emb_proj(nn.swish(lowerCAmelCase__				)				)
													SCREAMING_SNAKE_CASE					:      Dict     					=			jnp.expand_dims(jnp.expand_dims(lowerCAmelCase__				,    1				)				,    1				)
													SCREAMING_SNAKE_CASE					:      Tuple     					=			hidden_states + temb

													SCREAMING_SNAKE_CASE					:      List[str]     					=			self.norma(lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      Tuple     					=			nn.swish(lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      Optional[int]     					=			self.dropout(lowerCAmelCase__				,    lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      Optional[int]     					=			self.conva(lowerCAmelCase__				)

													if self.conv_shortcut is not None:
																				SCREAMING_SNAKE_CASE					:      List[Any]     					=			self.conv_shortcut(lowerCAmelCase__				)

													return hidden_states + residual









'''simple docstring'''



from typing import TYPE_CHECKING

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


lowerCAmelCase_     :       Tuple								=						{
    'configuration_tapas': ['TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TapasConfig'],
    'tokenization_tapas': ['TapasTokenizer'],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCAmelCase_     :       Dict								=						[
							    'TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'TapasForMaskedLM',
							    'TapasForQuestionAnswering',
							    'TapasForSequenceClassification',
							    'TapasModel',
							    'TapasPreTrainedModel',
							    'load_tf_weights_in_tapas',
							]
try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCAmelCase_     :       str								=						[
							    'TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'TFTapasForMaskedLM',
							    'TFTapasForQuestionAnswering',
							    'TFTapasForSequenceClassification',
							    'TFTapasModel',
							    'TFTapasPreTrainedModel',
							]


if TYPE_CHECKING:
							from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
							from .tokenization_tapas import TapasTokenizer

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tapas import (
														    TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TapasForMaskedLM,
														    TapasForQuestionAnswering,
														    TapasForSequenceClassification,
														    TapasModel,
														    TapasPreTrainedModel,
														    load_tf_weights_in_tapas,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_tapas import (
														    TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TFTapasForMaskedLM,
														    TFTapasForQuestionAnswering,
														    TFTapasForSequenceClassification,
														    TFTapasModel,
														    TFTapasPreTrainedModel,
														)


else:
							import sys

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




