code: 
  
  
  
  
  
  
  """simple docstring"""
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  
  
  def        A	(							__snake_case:   Dict   ,    __snake_case:   Union[str, Any]   ,    __snake_case:   List[Any]   ,    __snake_case:   List[str]   ,    __snake_case:   int = 4_0_9_6     )						->		Any:
  
         """simple docstring"""
  
  
  
  
  
         __magic_name__							,     __magic_name__							,     __magic_name__										=						key.shape[-3:]
         __magic_name__										=						value.shape[-1]
         __magic_name__										=						min(__snake_case   ,    __snake_case     )
         __magic_name__										=						query / jnp.sqrt(__snake_case     )
  
         @functools.partial(jax.checkpoint   ,    prevent_cse=__snake_case     )
         def summarize_chunk(__snake_case:   int   ,    __snake_case:   Tuple   ,    __snake_case:   List[str]     ):
                __magic_name__										=						jnp.einsum('...qhd,...khd->...qhk'   ,    __snake_case   ,    __snake_case   ,    precision=__snake_case     )
  
                __magic_name__										=						jnp.max(__snake_case   ,    axis=-1   ,    keepdims=__snake_case     )
                __magic_name__										=						jax.lax.stop_gradient(__snake_case     )
                __magic_name__										=						jnp.exp(attn_weights - max_score     )
  
                __magic_name__										=						jnp.einsum('...vhf,...qhv->...qhf'   ,    __snake_case   ,    __snake_case   ,    precision=__snake_case     )
                __magic_name__										=						jnp.einsum('...qhk->...qh'   ,    __snake_case     )
  
                return (exp_values, exp_weights.sum(axis=-1     ), max_score)
  
         def chunk_scanner(__snake_case:   str     ):
                # julienne key array
                __magic_name__										=						jax.lax.dynamic_slice(
                    operand=__snake_case   ,    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
                __magic_name__										=						jax.lax.dynamic_slice(
                    operand=__snake_case   ,    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(__snake_case   ,    __snake_case   ,    __snake_case     )
  
         __magic_name__							,     __magic_name__							,     __magic_name__										=						jax.lax.map(f=__snake_case   ,    xs=jnp.arange(0   ,    __snake_case   ,    __snake_case     )     )
  
         __magic_name__										=						jnp.max(__snake_case   ,    axis=0   ,    keepdims=__snake_case     )
         __magic_name__										=						jnp.exp(chunk_max - global_max     )
  
         chunk_values *= jnp.expand_dims(__snake_case   ,    axis=-1     )
         chunk_weights *= max_diffs
  
         __magic_name__										=						chunk_values.sum(axis=0     )
         __magic_name__										=						jnp.expand_dims(__snake_case   ,    -1     ).sum(axis=0     )
  
         return all_values / all_weights
  
  
  def        A	(							__snake_case:   Any   ,    __snake_case:   Dict   ,    __snake_case:   Optional[Any]   ,    __snake_case:   Tuple=jax.lax.Precision.HIGHEST   ,    __snake_case:   int = 1_0_2_4   ,    __snake_case:   int = 4_0_9_6     )						->		Union[str, Any]:
  
         """simple docstring"""
  
  
  
  
  
         __magic_name__							,     __magic_name__							,     __magic_name__										=						query.shape[-3:]
  
         def chunk_scanner(__snake_case:   List[Any]   ,    __snake_case:   List[str]     ):
                # julienne query array
                __magic_name__										=						jax.lax.dynamic_slice(
                    operand=__snake_case   ,    start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]   ,    slice_sizes=list(query.shape[:-3]     ) + [min(__snake_case   ,    __snake_case     ), num_heads, q_features]   ,    )
  
                return (
                    chunk_idx + query_chunk_size,  # unused ignore it
                    _query_chunk_attention(
                        query=__snake_case   ,    key=__snake_case   ,    value=__snake_case   ,    precision=__snake_case   ,    key_chunk_size=__snake_case     ),
                )
  
         __magic_name__							,     __magic_name__										=						jax.lax.scan(
             f=__snake_case   ,    init=0   ,    xs=__snake_case   ,    length=math.ceil(num_q / query_chunk_size     )     )
  
         return jnp.concatenate(__snake_case   ,    axis=-3     )  # fuse the chunked result back
  
  
  
  class 					UpperCamelCase__   (							nn.Module):
  
  
  
  
  
  
  
      """simple docstring"""
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   8
      __UpperCAmelCase											=   64
      __UpperCAmelCase											=   0.0
      __UpperCAmelCase											=   False
      __UpperCAmelCase											=   jnp.floataa
  
  
  
  
  
  
  
      def 							a__ (		self   :  Union[str, Any]   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						self.dim_head * self.heads
             __magic_name__										=						self.dim_head**-0.5
  
             # Weights were exported with old names {to_q, to_k, to_v, to_out}
             __magic_name__										=						nn.Dense(UpperCamelCase_							,						use_bias=UpperCamelCase_							,						dtype=self.dtype							,						name='to_q'   )
             __magic_name__										=						nn.Dense(UpperCamelCase_							,						use_bias=UpperCamelCase_							,						dtype=self.dtype							,						name='to_k'   )
             __magic_name__										=						nn.Dense(UpperCamelCase_							,						use_bias=UpperCamelCase_							,						dtype=self.dtype							,						name='to_v'   )
  
             __magic_name__										=						nn.Dense(self.query_dim							,						dtype=self.dtype							,						name='to_out_0'   )
             __magic_name__										=						nn.Dropout(rate=self.dropout   )
  
  
  
  
  
  
  
      def 							a__ (		self   :  List[Any]							,						UpperCamelCase_   :  str   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__							,     __magic_name__							,     __magic_name__										=						tensor.shape
             __magic_name__										=						self.heads
             __magic_name__										=						tensor.reshape(UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_							,						dim // head_size   )
             __magic_name__										=						jnp.transpose(UpperCamelCase_							,						(0, 2, 1, 3)   )
             __magic_name__										=						tensor.reshape(batch_size * head_size							,						UpperCamelCase_							,						dim // head_size   )
             return tensor
  
  
  
  
  
  
  
      def 							a__ (		self   :  Optional[Any]							,						UpperCamelCase_   :  List[Any]   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__							,     __magic_name__							,     __magic_name__										=						tensor.shape
             __magic_name__										=						self.heads
             __magic_name__										=						tensor.reshape(batch_size // head_size							,						UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_   )
             __magic_name__										=						jnp.transpose(UpperCamelCase_							,						(0, 2, 1, 3)   )
             __magic_name__										=						tensor.reshape(batch_size // head_size							,						UpperCamelCase_							,						dim * head_size   )
             return tensor
  
  
  
  
  
  
  
      def __call__(		self   :  Optional[int]							,						UpperCamelCase_   :  Optional[int]							,						UpperCamelCase_   :  int=None							,						UpperCamelCase_   :  Any=True   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						hidden_states if context is None else context
  
             __magic_name__										=						self.query(UpperCamelCase_   )
             __magic_name__										=						self.key(UpperCamelCase_   )
             __magic_name__										=						self.value(UpperCamelCase_   )
  
             __magic_name__										=						self.reshape_heads_to_batch_dim(UpperCamelCase_   )
             __magic_name__										=						self.reshape_heads_to_batch_dim(UpperCamelCase_   )
             __magic_name__										=						self.reshape_heads_to_batch_dim(UpperCamelCase_   )
  
             if self.use_memory_efficient_attention:
                    __magic_name__										=						query_states.transpose(1							,						0							,						2   )
                    __magic_name__										=						key_states.transpose(1							,						0							,						2   )
                    __magic_name__										=						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
  
                    __magic_name__										=						query_states.shape[-3]
                    if flatten_latent_dim % 6_4 == 0:
                           __magic_name__										=						int(flatten_latent_dim / 6_4   )
                    elif flatten_latent_dim % 1_6 == 0:
                           __magic_name__										=						int(flatten_latent_dim / 1_6   )
                    elif flatten_latent_dim % 4 == 0:
                           __magic_name__										=						int(flatten_latent_dim / 4   )
                    else:
                           __magic_name__										=						int(UpperCamelCase_   )
  
                    __magic_name__										=						jax_memory_efficient_attention(
                        UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_							,						query_chunk_size=UpperCamelCase_							,						key_chunk_size=4_0_9_6 * 4   )
  
                    __magic_name__										=						hidden_states.transpose(1							,						0							,						2   )
             else:
                    # compute attentions
                    __magic_name__										=						jnp.einsum('b i d, b j d->b i j'							,						UpperCamelCase_							,						UpperCamelCase_   )
                    __magic_name__										=						attention_scores * self.scale
                    __magic_name__										=						nn.softmax(UpperCamelCase_							,						axis=2   )
  
                    # attend to values
                    __magic_name__										=						jnp.einsum('b i j, b j d -> b i d'							,						UpperCamelCase_							,						UpperCamelCase_   )
  
             __magic_name__										=						self.reshape_batch_dim_to_heads(UpperCamelCase_   )
             __magic_name__										=						self.proj_attn(UpperCamelCase_   )
             return self.dropout_layer(UpperCamelCase_							,						deterministic=UpperCamelCase_   )
  
  
  
  class 					UpperCamelCase__   (							nn.Module):
  
  
  
  
  
  
  
      """simple docstring"""
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   0.0
      __UpperCAmelCase											=   False
      __UpperCAmelCase											=   jnp.floataa
      __UpperCAmelCase											=   False
  
  
  
  
  
  
  
      def 							a__ (		self   :  List[Any]   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						FlaxAttention(
                 self.dim							,						self.n_heads							,						self.d_head							,						self.dropout							,						self.use_memory_efficient_attention							,						dtype=self.dtype   )
             # cross attention
             __magic_name__										=						FlaxAttention(
                 self.dim							,						self.n_heads							,						self.d_head							,						self.dropout							,						self.use_memory_efficient_attention							,						dtype=self.dtype   )
             __magic_name__										=						FlaxFeedForward(dim=self.dim							,						dropout=self.dropout							,						dtype=self.dtype   )
             __magic_name__										=						nn.LayerNorm(epsilon=1e-5							,						dtype=self.dtype   )
             __magic_name__										=						nn.LayerNorm(epsilon=1e-5							,						dtype=self.dtype   )
             __magic_name__										=						nn.LayerNorm(epsilon=1e-5							,						dtype=self.dtype   )
             __magic_name__										=						nn.Dropout(rate=self.dropout   )
  
  
  
  
  
  
  
      def __call__(		self   :  Optional[int]							,						UpperCamelCase_   :  List[Any]							,						UpperCamelCase_   :  int							,						UpperCamelCase_   :  Union[str, Any]=True   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						hidden_states
             if self.only_cross_attention:
                    __magic_name__										=						self.attna(self.norma(UpperCamelCase_   )							,						UpperCamelCase_							,						deterministic=UpperCamelCase_   )
             else:
                    __magic_name__										=						self.attna(self.norma(UpperCamelCase_   )							,						deterministic=UpperCamelCase_   )
             __magic_name__										=						hidden_states + residual
  
             # cross attention
             __magic_name__										=						hidden_states
             __magic_name__										=						self.attna(self.norma(UpperCamelCase_   )							,						UpperCamelCase_							,						deterministic=UpperCamelCase_   )
             __magic_name__										=						hidden_states + residual
  
             # feed forward
             __magic_name__										=						hidden_states
             __magic_name__										=						self.ff(self.norma(UpperCamelCase_   )							,						deterministic=UpperCamelCase_   )
             __magic_name__										=						hidden_states + residual
  
             return self.dropout_layer(UpperCamelCase_							,						deterministic=UpperCamelCase_   )
  
  
  
  class 					UpperCamelCase__   (							nn.Module):
  
  
  
  
  
  
  
      """simple docstring"""
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   1
      __UpperCAmelCase											=   0.0
      __UpperCAmelCase											=   False
      __UpperCAmelCase											=   False
      __UpperCAmelCase											=   jnp.floataa
      __UpperCAmelCase											=   False
  
  
  
  
  
  
  
      def 							a__ (		self   :  int   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						nn.GroupNorm(num_groups=3_2							,						epsilon=1e-5   )
  
             __magic_name__										=						self.n_heads * self.d_head
             if self.use_linear_projection:
                    __magic_name__										=						nn.Dense(UpperCamelCase_							,						dtype=self.dtype   )
             else:
                    __magic_name__										=						nn.Conv(
                        UpperCamelCase_							,						kernel_size=(1, 1)							,						strides=(1, 1)							,						padding='VALID'							,						dtype=self.dtype							,						)
  
             __magic_name__										=						[
                 FlaxBasicTransformerBlock(
                     UpperCamelCase_							,						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:
                    __magic_name__										=						nn.Dense(UpperCamelCase_							,						dtype=self.dtype   )
             else:
                    __magic_name__										=						nn.Conv(
                        UpperCamelCase_							,						kernel_size=(1, 1)							,						strides=(1, 1)							,						padding='VALID'							,						dtype=self.dtype							,						)
  
             __magic_name__										=						nn.Dropout(rate=self.dropout   )
  
  
  
  
  
  
  
      def __call__(		self   :  str							,						UpperCamelCase_   :  List[str]							,						UpperCamelCase_   :  Optional[Any]							,						UpperCamelCase_   :  Tuple=True   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__							,     __magic_name__							,     __magic_name__							,     __magic_name__										=						hidden_states.shape
             __magic_name__										=						hidden_states
             __magic_name__										=						self.norm(UpperCamelCase_   )
             if self.use_linear_projection:
                    __magic_name__										=						hidden_states.reshape(UpperCamelCase_							,						height * width							,						UpperCamelCase_   )
                    __magic_name__										=						self.proj_in(UpperCamelCase_   )
             else:
                    __magic_name__										=						self.proj_in(UpperCamelCase_   )
                    __magic_name__										=						hidden_states.reshape(UpperCamelCase_							,						height * width							,						UpperCamelCase_   )
  
             for transformer_block in self.transformer_blocks:
                    __magic_name__										=						transformer_block(UpperCamelCase_							,						UpperCamelCase_							,						deterministic=UpperCamelCase_   )
  
             if self.use_linear_projection:
                    __magic_name__										=						self.proj_out(UpperCamelCase_   )
                    __magic_name__										=						hidden_states.reshape(UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_   )
             else:
                    __magic_name__										=						hidden_states.reshape(UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_							,						UpperCamelCase_   )
                    __magic_name__										=						self.proj_out(UpperCamelCase_   )
  
             __magic_name__										=						hidden_states + residual
             return self.dropout_layer(UpperCamelCase_							,						deterministic=UpperCamelCase_   )
  
  
  
  class 					UpperCamelCase__   (							nn.Module):
  
  
  
  
  
  
  
      """simple docstring"""
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   0.0
      __UpperCAmelCase											=   jnp.floataa
  
  
  
  
  
  
  
      def 							a__ (		self   :  Optional[Any]   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						FlaxGEGLU(self.dim							,						self.dropout							,						self.dtype   )
             __magic_name__										=						nn.Dense(self.dim							,						dtype=self.dtype   )
  
  
  
  
  
  
  
      def __call__(		self   :  Union[str, Any]							,						UpperCamelCase_   :  List[str]							,						UpperCamelCase_   :  Union[str, Any]=True   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						self.net_a(UpperCamelCase_							,						deterministic=UpperCamelCase_   )
             __magic_name__										=						self.net_a(UpperCamelCase_   )
             return hidden_states
  
  
  
  class 					UpperCamelCase__   (							nn.Module):
  
  
  
  
  
  
  
      """simple docstring"""
      __UpperCAmelCase											=   42
      __UpperCAmelCase											=   0.0
      __UpperCAmelCase											=   jnp.floataa
  
  
  
  
  
  
  
      def 							a__ (		self   :  Any   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						self.dim * 4
             __magic_name__										=						nn.Dense(inner_dim * 2							,						dtype=self.dtype   )
             __magic_name__										=						nn.Dropout(rate=self.dropout   )
  
  
  
  
  
  
  
      def __call__(		self   :  Tuple							,						UpperCamelCase_   :  Tuple							,						UpperCamelCase_   :  Dict=True   ):
  
  
  
  
  
             '''simple docstring'''
  
             __magic_name__										=						self.proj(UpperCamelCase_   )
             __magic_name__							,     __magic_name__										=						jnp.split(UpperCamelCase_							,						2							,						axis=2   )
             return self.dropout_layer(hidden_linear * nn.gelu(UpperCamelCase_   )							,						deterministic=UpperCamelCase_   )
style_context: 
  
  
  
  
  
  
  """simple docstring"""
  from typing import TYPE_CHECKING
  
  from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
  
  
  snake_case   :   str    		=		{"""configuration_sew_d""": ["""SEW_D_PRETRAINED_CONFIG_ARCHIVE_MAP""", """SEWDConfig"""]}
  
  try:
   if not is_torch_available():
    raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
   pass
  else:
   snake_case   :   Union[str, Any]    		=		[
       """SEW_D_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """SEWDForCTC""",
       """SEWDForSequenceClassification""",
       """SEWDModel""",
       """SEWDPreTrainedModel""",
   ]
  
  if TYPE_CHECKING:
   from .configuration_sew_d import SEW_D_PRETRAINED_CONFIG_ARCHIVE_MAP, SEWDConfig
  
   try:
    if not is_torch_available():
     raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
    pass
   else:
    from .modeling_sew_d import (
        SEW_D_PRETRAINED_MODEL_ARCHIVE_LIST,
        SEWDForCTC,
        SEWDForSequenceClassification,
        SEWDModel,
        SEWDPreTrainedModel,
    )
  
  
  else:
   import sys
  
   snake_case   :   Union[str, Any]    		=		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)