



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



class     _UpperCAmelCase			(       nn.Module     ):
     UpperCamelCase__		    =				42
     UpperCamelCase__		    =				jnp.floataa




     def 		snake_case_		(				self):
         A__		 =					nn.Conv(
             self.out_channels ,						kernel_size=(3, 3) ,						strides=(1, 1) ,						padding=((1, 1), (1, 1)) ,						dtype=self.dtype ,						)




     def __call__(				self ,						a__):
         A__ ,	A__ ,	A__ ,	A__		 =					hidden_states.shape
         A__		 =					jax.image.resize(
             __A ,						shape=(batch, height * 2, width * 2, channels) ,						method='''nearest''' ,						)
         A__		 =					self.conv(__A)
         return hidden_states







class     _UpperCAmelCase			(       nn.Module     ):
     UpperCamelCase__		    =				42
     UpperCamelCase__		    =				jnp.floataa




     def 		snake_case_		(				self):
         A__		 =					nn.Conv(
             self.out_channels ,						kernel_size=(3, 3) ,						strides=(2, 2) ,						padding=((1, 1), (1, 1)) ,						dtype=self.dtype ,						)




     def __call__(				self ,						a__):
         # 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(__A)
         return hidden_states







class     _UpperCAmelCase			(       nn.Module     ):
     UpperCamelCase__		    =				42
     UpperCamelCase__		    =				None
     UpperCamelCase__		    =				0.0
     UpperCamelCase__		    =				None
     UpperCamelCase__		    =				jnp.floataa




     def 		snake_case_		(				self):
         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(
             __A ,						kernel_size=(3, 3) ,						strides=(1, 1) ,						padding=((1, 1), (1, 1)) ,						dtype=self.dtype ,						)

         A__		 =					nn.Dense(__A ,						dtype=self.dtype)

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




     def __call__(				self ,						a__ ,						a__ ,						a__=True):
         A__		 =					hidden_states
         A__		 =					self.norma(__A)
         A__		 =					nn.swish(__A)
         A__		 =					self.conva(__A)

         A__		 =					self.time_emb_proj(nn.swish(__A))
         A__		 =					jnp.expand_dims(jnp.expand_dims(__A ,						1) ,						1)
         A__		 =					hidden_states + temb

         A__		 =					self.norma(__A)
         A__		 =					nn.swish(__A)
         A__		 =					self.dropout(__A ,						__A)
         A__		 =					self.conva(__A)

         if self.conv_shortcut is not None:
             A__		 =					self.conv_shortcut(__A)

         return hidden_states + residual







from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowercase						  =		logging.get_logger(__name__)



class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				'''timm_backbone'''




     def __init__(				self ,						a__=None ,						a__=3 ,						a__=True ,						a__=True ,						a__=None ,						**a__ ,						):
         super().__init__(**a__)
         A__		 =					backbone
         A__		 =					num_channels
         A__		 =					features_only
         A__		 =					use_pretrained_backbone
         A__		 =					True
         A__		 =					out_indices if out_indices is not None else (-1,)
