


'''simple docstring'''




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

class 		UpperCAmelCase_      (  nn.Module		):
     lowerCamelCase				:	Optional[Any]		=     42
     lowerCamelCase				:	List[Any]		=     jnp.floataa





     def 							__UpperCAmelCase      ( self       :      Any    )				->  Tuple:
        lowerCAmelCase     =			nn.Conv(
            self.out_channels    ,      kernel_size=(3, 3)    ,      strides=(1, 1)    ,      padding=((1, 1), (1, 1))    ,      dtype=self.dtype    ,      )





     def __call__( self       :      List[Any]    ,      UpperCAmelCase__       :      Dict    )				->  Dict:
        lowerCAmelCase     =			hidden_states.shape
        lowerCAmelCase     =			jax.image.resize(
            a_    ,      shape=(batch, height * 2, width * 2, channels)    ,      method='nearest'    ,      )
        lowerCAmelCase     =			self.conv(a_    )
        return hidden_states




class 		UpperCAmelCase_      (  nn.Module		):
     lowerCamelCase				:	Any		=     42
     lowerCamelCase				:	List[Any]		=     jnp.floataa





     def 							__UpperCAmelCase      ( self       :      Optional[Any]    )				->  int:
        lowerCAmelCase     =			nn.Conv(
            self.out_channels    ,      kernel_size=(3, 3)    ,      strides=(2, 2)    ,      padding=((1, 1), (1, 1))    ,      dtype=self.dtype    ,      )





     def __call__( self       :      Optional[Any]    ,      UpperCAmelCase__       :      Any    )				->  Tuple:
        lowerCAmelCase     =			self.conv(a_    )
        return hidden_states




class 		UpperCAmelCase_      (  nn.Module		):
     lowerCamelCase				:	Tuple		=     42
     lowerCamelCase				:	Any		=     None
     lowerCamelCase				:	Any		=     0.0
     lowerCamelCase				:	Any		=     None
     lowerCamelCase				:	int		=     jnp.floataa





     def 							__UpperCAmelCase      ( self       :      Optional[Any]    )				->  Union[str, Any]:
        lowerCAmelCase     =			self.in_channels if self.out_channels is None else self.out_channels

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

        lowerCAmelCase     =			nn.Dense(a_    ,      dtype=self.dtype    )

        lowerCAmelCase     =			nn.GroupNorm(num_groups=3_2    ,      epsilon=1E-5    )
        lowerCAmelCase     =			nn.Dropout(self.dropout_prob    )
        lowerCAmelCase     =			nn.Conv(
            a_    ,      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(
               a_    ,      kernel_size=(1, 1)    ,      strides=(1, 1)    ,      padding='VALID'    ,      dtype=self.dtype    ,      )





     def __call__( self       :      Union[str, Any]    ,      UpperCAmelCase__       :      Any    ,      UpperCAmelCase__       :      List[str]    ,      UpperCAmelCase__       :      int=True    )				->  Optional[int]:
        lowerCAmelCase     =			hidden_states
        lowerCAmelCase     =			self.norma(a_    )
        lowerCAmelCase     =			nn.swish(a_    )
        lowerCAmelCase     =			self.conva(a_    )

        lowerCAmelCase     =			self.time_emb_proj(nn.swish(a_    )    )
        lowerCAmelCase     =			jnp.expand_dims(jnp.expand_dims(a_    ,      1    )    ,      1    )
        lowerCAmelCase     =			hidden_states + temb

        lowerCAmelCase     =			self.norma(a_    )
        lowerCAmelCase     =			nn.swish(a_    )
        lowerCAmelCase     =			self.dropout(a_    ,      a_    )
        lowerCAmelCase     =			self.conva(a_    )

        if self.conv_shortcut is not None:
           lowerCAmelCase     =			self.conv_shortcut(a_    )

        return hidden_states + residual










import unittest

from transformers import load_tool
from transformers.utils import is_torch_available


if is_torch_available():
				import torch

from transformers.testing_utils import require_torch

from .test_tools_common import ToolTesterMixin

@require_torch
class 		UpperCAmelCase__			(			unittest.TestCase						,__UpperCamelCase   ):




			'''simple docstring'''






			def  snake_case__     (   self							: List[Any] ):
					'''simple docstring'''





					__UpperCAmelCase    :     Tuple													=	load_tool('''text-to-speech''' )
					self.tool.setup()




			def  snake_case__     (   self							: Union[str, Any] ):
					'''simple docstring'''





					torch.manual_seed(0 )
					__UpperCAmelCase    :     List[Any]													=	self.tool('''hey''' )
					__UpperCAmelCase    :     Union[str, Any]													=	result.to_raw()
					self.assertTrue(
					    torch.allclose(
					        resulting_tensor[:3]       ,       torch.tensor([-0.0_0_0_5_9_6_6_6_6_8_8_3_2_1_1_5_8_2_9, -0.0_0_0_3_6_5_7_6_4_0_1_9_0_7_9_5_0_6_4, -0.0_0_0_1_3_4_3_9_5_0_2_7_9_9_8_8_3_4_8_5] )       ,       ) )




			def  snake_case__     (   self							: Optional[Any] ):
					'''simple docstring'''





					torch.manual_seed(0 )
					__UpperCAmelCase    :     Optional[int]													=	self.tool('''hey''' )
					__UpperCAmelCase    :     List[str]													=	result.to_raw()
					self.assertTrue(
					    torch.allclose(
					        resulting_tensor[:3]       ,       torch.tensor([-0.0_0_0_5_9_6_6_6_6_8_8_3_2_1_1_5_8_2_9, -0.0_0_0_3_6_5_7_6_4_0_1_9_0_7_9_5_0_6_4, -0.0_0_0_1_3_4_3_9_5_0_2_7_9_9_8_8_3_4_8_5] )       ,       ) )





