



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



class  lowerCamelCase__    (							nn.Module):




		'''simple docstring'''



		snake_case_			=42
		snake_case_			=jnp.floataa



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

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  nn.Conv(
						    self.out_channels	,kernel_size=(3, 3)	,strides=(1, 1)	,padding=((1, 1), (1, 1))	,dtype=self.dtype	,)



		def __call__(self	,__lowerCamelCase				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__  :       Optional[Any]				=  hidden_states.shape
						lowerCAmelCase__  :       str				=  jax.image.resize(
						    __lowerCamelCase	,shape=(batch, height * 2, width * 2, channels)	,method='''nearest'''	,)
						lowerCAmelCase__  :       Tuple				=  self.conv(__lowerCamelCase				)
						return hidden_states



class  lowerCamelCase__    (							nn.Module):




		'''simple docstring'''



		snake_case_			=42
		snake_case_			=jnp.floataa



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

						"""simple docstring"""






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



		def __call__(self	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  self.conv(__lowerCamelCase				)
						return hidden_states



class  lowerCamelCase__    (							nn.Module):




		'''simple docstring'''



		snake_case_			=42
		snake_case_			=None
		snake_case_			=0.0
		snake_case_			=None
		snake_case_			=jnp.floataa



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  self.in_channels if self.out_channels is None else self.out_channels

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

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

						lowerCAmelCase__  :       Dict				=  nn.GroupNorm(num_groups=32	,epsilon=1e-5				)
						lowerCAmelCase__  :       Dict				=  nn.Dropout(self.dropout_prob				)
						lowerCAmelCase__  :       Any				=  nn.Conv(
						    __lowerCamelCase	,kernel_size=(3, 3)	,strides=(1, 1)	,padding=((1, 1), (1, 1))	,dtype=self.dtype	,)

						lowerCAmelCase__  :       Dict				=  self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

						lowerCAmelCase__  :       Union[str, Any]				=  None
						if use_nin_shortcut:
										lowerCAmelCase__  :       Dict				=  nn.Conv(
										    __lowerCamelCase	,kernel_size=(1, 1)	,strides=(1, 1)	,padding='''VALID'''	,dtype=self.dtype	,)



		def __call__(self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase=True				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  hidden_states
						lowerCAmelCase__  :       Any				=  self.norma(__lowerCamelCase				)
						lowerCAmelCase__  :       Dict				=  nn.swish(__lowerCamelCase				)
						lowerCAmelCase__  :       Tuple				=  self.conva(__lowerCamelCase				)

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

						lowerCAmelCase__  :       Optional[Any]				=  self.norma(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  nn.swish(__lowerCamelCase				)
						lowerCAmelCase__  :       Dict				=  self.dropout(__lowerCamelCase	,__lowerCamelCase				)
						lowerCAmelCase__  :       Dict				=  self.conva(__lowerCamelCase				)

						if self.conv_shortcut is not None:
										lowerCAmelCase__  :       List[str]				=  self.conv_shortcut(__lowerCamelCase				)

						return hidden_states + residual






from collections.abc import Callable

import numpy as np






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Callable ,lowerCamelCase_	:	float ,lowerCamelCase_	:	float ,lowerCamelCase_	:	float ,lowerCamelCase_	:	float):





				'''simple docstring'''

				lowerCAmelCase__  :       Dict				=  int(np.ceil((x_end - xa) / step_size))
				lowerCAmelCase__  :       Any				=  np.zeros((n + 1,))
				lowerCAmelCase__  :       Union[str, Any]				=  ya
				lowerCAmelCase__  :       str				=  xa

				for k in range(lowerCamelCase_):
								lowerCAmelCase__  :       Tuple				=  y[k] + step_size * ode_func(lowerCamelCase_ ,y[k])
								lowerCAmelCase__  :       str				=  y[k] + (
								    (step_size / 2) * (ode_func(lowerCamelCase_ ,y[k]) + ode_func(x + step_size ,lowerCamelCase_))
								)
								x += step_size

				return y


if __name__ == "__main__":
							import doctest

							doctest.testmod()

