
'''simple docstring'''







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


class 		__UpperCAmelCase   (      nn.Module		):
		__lowercase   	= 42
		__lowercase   	= jnp.floataa


		def 				lowerCamelCase       ( self						):





			"""simple docstring"""






			_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					,     lowerCAmelCase_						):





			"""simple docstring"""






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


class 		__UpperCAmelCase   (      nn.Module		):
		__lowercase   	= 42
		__lowercase   	= jnp.floataa


		def 				lowerCamelCase       ( self						):





			"""simple docstring"""






			_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					,     lowerCAmelCase_						):





			"""simple docstring"""






			_snake_case     		=	self.conv(lowerCAmelCase_						)
			return hidden_states


class 		__UpperCAmelCase   (      nn.Module		):
		__lowercase   	= 42
		__lowercase   	= None
		__lowercase   	= 0.0
		__lowercase   	= None
		__lowercase   	= jnp.floataa


		def 				lowerCamelCase       ( self						):





			"""simple docstring"""






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

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

			_snake_case     		=	nn.Dense(lowerCAmelCase_					,     dtype=self.dtype						)

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

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

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




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





			"""simple docstring"""






			_snake_case     		=	hidden_states
			_snake_case     		=	self.norma(lowerCAmelCase_						)
			_snake_case     		=	nn.swish(lowerCAmelCase_						)
			_snake_case     		=	self.conva(lowerCAmelCase_						)

			_snake_case     		=	self.time_emb_proj(nn.swish(lowerCAmelCase_						)						)
			_snake_case     		=	jnp.expand_dims(jnp.expand_dims(lowerCAmelCase_					,     1						)					,     1						)
			_snake_case     		=	hidden_states + temb

			_snake_case     		=	self.norma(lowerCAmelCase_						)
			_snake_case     		=	nn.swish(lowerCAmelCase_						)
			_snake_case     		=	self.dropout(lowerCAmelCase_					,     lowerCAmelCase_						)
			_snake_case     		=	self.conva(lowerCAmelCase_						)

			if self.conv_shortcut is not None:
				_snake_case     		=	self.conv_shortcut(lowerCAmelCase_						)

			return hidden_states + residual



'''simple docstring'''







lowercase					:							Tuple   =      {
    "a": "AAAAA",
    "b": "AAAAB",
    "c": "AAABA",
    "d": "AAABB",
    "e": "AABAA",
    "f": "AABAB",
    "g": "AABBA",
    "h": "AABBB",
    "i": "ABAAA",
    "j": "BBBAA",
    "k": "ABAAB",
    "l": "ABABA",
    "m": "ABABB",
    "n": "ABBAA",
    "o": "ABBAB",
    "p": "ABBBA",
    "q": "ABBBB",
    "r": "BAAAA",
    "s": "BAAAB",
    "t": "BAABA",
    "u": "BAABB",
    "v": "BBBAB",
    "w": "BABAA",
    "x": "BABAB",
    "y": "BABBA",
    "z": "BABBB",
    " ": " ",
}


lowercase					:							Union[str, Any]   =      {value: key for key, value in encode_dict.items()}







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			str:
	_snake_case     		=	''
	for letter in word.lower():
		if letter.isalpha() or letter == " ":
			encoded += encode_dict[letter]
		else:
			raise Exception('encode() accepts only letters of the alphabet and spaces'			)
	return encoded







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			str:
	if set(__A			) - {"A", "B", " "} != set():
		raise Exception('decode() accepts only \'A\', \'B\' and spaces'			)
	_snake_case     		=	''
	for word in coded.split():
		while len(__A			) != 0:
			decoded += decode_dict[word[:5]]
			_snake_case     		=	word[5:]
		decoded += " "
	return decoded.strip()


if __name__ == "__main__":
					from doctest import testmod

					testmod()

