






'''simple docstring'''
import functools
import math

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




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase = 4096							):


   """simple docstring"""




   lowercase_		:      str          =       key.shape[-3:]
   lowercase_		:      Optional[int]          =       value.shape[-1]
   lowercase_		:      List[Any]          =       min(_UpperCamelCase  ,						_UpperCamelCase							)
   lowercase_		:      Any          =       query / jnp.sqrt(_UpperCamelCase							)

   @functools.partial(jax.checkpoint  ,						prevent_cse=_UpperCamelCase							)
   def summarize_chunk(_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							):
      lowercase_		:      Dict          =       jnp.einsum("...qhd,...khd->...qhk"  ,						_UpperCamelCase  ,						_UpperCamelCase  ,						precision=_UpperCamelCase							)

      lowercase_		:      Any          =       jnp.max(_UpperCamelCase  ,						axis=-1  ,						keepdims=_UpperCamelCase							)
      lowercase_		:      List[Any]          =       jax.lax.stop_gradient(_UpperCamelCase							)
      lowercase_		:      Any          =       jnp.exp(attn_weights - max_score							)

      lowercase_		:      List[Any]          =       jnp.einsum("...vhf,...qhv->...qhf"  ,						_UpperCamelCase  ,						_UpperCamelCase  ,						precision=_UpperCamelCase							)
      lowercase_		:      List[Any]          =       jnp.einsum("...qhk->...qh"  ,						_UpperCamelCase							)

      return (exp_values, exp_weights.sum(axis=-1							), max_score)

   def chunk_scanner(_UpperCamelCase							):
      # julienne key array
      lowercase_		:      str          =       jax.lax.dynamic_slice(
          operand=_UpperCamelCase  ,						start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]  ,						slice_sizes=list(key.shape[:-3]							) + [key_chunk_size, num_heads, k_features]  ,						)

      # julienne value array
      lowercase_		:      str          =       jax.lax.dynamic_slice(
          operand=_UpperCamelCase  ,						start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]  ,						slice_sizes=list(value.shape[:-3]							) + [key_chunk_size, num_heads, v_features]  ,						)

      return summarize_chunk(_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							)

   lowercase_		:      Dict          =       jax.lax.map(f=_UpperCamelCase  ,						xs=jnp.arange(0  ,						_UpperCamelCase  ,						_UpperCamelCase							)							)

   lowercase_		:      Any          =       jnp.max(_UpperCamelCase  ,						axis=0  ,						keepdims=_UpperCamelCase							)
   lowercase_		:      Tuple          =       jnp.exp(chunk_max - global_max							)

   chunk_values *= jnp.expand_dims(_UpperCamelCase  ,						axis=-1							)
   chunk_weights *= max_diffs

   lowercase_		:      Optional[Any]          =       chunk_values.sum(axis=0							)
   lowercase_		:      Optional[int]          =       jnp.expand_dims(_UpperCamelCase  ,						-1							).sum(axis=0							)

   return all_values / all_weights




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase=jax.lax.Precision.HIGHEST  ,						_UpperCamelCase = 1024  ,						_UpperCamelCase = 4096							):


   """simple docstring"""




   lowercase_		:      Union[str, Any]          =       query.shape[-3:]

   def chunk_scanner(_UpperCamelCase  ,						_UpperCamelCase							):
      # julienne query array
      lowercase_		:      List[str]          =       jax.lax.dynamic_slice(
          operand=_UpperCamelCase  ,						start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]  ,						slice_sizes=list(query.shape[:-3]							) + [min(_UpperCamelCase  ,						_UpperCamelCase							), num_heads, q_features]  ,						)

      return (
          chunk_idx + query_chunk_size,  # unused ignore it
          _query_chunk_attention(
              query=_UpperCamelCase  ,						key=_UpperCamelCase  ,						value=_UpperCamelCase  ,						precision=_UpperCamelCase  ,						key_chunk_size=_UpperCamelCase							),
      )

   lowercase_		:      List[str]          =       jax.lax.scan(
       f=_UpperCamelCase  ,						init=0  ,						xs=_UpperCamelCase  ,						length=math.ceil(num_q / query_chunk_size							)							)

   return jnp.concatenate(_UpperCamelCase  ,						axis=-3							)  # fuse the chunked result back


