





"""simple docstring"""
import functools
import math

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






def    A__							(     UpperCamelCase   ,    UpperCamelCase   ,    UpperCamelCase   ,    UpperCamelCase   ,    UpperCamelCase = 4_096			):
							A,      A,      A						=    key.shape[-3:]
							A						=    value.shape[-1]
							A						=    min(UpperCamelCase   ,    UpperCamelCase			)
							A						=    query / jnp.sqrt(UpperCamelCase			)

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

														A						=    jnp.max(UpperCamelCase   ,    axis=-1   ,    keepdims=UpperCamelCase			)
														A						=    jax.lax.stop_gradient(UpperCamelCase			)
														A						=    jnp.exp(attn_weights - max_score			)

														A						=    jnp.einsum("...vhf,...qhv->...qhf"   ,    UpperCamelCase   ,    UpperCamelCase   ,    precision=UpperCamelCase			)
														A						=    jnp.einsum("...qhk->...qh"   ,    UpperCamelCase			)

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

							def chunk_scanner(UpperCamelCase			):
														# julienne key array
														A						=    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
														A						=    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			)

							A,      A,      A						=    jax.lax.map(f=UpperCamelCase   ,    xs=jnp.arange(0   ,    UpperCamelCase   ,    UpperCamelCase			)			)

							A						=    jnp.max(UpperCamelCase   ,    axis=0   ,    keepdims=UpperCamelCase			)
							A						=    jnp.exp(chunk_max - global_max			)

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

							A						=    chunk_values.sum(axis=0			)
							A						=    jnp.expand_dims(UpperCamelCase   ,    -1			).sum(axis=0			)

							return all_values / all_weights






