


'''simple docstring'''






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

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class    __magic_name__				(						nn.Module							):
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =0.0
       UpperCAmelCase		 =1
       UpperCAmelCase		 =1
       UpperCAmelCase		 =True
       UpperCAmelCase		 =False
       UpperCAmelCase		 =False
       UpperCAmelCase		 =False
       UpperCAmelCase		 =jnp.floataa






       def      lowerCAmelCase						(     self)		->    List[str]:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =[]
          _UpperCAmelCase  :  str				       =[]

          for i in range(self.num_layers):
             _UpperCAmelCase  :  Any				       =self.in_channels if i == 0 else self.out_channels

             _UpperCAmelCase  :  Optional[Any]				       =FlaxResnetBlockaD(
                 in_channels=snake_case    ,				out_channels=self.out_channels    ,				dropout_prob=self.dropout    ,				dtype=self.dtype    ,				)
             resnets.append(snake_case)

             _UpperCAmelCase  :  Any				       =FlaxTransformeraDModel(
                 in_channels=self.out_channels    ,				n_heads=self.num_attention_heads    ,				d_head=self.out_channels // self.num_attention_heads    ,				depth=1    ,				use_linear_projection=self.use_linear_projection    ,				only_cross_attention=self.only_cross_attention    ,				use_memory_efficient_attention=self.use_memory_efficient_attention    ,				dtype=self.dtype    ,				)
             attentions.append(snake_case)

          _UpperCAmelCase  :  List[Any]				       =resnets
          _UpperCAmelCase  :  Optional[Any]				       =attentions

          if self.add_downsample:
             _UpperCAmelCase  :  Any				       =FlaxDownsampleaD(self.out_channels    ,				dtype=self.dtype)






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




          '''simple docstring'''
          _UpperCAmelCase  :  str				       =()

          for resnet, attn in zip(self.resnets    ,				self.attentions):
             _UpperCAmelCase  :  Tuple				       =resnet(snake_case    ,				snake_case    ,				deterministic=snake_case)
             _UpperCAmelCase  :  List[Any]				       =attn(snake_case    ,				snake_case    ,				deterministic=snake_case)
             output_states += (hidden_states,)

          if self.add_downsample:
             _UpperCAmelCase  :  List[str]				       =self.downsamplers_a(snake_case)
             output_states += (hidden_states,)

          return hidden_states, output_states






class    __magic_name__				(						nn.Module							):
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =0.0
       UpperCAmelCase		 =1
       UpperCAmelCase		 =True
       UpperCAmelCase		 =jnp.floataa






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =[]

          for i in range(self.num_layers):
             _UpperCAmelCase  :  Union[str, Any]				       =self.in_channels if i == 0 else self.out_channels

             _UpperCAmelCase  :  Any				       =FlaxResnetBlockaD(
                 in_channels=snake_case    ,				out_channels=self.out_channels    ,				dropout_prob=self.dropout    ,				dtype=self.dtype    ,				)
             resnets.append(snake_case)
          _UpperCAmelCase  :  Any				       =resnets

          if self.add_downsample:
             _UpperCAmelCase  :  Tuple				       =FlaxDownsampleaD(self.out_channels    ,				dtype=self.dtype)






       def __call__(     self    ,				snake_case    ,				snake_case    ,				snake_case=True)		->    Any:




          '''simple docstring'''
          _UpperCAmelCase  :  List[Any]				       =()

          for resnet in self.resnets:
             _UpperCAmelCase  :  List[str]				       =resnet(snake_case    ,				snake_case    ,				deterministic=snake_case)
             output_states += (hidden_states,)

          if self.add_downsample:
             _UpperCAmelCase  :  Any				       =self.downsamplers_a(snake_case)
             output_states += (hidden_states,)

          return hidden_states, output_states






class    __magic_name__				(						nn.Module							):
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =0.0
       UpperCAmelCase		 =1
       UpperCAmelCase		 =1
       UpperCAmelCase		 =True
       UpperCAmelCase		 =False
       UpperCAmelCase		 =False
       UpperCAmelCase		 =False
       UpperCAmelCase		 =jnp.floataa






       def      lowerCAmelCase						(     self)		->    List[str]:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[Any]				       =[]
          _UpperCAmelCase  :  str				       =[]

          for i in range(self.num_layers):
             _UpperCAmelCase  :  Dict				       =self.in_channels if (i == self.num_layers - 1) else self.out_channels
             _UpperCAmelCase  :  Dict				       =self.prev_output_channel if i == 0 else self.out_channels

             _UpperCAmelCase  :  List[Any]				       =FlaxResnetBlockaD(
                 in_channels=resnet_in_channels + res_skip_channels    ,				out_channels=self.out_channels    ,				dropout_prob=self.dropout    ,				dtype=self.dtype    ,				)
             resnets.append(snake_case)

             _UpperCAmelCase  :  Union[str, Any]				       =FlaxTransformeraDModel(
                 in_channels=self.out_channels    ,				n_heads=self.num_attention_heads    ,				d_head=self.out_channels // self.num_attention_heads    ,				depth=1    ,				use_linear_projection=self.use_linear_projection    ,				only_cross_attention=self.only_cross_attention    ,				use_memory_efficient_attention=self.use_memory_efficient_attention    ,				dtype=self.dtype    ,				)
             attentions.append(snake_case)

          _UpperCAmelCase  :  Optional[Any]				       =resnets
          _UpperCAmelCase  :  str				       =attentions

          if self.add_upsample:
             _UpperCAmelCase  :  List[str]				       =FlaxUpsampleaD(self.out_channels    ,				dtype=self.dtype)






       def __call__(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case=True)		->    str:




          '''simple docstring'''
          for resnet, attn in zip(self.resnets    ,				self.attentions):
             # pop res hidden states
             _UpperCAmelCase  :  Any				       =res_hidden_states_tuple[-1]
             _UpperCAmelCase  :  Optional[int]				       =res_hidden_states_tuple[:-1]
             _UpperCAmelCase  :  str				       =jnp.concatenate((hidden_states, res_hidden_states)    ,				axis=-1)

             _UpperCAmelCase  :  int				       =resnet(snake_case    ,				snake_case    ,				deterministic=snake_case)
             _UpperCAmelCase  :  List[Any]				       =attn(snake_case    ,				snake_case    ,				deterministic=snake_case)

          if self.add_upsample:
             _UpperCAmelCase  :  List[str]				       =self.upsamplers_a(snake_case)

          return hidden_states






class    __magic_name__				(						nn.Module							):
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =42
       UpperCAmelCase		 =0.0
       UpperCAmelCase		 =1
       UpperCAmelCase		 =True
       UpperCAmelCase		 =jnp.floataa






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =[]

          for i in range(self.num_layers):
             _UpperCAmelCase  :  int				       =self.in_channels if (i == self.num_layers - 1) else self.out_channels
             _UpperCAmelCase  :  Dict				       =self.prev_output_channel if i == 0 else self.out_channels

             _UpperCAmelCase  :  Union[str, Any]				       =FlaxResnetBlockaD(
                 in_channels=resnet_in_channels + res_skip_channels    ,				out_channels=self.out_channels    ,				dropout_prob=self.dropout    ,				dtype=self.dtype    ,				)
             resnets.append(snake_case)

          _UpperCAmelCase  :  Any				       =resnets

          if self.add_upsample:
             _UpperCAmelCase  :  Tuple				       =FlaxUpsampleaD(self.out_channels    ,				dtype=self.dtype)






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




          '''simple docstring'''
          for resnet in self.resnets:
             # pop res hidden states
             _UpperCAmelCase  :  Optional[int]				       =res_hidden_states_tuple[-1]
             _UpperCAmelCase  :  List[str]				       =res_hidden_states_tuple[:-1]
             _UpperCAmelCase  :  Any				       =jnp.concatenate((hidden_states, res_hidden_states)    ,				axis=-1)

             _UpperCAmelCase  :  Optional[int]				       =resnet(snake_case    ,				snake_case    ,				deterministic=snake_case)

          if self.add_upsample:
             _UpperCAmelCase  :  str				       =self.upsamplers_a(snake_case)

          return hidden_states






class    __magic_name__				(						nn.Module							):
       UpperCAmelCase		 =42
       UpperCAmelCase		 =0.0
       UpperCAmelCase		 =1
       UpperCAmelCase		 =1
       UpperCAmelCase		 =False
       UpperCAmelCase		 =False
       UpperCAmelCase		 =jnp.floataa






       def      lowerCAmelCase						(     self)		->    int:




          '''simple docstring'''
          # there is always at least one resnet
          _UpperCAmelCase  :  Any				       =[
              FlaxResnetBlockaD(
                  in_channels=self.in_channels    ,				out_channels=self.in_channels    ,				dropout_prob=self.dropout    ,				dtype=self.dtype    ,				)
          ]

          _UpperCAmelCase  :  List[str]				       =[]

          for _ in range(self.num_layers):
             _UpperCAmelCase  :  Union[str, Any]				       =FlaxTransformeraDModel(
                 in_channels=self.in_channels    ,				n_heads=self.num_attention_heads    ,				d_head=self.in_channels // self.num_attention_heads    ,				depth=1    ,				use_linear_projection=self.use_linear_projection    ,				use_memory_efficient_attention=self.use_memory_efficient_attention    ,				dtype=self.dtype    ,				)
             attentions.append(snake_case)

             _UpperCAmelCase  :  Dict				       =FlaxResnetBlockaD(
                 in_channels=self.in_channels    ,				out_channels=self.in_channels    ,				dropout_prob=self.dropout    ,				dtype=self.dtype    ,				)
             resnets.append(snake_case)

          _UpperCAmelCase  :  List[Any]				       =resnets
          _UpperCAmelCase  :  List[Any]				       =attentions






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




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =self.resnets[0](snake_case    ,				snake_case)
          for attn, resnet in zip(self.attentions    ,				self.resnets[1:]):
             _UpperCAmelCase  :  str				       =attn(snake_case    ,				snake_case    ,				deterministic=snake_case)
             _UpperCAmelCase  :  Optional[Any]				       =resnet(snake_case    ,				snake_case    ,				deterministic=snake_case)

          return hidden_states





'''simple docstring'''






import warnings
from functools import wraps
from typing import Callable





def lowerCamelCase__     (		__lowerCamelCase			:				Callable				):







   '''simple docstring'''




   @wraps(__lowerCamelCase				)
   def _inner_fn(*__lowerCamelCase			:				int			,       **__lowerCamelCase			:				Optional[Any]				):
      warnings.warn(
          (f"'{fn.__name__}' is experimental and might be subject to breaking changes in the future.")			,       __lowerCamelCase			,       )
      return fn(*__lowerCamelCase			,       **__lowerCamelCase				)

   return _inner_fn