class     _UpperCAmelCase  (						nn.Module     ):
    __lowerCamelCase:				int
    __lowerCamelCase:				int						    =      8
    __lowerCamelCase:				int						    =      64
    __lowerCamelCase:				float						    =      0.0
    __lowerCamelCase:				bool						    =      False
    __lowerCamelCase:				jnp.dtype						    =      jnp.floataa




    def 			lowerCAmelCase__			(     self					:    Dict     ):






       '''simple docstring'''






       lowercase_		:      Dict          =       self.dim_head * self.heads
       lowercase_		:      Union[str, Any]          =       self.dim_head**-0.5

       # Weights were exported with old names {to_q, to_k, to_v, to_out}
       lowercase_		:      Union[str, Any]          =       nn.Dense(a					,       use_bias=a					,       dtype=self.dtype					,       name="to_q"     )
       lowercase_		:      Any          =       nn.Dense(a					,       use_bias=a					,       dtype=self.dtype					,       name="to_k"     )
       lowercase_		:      List[str]          =       nn.Dense(a					,       use_bias=a					,       dtype=self.dtype					,       name="to_v"     )

       lowercase_		:      List[Any]          =       nn.Dense(self.query_dim					,       dtype=self.dtype					,       name="to_out_0"     )
       lowercase_		:      List[str]          =       nn.Dropout(rate=self.dropout     )




    def 			lowerCAmelCase__			(     self					:    Optional[int]					,       a					:    Tuple     ):






       '''simple docstring'''






       lowercase_		:      Tuple          =       tensor.shape
       lowercase_		:      Optional[Any]          =       self.heads
       lowercase_		:      Tuple          =       tensor.reshape(a					,       a					,       a					,       dim // head_size     )
       lowercase_		:      int          =       jnp.transpose(a					,       (0, 2, 1, 3)     )
       lowercase_		:      List[Any]          =       tensor.reshape(batch_size * head_size					,       a					,       dim // head_size     )
       return tensor




    def 			lowerCAmelCase__			(     self					:    List[str]					,       a					:    List[str]     ):






       '''simple docstring'''






       lowercase_		:      str          =       tensor.shape
       lowercase_		:      int          =       self.heads
       lowercase_		:      Union[str, Any]          =       tensor.reshape(batch_size // head_size					,       a					,       a					,       a     )
       lowercase_		:      List[str]          =       jnp.transpose(a					,       (0, 2, 1, 3)     )
       lowercase_		:      Optional[Any]          =       tensor.reshape(batch_size // head_size					,       a					,       dim * head_size     )
       return tensor





    def __call__(     self					:    Dict					,       a					:    Dict					,       a					:    List[Any]=None					,       a					:    int=True     ):






       '''simple docstring'''






       lowercase_		:      int          =       hidden_states if context is None else context

       lowercase_		:      Dict          =       self.query(a     )
       lowercase_		:      List[Any]          =       self.key(a     )
       lowercase_		:      int          =       self.value(a     )

       lowercase_		:      str          =       self.reshape_heads_to_batch_dim(a     )
       lowercase_		:      Union[str, Any]          =       self.reshape_heads_to_batch_dim(a     )
       lowercase_		:      str          =       self.reshape_heads_to_batch_dim(a     )

       if self.use_memory_efficient_attention:
          lowercase_		:      str          =       query_states.transpose(1					,       0					,       2     )
          lowercase_		:      Tuple          =       key_states.transpose(1					,       0					,       2     )
          lowercase_		:      str          =       value_states.transpose(1					,       0					,       2     )

          # this if statement create a chunk size for each layer of the unet
          # the chunk size is equal to the query_length dimension of the deepest layer of the unet

          lowercase_		:      Optional[Any]          =       query_states.shape[-3]
          if flatten_latent_dim % 6_4 == 0:
             lowercase_		:      Tuple          =       int(flatten_latent_dim / 6_4     )
          elif flatten_latent_dim % 1_6 == 0:
             lowercase_		:      Tuple          =       int(flatten_latent_dim / 1_6     )
          elif flatten_latent_dim % 4 == 0:
             lowercase_		:      Union[str, Any]          =       int(flatten_latent_dim / 4     )
          else:
             lowercase_		:      Union[str, Any]          =       int(a     )

          lowercase_		:      Any          =       jax_memory_efficient_attention(
              a					,       a					,       a					,       query_chunk_size=a					,       key_chunk_size=4_0_9_6 * 4     )

          lowercase_		:      Dict          =       hidden_states.transpose(1					,       0					,       2     )
       else:
          # compute attentions
          lowercase_		:      Optional[int]          =       jnp.einsum("b i d, b j d->b i j"					,       a					,       a     )
          lowercase_		:      List[Any]          =       attention_scores * self.scale
          lowercase_		:      Union[str, Any]          =       nn.softmax(a					,       axis=2     )

          # attend to values
          lowercase_		:      Dict          =       jnp.einsum("b i j, b j d -> b i d"					,       a					,       a     )

       lowercase_		:      Optional[Any]          =       self.reshape_batch_dim_to_heads(a     )
       lowercase_		:      str          =       self.proj_attn(a     )
       return self.dropout_layer(a					,       deterministic=a     )


class     _UpperCAmelCase  (						nn.Module     ):
    __lowerCamelCase:				int
    __lowerCamelCase:				int
    __lowerCamelCase:				int
    __lowerCamelCase:				float						    =      0.0
    __lowerCamelCase:				bool						    =      False
    __lowerCamelCase:				jnp.dtype						    =      jnp.floataa
    __lowerCamelCase:				bool						    =      False




    def 			lowerCAmelCase__			(     self					:    Any     ):






       '''simple docstring'''






       lowercase_		:      str          =       FlaxAttention(
           self.dim					,       self.n_heads					,       self.d_head					,       self.dropout					,       self.use_memory_efficient_attention					,       dtype=self.dtype     )
       # cross attention
       lowercase_		:      List[Any]          =       FlaxAttention(
           self.dim					,       self.n_heads					,       self.d_head					,       self.dropout					,       self.use_memory_efficient_attention					,       dtype=self.dtype     )
       lowercase_		:      List[Any]          =       FlaxFeedForward(dim=self.dim					,       dropout=self.dropout					,       dtype=self.dtype     )
       lowercase_		:      Tuple          =       nn.LayerNorm(epsilon=1e-5					,       dtype=self.dtype     )
       lowercase_		:      Union[str, Any]          =       nn.LayerNorm(epsilon=1e-5					,       dtype=self.dtype     )
       lowercase_		:      List[str]          =       nn.LayerNorm(epsilon=1e-5					,       dtype=self.dtype     )
       lowercase_		:      Optional[int]          =       nn.Dropout(rate=self.dropout     )





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






       '''simple docstring'''






       lowercase_		:      Any          =       hidden_states
       if self.only_cross_attention:
          lowercase_		:      str          =       self.attna(self.norma(a     )					,       a					,       deterministic=a     )
       else:
          lowercase_		:      Dict          =       self.attna(self.norma(a     )					,       deterministic=a     )
       lowercase_		:      Dict          =       hidden_states + residual

       # cross attention
       lowercase_		:      Union[str, Any]          =       hidden_states
       lowercase_		:      Dict          =       self.attna(self.norma(a     )					,       a					,       deterministic=a     )
       lowercase_		:      int          =       hidden_states + residual

       # feed forward
       lowercase_		:      Union[str, Any]          =       hidden_states
       lowercase_		:      Tuple          =       self.ff(self.norma(a     )					,       deterministic=a     )
       lowercase_		:      Optional[int]          =       hidden_states + residual

       return self.dropout_layer(a					,       deterministic=a     )


class     _UpperCAmelCase  (						nn.Module     ):
    __lowerCamelCase:				int
    __lowerCamelCase:				int
    __lowerCamelCase:				int
    __lowerCamelCase:				int						    =      1
    __lowerCamelCase:				float						    =      0.0
    __lowerCamelCase:				bool						    =      False
    __lowerCamelCase:				bool						    =      False
    __lowerCamelCase:				jnp.dtype						    =      jnp.floataa
    __lowerCamelCase:				bool						    =      False




    def 			lowerCAmelCase__			(     self					:    Optional[Any]     ):






       '''simple docstring'''






       lowercase_		:      str          =       nn.GroupNorm(num_groups=3_2					,       epsilon=1e-5     )

       lowercase_		:      Optional[Any]          =       self.n_heads * self.d_head
       if self.use_linear_projection:
          lowercase_		:      Union[str, Any]          =       nn.Dense(a					,       dtype=self.dtype     )
       else:
          lowercase_		:      List[Any]          =       nn.Conv(
              a					,       kernel_size=(1, 1)					,       strides=(1, 1)					,       padding="VALID"					,       dtype=self.dtype					,       )

       lowercase_		:      List[str]          =       [
           FlaxBasicTransformerBlock(
               a					,       self.n_heads					,       self.d_head					,       dropout=self.dropout					,       only_cross_attention=self.only_cross_attention					,       dtype=self.dtype					,       use_memory_efficient_attention=self.use_memory_efficient_attention					,       )
           for _ in range(self.depth     )
       ]

       if self.use_linear_projection:
          lowercase_		:      Optional[Any]          =       nn.Dense(a					,       dtype=self.dtype     )
       else:
          lowercase_		:      int          =       nn.Conv(
              a					,       kernel_size=(1, 1)					,       strides=(1, 1)					,       padding="VALID"					,       dtype=self.dtype					,       )

       lowercase_		:      Optional[Any]          =       nn.Dropout(rate=self.dropout     )





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






       '''simple docstring'''






       lowercase_		:      Optional[int]          =       hidden_states.shape
       lowercase_		:      Union[str, Any]          =       hidden_states
       lowercase_		:      str          =       self.norm(a     )
       if self.use_linear_projection:
          lowercase_		:      List[str]          =       hidden_states.reshape(a					,       height * width					,       a     )
          lowercase_		:      Optional[Any]          =       self.proj_in(a     )
       else:
          lowercase_		:      List[Any]          =       self.proj_in(a     )
          lowercase_		:      Any          =       hidden_states.reshape(a					,       height * width					,       a     )

       for transformer_block in self.transformer_blocks:
          lowercase_		:      int          =       transformer_block(a					,       a					,       deterministic=a     )

       if self.use_linear_projection:
          lowercase_		:      Union[str, Any]          =       self.proj_out(a     )
          lowercase_		:      Tuple          =       hidden_states.reshape(a					,       a					,       a					,       a     )
       else:
          lowercase_		:      List[Any]          =       hidden_states.reshape(a					,       a					,       a					,       a     )
          lowercase_		:      Tuple          =       self.proj_out(a     )

       lowercase_		:      Tuple          =       hidden_states + residual
       return self.dropout_layer(a					,       deterministic=a     )


class     _UpperCAmelCase  (						nn.Module     ):
    __lowerCamelCase:				int
    __lowerCamelCase:				float						    =      0.0
    __lowerCamelCase:				jnp.dtype						    =      jnp.floataa




    def 			lowerCAmelCase__			(     self					:    Tuple     ):






       '''simple docstring'''






       lowercase_		:      List[Any]          =       FlaxGEGLU(self.dim					,       self.dropout					,       self.dtype     )
       lowercase_		:      List[Any]          =       nn.Dense(self.dim					,       dtype=self.dtype     )





    def __call__(     self					:    Tuple					,       a					:    Optional[int]					,       a					:    Any=True     ):






       '''simple docstring'''






       lowercase_		:      str          =       self.net_a(a					,       deterministic=a     )
       lowercase_		:      Optional[Any]          =       self.net_a(a     )
       return hidden_states


class     _UpperCAmelCase  (						nn.Module     ):
    __lowerCamelCase:				int
    __lowerCamelCase:				float						    =      0.0
    __lowerCamelCase:				jnp.dtype						    =      jnp.floataa




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






       '''simple docstring'''






       lowercase_		:      Dict          =       self.dim * 4
       lowercase_		:      List[Any]          =       nn.Dense(inner_dim * 2					,       dtype=self.dtype     )
       lowercase_		:      List[str]          =       nn.Dropout(rate=self.dropout     )





    def __call__(     self					:    Any					,       a					:    List[Any]					,       a					:    Tuple=True     ):






       '''simple docstring'''






       lowercase_		:      Dict          =       self.proj(a     )
       lowercase_		:      Dict          =       jnp.split(a					,       2					,       axis=2     )
       return self.dropout_layer(hidden_linear * nn.gelu(a     )					,       deterministic=a     )











'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


UpperCamelCase__	=     {
    'configuration_unispeech_sat': ['UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'UniSpeechSatConfig'],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCamelCase__	=     [
     'UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST',
     'UniSpeechSatForAudioFrameClassification',
     'UniSpeechSatForCTC',
     'UniSpeechSatForPreTraining',
     'UniSpeechSatForSequenceClassification',
     'UniSpeechSatForXVector',
     'UniSpeechSatModel',
     'UniSpeechSatPreTrainedModel',
 ]

if TYPE_CHECKING:
 from .configuration_unispeech_sat import UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP, UniSpeechSatConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_unispeech_sat import (
      UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST,
      UniSpeechSatForAudioFrameClassification,
      UniSpeechSatForCTC,
      UniSpeechSatForPreTraining,
      UniSpeechSatForSequenceClassification,
      UniSpeechSatForXVector,
      UniSpeechSatModel,
      UniSpeechSatPreTrainedModel,
  )

else:
 import sys

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