def    A__							(     UpperCamelCase   ,    UpperCamelCase   ,    UpperCamelCase   ,    UpperCamelCase=jax.lax.Precision.HIGHEST   ,    UpperCamelCase = 1_024   ,    UpperCamelCase = 4_096			):
							A,      A,      A						=    query.shape[-3:]

							def chunk_scanner(UpperCamelCase   ,    UpperCamelCase			):
														# julienne query array
														A						=    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			),
														)

							A,      A						=    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					):
					UpperCamelCase						=				42
					UpperCamelCase						=				8
					UpperCamelCase						=				6_4
					UpperCamelCase						=				0.0
					UpperCamelCase						=				False
					UpperCamelCase						=				jnp.floataa


					def       lowerCamelCase			(   self     :Union[str, Any]						):
												A						=    self.dim_head * self.heads
												A						=    self.dim_head**-0.5

												# Weights were exported with old names {to_q, to_k, to_v, to_out}
												A						=    nn.Dense(__UpperCamelCase	,			use_bias=__UpperCamelCase	,			dtype=self.dtype	,			name="to_q"						)
												A						=    nn.Dense(__UpperCamelCase	,			use_bias=__UpperCamelCase	,			dtype=self.dtype	,			name="to_k"						)
												A						=    nn.Dense(__UpperCamelCase	,			use_bias=__UpperCamelCase	,			dtype=self.dtype	,			name="to_v"						)

												A						=    nn.Dense(self.query_dim	,			dtype=self.dtype	,			name="to_out_0"						)
												A						=    nn.Dropout(rate=self.dropout						)


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


					def       lowerCamelCase			(   self     :Union[str, Any]	,			__UpperCamelCase     :Optional[Any]						):
												A,      A,      A						=    tensor.shape
												A						=    self.heads
												A						=    tensor.reshape(batch_size // head_size	,			__UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase						)
												A						=    jnp.transpose(__UpperCamelCase	,			(0, 2, 1, 3)						)
												A						=    tensor.reshape(batch_size // head_size	,			__UpperCamelCase	,			dim * head_size						)
												return tensor


					def __call__(   self     :List[str]	,			__UpperCamelCase     :Dict	,			__UpperCamelCase     :Union[str, Any]=None	,			__UpperCamelCase     :Any=True						):
												A						=    hidden_states if context is None else context

												A						=    self.query(__UpperCamelCase						)
												A						=    self.key(__UpperCamelCase						)
												A						=    self.value(__UpperCamelCase						)

												A						=    self.reshape_heads_to_batch_dim(__UpperCamelCase						)
												A						=    self.reshape_heads_to_batch_dim(__UpperCamelCase						)
												A						=    self.reshape_heads_to_batch_dim(__UpperCamelCase						)

												if self.use_memory_efficient_attention:
																			A						=    query_states.transpose(1	,			0	,			2						)
																			A						=    key_states.transpose(1	,			0	,			2						)
																			A						=    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

																			A						=    query_states.shape[-3]
																			if flatten_latent_dim % 64 == 0:
																										A						=    int(flatten_latent_dim / 64						)
																			elif flatten_latent_dim % 16 == 0:
																										A						=    int(flatten_latent_dim / 16						)
																			elif flatten_latent_dim % 4 == 0:
																										A						=    int(flatten_latent_dim / 4						)
																			else:
																										A						=    int(__UpperCamelCase						)

																			A						=    jax_memory_efficient_attention(
																			    __UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase	,			query_chunk_size=__UpperCamelCase	,			key_chunk_size=40_96 * 4						)

																			A						=    hidden_states.transpose(1	,			0	,			2						)
												else:
																			# compute attentions
																			A						=    jnp.einsum("b i d, b j d->b i j"	,			__UpperCamelCase	,			__UpperCamelCase						)
																			A						=    attention_scores * self.scale
																			A						=    nn.softmax(__UpperCamelCase	,			axis=2						)

																			# attend to values
																			A						=    jnp.einsum("b i j, b j d -> b i d"	,			__UpperCamelCase	,			__UpperCamelCase						)

												A						=    self.reshape_batch_dim_to_heads(__UpperCamelCase						)
												A						=    self.proj_attn(__UpperCamelCase						)
												return self.dropout_layer(__UpperCamelCase	,			deterministic=__UpperCamelCase						)







class    _UpperCAmelCase  (      nn.Module					):
					UpperCamelCase						=				42
					UpperCamelCase						=				42
					UpperCamelCase						=				42
					UpperCamelCase						=				0.0
					UpperCamelCase						=				False
					UpperCamelCase						=				jnp.floataa
					UpperCamelCase						=				False


					def       lowerCamelCase			(   self     :int						):
												# self attention (or cross_attention if only_cross_attention is True)
												A						=    FlaxAttention(
												    self.dim	,			self.n_heads	,			self.d_head	,			self.dropout	,			self.use_memory_efficient_attention	,			dtype=self.dtype						)
												# cross attention
												A						=    FlaxAttention(
												    self.dim	,			self.n_heads	,			self.d_head	,			self.dropout	,			self.use_memory_efficient_attention	,			dtype=self.dtype						)
												A						=    FlaxFeedForward(dim=self.dim	,			dropout=self.dropout	,			dtype=self.dtype						)
												A						=    nn.LayerNorm(epsilon=1e-5	,			dtype=self.dtype						)
												A						=    nn.LayerNorm(epsilon=1e-5	,			dtype=self.dtype						)
												A						=    nn.LayerNorm(epsilon=1e-5	,			dtype=self.dtype						)
												A						=    nn.Dropout(rate=self.dropout						)


					def __call__(   self     :List[str]	,			__UpperCamelCase     :int	,			__UpperCamelCase     :Any	,			__UpperCamelCase     :Any=True						):
												# self attention
												A						=    hidden_states
												if self.only_cross_attention:
																			A						=    self.attna(self.norma(__UpperCamelCase						)	,			__UpperCamelCase	,			deterministic=__UpperCamelCase						)
												else:
																			A						=    self.attna(self.norma(__UpperCamelCase						)	,			deterministic=__UpperCamelCase						)
												A						=    hidden_states + residual

												# cross attention
												A						=    hidden_states
												A						=    self.attna(self.norma(__UpperCamelCase						)	,			__UpperCamelCase	,			deterministic=__UpperCamelCase						)
												A						=    hidden_states + residual

												# feed forward
												A						=    hidden_states
												A						=    self.ff(self.norma(__UpperCamelCase						)	,			deterministic=__UpperCamelCase						)
												A						=    hidden_states + residual

												return self.dropout_layer(__UpperCamelCase	,			deterministic=__UpperCamelCase						)







class    _UpperCAmelCase  (      nn.Module					):
					UpperCamelCase						=				42
					UpperCamelCase						=				42
					UpperCamelCase						=				42
					UpperCamelCase						=				1
					UpperCamelCase						=				0.0
					UpperCamelCase						=				False
					UpperCamelCase						=				False
					UpperCamelCase						=				jnp.floataa
					UpperCamelCase						=				False


					def       lowerCamelCase			(   self     :Optional[Any]						):
												A						=    nn.GroupNorm(num_groups=32	,			epsilon=1e-5						)

												A						=    self.n_heads * self.d_head
												if self.use_linear_projection:
																			A						=    nn.Dense(__UpperCamelCase	,			dtype=self.dtype						)
												else:
																			A						=    nn.Conv(
																			    __UpperCamelCase	,			kernel_size=(1, 1)	,			strides=(1, 1)	,			padding="VALID"	,			dtype=self.dtype	,			)

												A						=    [
												    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:
																			A						=    nn.Dense(__UpperCamelCase	,			dtype=self.dtype						)
												else:
																			A						=    nn.Conv(
																			    __UpperCamelCase	,			kernel_size=(1, 1)	,			strides=(1, 1)	,			padding="VALID"	,			dtype=self.dtype	,			)

												A						=    nn.Dropout(rate=self.dropout						)


					def __call__(   self     :Dict	,			__UpperCamelCase     :List[str]	,			__UpperCamelCase     :Union[str, Any]	,			__UpperCamelCase     :List[Any]=True						):
												A,      A,      A,      A						=    hidden_states.shape
												A						=    hidden_states
												A						=    self.norm(__UpperCamelCase						)
												if self.use_linear_projection:
																			A						=    hidden_states.reshape(__UpperCamelCase	,			height * width	,			__UpperCamelCase						)
																			A						=    self.proj_in(__UpperCamelCase						)
												else:
																			A						=    self.proj_in(__UpperCamelCase						)
																			A						=    hidden_states.reshape(__UpperCamelCase	,			height * width	,			__UpperCamelCase						)

												for transformer_block in self.transformer_blocks:
																			A						=    transformer_block(__UpperCamelCase	,			__UpperCamelCase	,			deterministic=__UpperCamelCase						)

												if self.use_linear_projection:
																			A						=    self.proj_out(__UpperCamelCase						)
																			A						=    hidden_states.reshape(__UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase						)
												else:
																			A						=    hidden_states.reshape(__UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase	,			__UpperCamelCase						)
																			A						=    self.proj_out(__UpperCamelCase						)

												A						=    hidden_states + residual
												return self.dropout_layer(__UpperCamelCase	,			deterministic=__UpperCamelCase						)







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


					def       lowerCamelCase			(   self     :Optional[Any]						):
												# The second linear layer needs to be called
												# net_2 for now to match the index of the Sequential layer
												A						=    FlaxGEGLU(self.dim	,			self.dropout	,			self.dtype						)
												A						=    nn.Dense(self.dim	,			dtype=self.dtype						)


					def __call__(   self     :Any	,			__UpperCamelCase     :Optional[int]	,			__UpperCamelCase     :Any=True						):
												A						=    self.net_a(__UpperCamelCase	,			deterministic=__UpperCamelCase						)
												A						=    self.net_a(__UpperCamelCase						)
												return hidden_states







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


					def       lowerCamelCase			(   self     :Any						):
												A						=    self.dim * 4
												A						=    nn.Dense(inner_dim * 2	,			dtype=self.dtype						)
												A						=    nn.Dropout(rate=self.dropout						)


					def __call__(   self     :Union[str, Any]	,			__UpperCamelCase     :Any	,			__UpperCamelCase     :List[str]=True						):
												A						=    self.proj(__UpperCamelCase						)
												A,      A						=    jnp.split(__UpperCamelCase	,			2	,			axis=2						)
												return self.dropout_layer(hidden_linear * nn.gelu(__UpperCamelCase						)	,			deterministic=__UpperCamelCase						)









"""simple docstring"""
from typing import TYPE_CHECKING

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


_snake_case	:						Any					=					{
    'configuration_unispeech_sat': ['UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'UniSpeechSatConfig'],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_snake_case	:						str					=					[
		    '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

		_snake_case	:						Union[str, Any]					=					_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


