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




class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






						_lowercase		:				int
						_lowercase		:				jnp.dtype									=       jnp.floataa




						def  _lowercase     ( self				):

												"""simple docstring"""



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




						def __call__( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase						,   _lowerCAmelCase						,   _lowerCAmelCase						,   _lowerCAmelCase  						=			hidden_states.shape
												_lowerCAmelCase  						=			jax.image.resize(
												    __lowerCAmelCase   ,       shape=(batch, height * 2, width * 2, channels)   ,       method="""nearest"""   ,       )
												_lowerCAmelCase  						=			self.conv(__lowerCAmelCase				)
												return hidden_states






class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






						_lowercase		:				int
						_lowercase		:				jnp.dtype									=       jnp.floataa




						def  _lowercase     ( self				):

												"""simple docstring"""



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




						def __call__( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.conv(__lowerCAmelCase				)
												return hidden_states






class   UpperCAmelCase_					(		nn.Module			):






						'''simple docstring'''






						_lowercase		:				int
						_lowercase		:				int									=       None
						_lowercase		:				float									=       0.0
						_lowercase		:				bool									=       None
						_lowercase		:				jnp.dtype									=       jnp.floataa




						def  _lowercase     ( self				):

												"""simple docstring"""



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

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

												_lowerCAmelCase  						=			nn.Dense(__lowerCAmelCase   ,       dtype=self.dtype				)

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

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

												_lowerCAmelCase  						=			None
												if use_nin_shortcut:
																		_lowerCAmelCase  						=			nn.Conv(
																		    __lowerCAmelCase   ,       kernel_size=(1, 1)   ,       strides=(1, 1)   ,       padding="""VALID"""   ,       dtype=self.dtype   ,       )




						def __call__( self   ,       _lowercase   ,       _lowercase   ,       _lowercase=True				):

												"""simple docstring"""



												_lowerCAmelCase  						=			hidden_states
												_lowerCAmelCase  						=			self.norma(__lowerCAmelCase				)
												_lowerCAmelCase  						=			nn.swish(__lowerCAmelCase				)
												_lowerCAmelCase  						=			self.conva(__lowerCAmelCase				)

												_lowerCAmelCase  						=			self.time_emb_proj(nn.swish(__lowerCAmelCase				)				)
												_lowerCAmelCase  						=			jnp.expand_dims(jnp.expand_dims(__lowerCAmelCase   ,       1				)   ,       1				)
												_lowerCAmelCase  						=			hidden_states + temb

												_lowerCAmelCase  						=			self.norma(__lowerCAmelCase				)
												_lowerCAmelCase  						=			nn.swish(__lowerCAmelCase				)
												_lowerCAmelCase  						=			self.dropout(__lowerCAmelCase   ,       __lowerCAmelCase				)
												_lowerCAmelCase  						=			self.conva(__lowerCAmelCase				)

												if self.conv_shortcut is not None:
																		_lowerCAmelCase  						=			self.conv_shortcut(__lowerCAmelCase				)

												return hidden_states + residual


'''simple docstring'''



from ...configuration_utils import PretrainedConfig




class   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE			):






						'''simple docstring'''






						_lowercase		:				Tuple									=       '''bert-generation'''




						def __init__( self   ,       _lowercase=50_358   ,       _lowercase=1_024   ,       _lowercase=24   ,       _lowercase=16   ,       _lowercase=4_096   ,       _lowercase="gelu"   ,       _lowercase=0.1   ,       _lowercase=0.1   ,       _lowercase=512   ,       _lowercase=0.02   ,       _lowercase=1e-12   ,       _lowercase=0   ,       _lowercase=2   ,       _lowercase=1   ,       _lowercase="absolute"   ,       _lowercase=True   ,       **_lowercase   ,       ):

												"""simple docstring"""



												super().__init__(pad_token_id=_lowercase   ,       bos_token_id=_lowercase   ,       eos_token_id=_lowercase   ,       **_lowercase				)

												_lowerCAmelCase  						=			vocab_size
												_lowerCAmelCase  						=			hidden_size
												_lowerCAmelCase  						=			num_hidden_layers
												_lowerCAmelCase  						=			num_attention_heads
												_lowerCAmelCase  						=			hidden_act
												_lowerCAmelCase  						=			intermediate_size
												_lowerCAmelCase  						=			hidden_dropout_prob
												_lowerCAmelCase  						=			attention_probs_dropout_prob
												_lowerCAmelCase  						=			max_position_embeddings
												_lowerCAmelCase  						=			initializer_range
												_lowerCAmelCase  						=			layer_norm_eps
												_lowerCAmelCase  						=			position_embedding_type
												_lowerCAmelCase  						=			use_cache
