






import flax.linen as nn
import jax
import jax.numpy as jnp
class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		snake_case    =   42
		snake_case    =   jnp.floataa
		def       A(  self		:       Dict		)    ->   List[str]:


						'''simple docstring'''
						A	 =      nn.Conv(
						    self.out_channels						,kernel_size=(3, 3)						,strides=(1, 1)						,padding=((1, 1), (1, 1))						,dtype=self.dtype						,)







		def __call__(  self		:       Union[str, Any]						,_SCREAMING_SNAKE_CASE		:       Any		)    ->   str:


						'''simple docstring'''
						A ,		A ,		A ,		A	 =      hidden_states.shape
						A	 =      jax.image.resize(
						    _SCREAMING_SNAKE_CASE						,shape=(batch, height * 2, width * 2, channels)						,method='nearest'						,)
						A	 =      self.conv(_SCREAMING_SNAKE_CASE		)
						return hidden_states


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		snake_case    =   42
		snake_case    =   jnp.floataa
		def       A(  self		:       str		)    ->   Optional[int]:


						'''simple docstring'''
						A	 =      nn.Conv(
						    self.out_channels						,kernel_size=(3, 3)						,strides=(2, 2)						,padding=((1, 1), (1, 1))						,dtype=self.dtype						,)







		def __call__(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       str		)    ->   Tuple:


						'''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)
						A	 =      self.conv(_SCREAMING_SNAKE_CASE		)
						return hidden_states


class      UpperCamelCase (				nn.Module ):







		"""simple docstring"""



		snake_case    =   42
		snake_case    =   None
		snake_case    =   0.0
		snake_case    =   None
		snake_case    =   jnp.floataa
		def       A(  self		:       Optional[Any]		)    ->   Optional[int]:


						'''simple docstring'''
						A	 =      self.in_channels if self.out_channels is None else self.out_channels

						A	 =      nn.GroupNorm(num_groups=3_2						,epsilon=1E-5		)
						A	 =      nn.Conv(
						    _SCREAMING_SNAKE_CASE						,kernel_size=(3, 3)						,strides=(1, 1)						,padding=((1, 1), (1, 1))						,dtype=self.dtype						,)

						A	 =      nn.Dense(_SCREAMING_SNAKE_CASE						,dtype=self.dtype		)

						A	 =      nn.GroupNorm(num_groups=3_2						,epsilon=1E-5		)
						A	 =      nn.Dropout(self.dropout_prob		)
						A	 =      nn.Conv(
						    _SCREAMING_SNAKE_CASE						,kernel_size=(3, 3)						,strides=(1, 1)						,padding=((1, 1), (1, 1))						,dtype=self.dtype						,)

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

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







		def __call__(  self		:       Tuple						,_SCREAMING_SNAKE_CASE		:       List[str]						,_SCREAMING_SNAKE_CASE		:       str						,_SCREAMING_SNAKE_CASE		:       Optional[int]=True		)    ->   str:


						'''simple docstring'''
						A	 =      hidden_states
						A	 =      self.norma(_SCREAMING_SNAKE_CASE		)
						A	 =      nn.swish(_SCREAMING_SNAKE_CASE		)
						A	 =      self.conva(_SCREAMING_SNAKE_CASE		)

						A	 =      self.time_emb_proj(nn.swish(_SCREAMING_SNAKE_CASE		)		)
						A	 =      jnp.expand_dims(jnp.expand_dims(_SCREAMING_SNAKE_CASE						,1		)						,1		)
						A	 =      hidden_states + temb

						A	 =      self.norma(_SCREAMING_SNAKE_CASE		)
						A	 =      nn.swish(_SCREAMING_SNAKE_CASE		)
						A	 =      self.dropout(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
						A	 =      self.conva(_SCREAMING_SNAKE_CASE		)

						if self.conv_shortcut is not None:
										A	 =      self.conv_shortcut(_SCREAMING_SNAKE_CASE		)

						return hidden_states + residual













def      snake_case		(							UpperCAmelCase	:       int   ):
				A	 =      0
				while num > 0:
								digit_sum += num % 10
								num //= 10
				return digit_sum







def      snake_case		(							UpperCAmelCase	:       int = 1_00   ):
				A	 =      1
				A	 =      2

				for i in range(2,						max_n + 1   ):
								A	 =      pre_numerator
								A	 =      2 * i // 3 if i % 3 == 0 else 1
								A	 =      cur_numerator
								A	 =      e_cont * pre_numerator + temp

				return sum_digits(UpperCAmelCase   )


if __name__ == "__main__":
					print(f'''{solution() = }''')





