

'''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 			lowercase_	(nn.Module      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  float							=			0.0
  SCREAMING_SNAKE_CASE    :  int							=			1
  SCREAMING_SNAKE_CASE    :  int							=			1
  SCREAMING_SNAKE_CASE    :  bool							=			True
  SCREAMING_SNAKE_CASE    :  bool							=			False
  SCREAMING_SNAKE_CASE    :  bool							=			False
  SCREAMING_SNAKE_CASE    :  bool							=			False
  SCREAMING_SNAKE_CASE    :  jnp.dtype							=			jnp.floataa






  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   __lowercase							  = []
   __lowercase							  = []

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

    __lowercase							  = FlaxResnetBlockaD(
        in_channels=lowercase__					,out_channels=self.out_channels					,dropout_prob=self.dropout					,dtype=self.dtype					,)
    resnets.append(lowercase__			)

    __lowercase							  = 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(lowercase__			)

   __lowercase							  = resnets
   __lowercase							  = attentions

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






  def __call__(      self   :     str					,lowercase__   :     Optional[int]					,lowercase__   :     Optional[Any]					,lowercase__   :     List[Any]					,lowercase__   :     int=True			):
   __lowercase							  = ()

   for resnet, attn in zip(self.resnets					,self.attentions			):
    __lowercase							  = resnet(lowercase__					,lowercase__					,deterministic=lowercase__			)
    __lowercase							  = attn(lowercase__					,lowercase__					,deterministic=lowercase__			)
    output_states += (hidden_states,)

   if self.add_downsample:
    __lowercase							  = self.downsamplers_a(lowercase__			)
    output_states += (hidden_states,)

   return hidden_states, output_states






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  float							=			0.0
  SCREAMING_SNAKE_CASE    :  int							=			1
  SCREAMING_SNAKE_CASE    :  bool							=			True
  SCREAMING_SNAKE_CASE    :  jnp.dtype							=			jnp.floataa






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase							  = []

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

    __lowercase							  = FlaxResnetBlockaD(
        in_channels=lowercase__					,out_channels=self.out_channels					,dropout_prob=self.dropout					,dtype=self.dtype					,)
    resnets.append(lowercase__			)
   __lowercase							  = resnets

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






  def __call__(      self   :     Union[str, Any]					,lowercase__   :     Any					,lowercase__   :     Any					,lowercase__   :     Optional[int]=True			):
   __lowercase							  = ()

   for resnet in self.resnets:
    __lowercase							  = resnet(lowercase__					,lowercase__					,deterministic=lowercase__			)
    output_states += (hidden_states,)

   if self.add_downsample:
    __lowercase							  = self.downsamplers_a(lowercase__			)
    output_states += (hidden_states,)

   return hidden_states, output_states






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  float							=			0.0
  SCREAMING_SNAKE_CASE    :  int							=			1
  SCREAMING_SNAKE_CASE    :  int							=			1
  SCREAMING_SNAKE_CASE    :  bool							=			True
  SCREAMING_SNAKE_CASE    :  bool							=			False
  SCREAMING_SNAKE_CASE    :  bool							=			False
  SCREAMING_SNAKE_CASE    :  bool							=			False
  SCREAMING_SNAKE_CASE    :  jnp.dtype							=			jnp.floataa






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   __lowercase							  = []
   __lowercase							  = []

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

    __lowercase							  = FlaxResnetBlockaD(
        in_channels=resnet_in_channels + res_skip_channels					,out_channels=self.out_channels					,dropout_prob=self.dropout					,dtype=self.dtype					,)
    resnets.append(lowercase__			)

    __lowercase							  = 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(lowercase__			)

   __lowercase							  = resnets
   __lowercase							  = attentions

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






  def __call__(      self   :     Tuple					,lowercase__   :     Optional[int]					,lowercase__   :     int					,lowercase__   :     str					,lowercase__   :     str					,lowercase__   :     int=True			):
   for resnet, attn in zip(self.resnets					,self.attentions			):
    # pop res hidden states
    __lowercase							  = res_hidden_states_tuple[-1]
    __lowercase							  = res_hidden_states_tuple[:-1]
    __lowercase							  = jnp.concatenate((hidden_states, res_hidden_states)					,axis=-1			)

    __lowercase							  = resnet(lowercase__					,lowercase__					,deterministic=lowercase__			)
    __lowercase							  = attn(lowercase__					,lowercase__					,deterministic=lowercase__			)

   if self.add_upsample:
    __lowercase							  = self.upsamplers_a(lowercase__			)

   return hidden_states






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  float							=			0.0
  SCREAMING_SNAKE_CASE    :  int							=			1
  SCREAMING_SNAKE_CASE    :  bool							=			True
  SCREAMING_SNAKE_CASE    :  jnp.dtype							=			jnp.floataa






  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   __lowercase							  = []

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

    __lowercase							  = FlaxResnetBlockaD(
        in_channels=resnet_in_channels + res_skip_channels					,out_channels=self.out_channels					,dropout_prob=self.dropout					,dtype=self.dtype					,)
    resnets.append(lowercase__			)

   __lowercase							  = resnets

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






  def __call__(      self   :     List[Any]					,lowercase__   :     Any					,lowercase__   :     List[str]					,lowercase__   :     Optional[Any]					,lowercase__   :     Any=True			):
   for resnet in self.resnets:
    # pop res hidden states
    __lowercase							  = res_hidden_states_tuple[-1]
    __lowercase							  = res_hidden_states_tuple[:-1]
    __lowercase							  = jnp.concatenate((hidden_states, res_hidden_states)					,axis=-1			)

    __lowercase							  = resnet(lowercase__					,lowercase__					,deterministic=lowercase__			)

   if self.add_upsample:
    __lowercase							  = self.upsamplers_a(lowercase__			)

   return hidden_states






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  int
  SCREAMING_SNAKE_CASE    :  float							=			0.0
  SCREAMING_SNAKE_CASE    :  int							=			1
  SCREAMING_SNAKE_CASE    :  int							=			1
  SCREAMING_SNAKE_CASE    :  bool							=			False
  SCREAMING_SNAKE_CASE    :  bool							=			False
  SCREAMING_SNAKE_CASE    :  jnp.dtype							=			jnp.floataa






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   # there is always at least one resnet
   __lowercase							  = [
       FlaxResnetBlockaD(
           in_channels=self.in_channels					,out_channels=self.in_channels					,dropout_prob=self.dropout					,dtype=self.dtype					,)
   ]

   __lowercase							  = []

   for _ in range(self.num_layers			):
    __lowercase							  = 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(lowercase__			)

    __lowercase							  = FlaxResnetBlockaD(
        in_channels=self.in_channels					,out_channels=self.in_channels					,dropout_prob=self.dropout					,dtype=self.dtype					,)
    resnets.append(lowercase__			)

   __lowercase							  = resnets
   __lowercase							  = attentions






  def __call__(      self   :     Union[str, Any]					,lowercase__   :     str					,lowercase__   :     str					,lowercase__   :     List[str]					,lowercase__   :     Optional[int]=True			):
   __lowercase							  = self.resnets[0](lowercase__					,lowercase__			)
   for attn, resnet in zip(self.attentions					,self.resnets[1:]			):
    __lowercase							  = attn(lowercase__					,lowercase__					,deterministic=lowercase__			)
    __lowercase							  = resnet(lowercase__					,lowercase__					,deterministic=lowercase__			)

   return hidden_states






'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


_UpperCamelCase											=						{'configuration_speech_encoder_decoder': ['SpeechEncoderDecoderConfig']}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCamelCase											=						['SpeechEncoderDecoderModel']

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCamelCase											=						['FlaxSpeechEncoderDecoderModel']

if TYPE_CHECKING:
   from .configuration_speech_encoder_decoder import SpeechEncoderDecoderConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_speech_encoder_decoder import SpeechEncoderDecoderModel

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_speech_encoder_decoder import FlaxSpeechEncoderDecoderModel

else:
   import sys

   _UpperCamelCase											=						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






