



import functools
import math

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

def      A_ (		__a       :			Optional[int]	,						__a       :			Tuple	,						__a       :			Tuple	,						__a       :			Optional[Any]	,						__a       :			int = 4_096						):

				"""simple docstring"""
				a__ ,     a__ ,     a__			 =						key.shape[-3:]
				a__			 =						value.shape[-1]
				a__			 =						min(__a	,						__a						)
				a__			 =						query / jnp.sqrt(__a						)

				@functools.partial(jax.checkpoint	,						prevent_cse=__a						)
				def summarize_chunk(__a       :			Tuple	,						__a       :			List[str]	,						__a       :			Tuple						):
								a__			 =						jnp.einsum("""...qhd,...khd->...qhk"""	,						__a	,						__a	,						precision=__a						)

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

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

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

				def chunk_scanner(__a       :			int						):
								# julienne key array
								a__			 =						jax.lax.dynamic_slice(
								    operand=__a	,						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=__a	,						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(__a	,						__a	,						__a						)

				a__ ,     a__ ,     a__			 =						jax.lax.map(f=__a	,						xs=jnp.arange(0	,						__a	,						__a						)						)

				a__			 =						jnp.max(__a	,						axis=0	,						keepdims=__a						)
				a__			 =						jnp.exp(chunk_max - global_max						)

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

				a__			 =						chunk_values.sum(axis=0						)
				a__			 =						jnp.expand_dims(__a	,						-1						).sum(axis=0						)

				return all_values / all_weights

def      A_ (		__a       :			int	,						__a       :			int	,						__a       :			Tuple	,						__a       :			int=jax.lax.Precision.HIGHEST	,						__a       :			int = 1_024	,						__a       :			int = 4_096						):

				"""simple docstring"""
				a__ ,     a__ ,     a__			 =						query.shape[-3:]

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

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

				a__ ,     a__			 =						jax.lax.scan(
				    f=__a	,						init=0	,						xs=__a	,						length=math.ceil(num_q / query_chunk_size						)						)

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

class 	__snake_case (				nn.Module):







							'''simple docstring'''






							UpperCamelCase__			: int
							UpperCamelCase__			: int             =   8
							UpperCamelCase__			: int             =   6_4
							UpperCamelCase__			: float             =   0.0
							UpperCamelCase__			: bool             =   False
							UpperCamelCase__			: jnp.dtype             =   jnp.floataa


							def _a (							self						):
											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(a_					,					use_bias=a_					,					dtype=self.dtype					,					name="""to_q"""						)
											a__			 =						nn.Dense(a_					,					use_bias=a_					,					dtype=self.dtype					,					name="""to_k"""						)
											a__			 =						nn.Dense(a_					,					use_bias=a_					,					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 _a (							self					,					a_						):
											a__ ,     a__ ,     a__			 =						tensor.shape
											a__			 =						self.heads
											a__			 =						tensor.reshape(a_					,					a_					,					a_					,					dim // head_size						)
											a__			 =						jnp.transpose(a_					,					(0, 2, 1, 3)						)
											a__			 =						tensor.reshape(batch_size * head_size					,					a_					,					dim // head_size						)
											return tensor


							def _a (							self					,					a_						):
											a__ ,     a__ ,     a__			 =						tensor.shape
											a__			 =						self.heads
											a__			 =						tensor.reshape(batch_size // head_size					,					a_					,					a_					,					a_						)
											a__			 =						jnp.transpose(a_					,					(0, 2, 1, 3)						)
											a__			 =						tensor.reshape(batch_size // head_size					,					a_					,					dim * head_size						)
											return tensor







							def __call__(							self					,					a_					,					a_=None					,					a_=True						):
											a__			 =						hidden_states if context is None else context

											a__			 =						self.query(a_						)
											a__			 =						self.key(a_						)
											a__			 =						self.value(a_						)

											a__			 =						self.reshape_heads_to_batch_dim(a_						)
											a__			 =						self.reshape_heads_to_batch_dim(a_						)
											a__			 =						self.reshape_heads_to_batch_dim(a_						)

											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(a_						)

															a__			 =						jax_memory_efficient_attention(
															    a_					,					a_					,					a_					,					query_chunk_size=a_					,					key_chunk_size=4_096 * 4						)

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

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

											a__			 =						self.reshape_batch_dim_to_heads(a_						)
											a__			 =						self.proj_attn(a_						)
											return self.dropout_layer(a_					,					deterministic=a_						)




class 	__snake_case (				nn.Module):







							'''simple docstring'''






							UpperCamelCase__			: int
							UpperCamelCase__			: int
							UpperCamelCase__			: int
							UpperCamelCase__			: float             =   0.0
							UpperCamelCase__			: bool             =   False
							UpperCamelCase__			: jnp.dtype             =   jnp.floataa
							UpperCamelCase__			: bool             =   False


							def _a (							self						):
											# 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					,					a_					,					a_					,					a_=True						):
											# self attention
											a__			 =						hidden_states
											if self.only_cross_attention:
															a__			 =						self.attna(self.norma(a_						)					,					a_					,					deterministic=a_						)
											else:
															a__			 =						self.attna(self.norma(a_						)					,					deterministic=a_						)
											a__			 =						hidden_states + residual

											# cross attention
											a__			 =						hidden_states
											a__			 =						self.attna(self.norma(a_						)					,					a_					,					deterministic=a_						)
											a__			 =						hidden_states + residual

											# feed forward
											a__			 =						hidden_states
											a__			 =						self.ff(self.norma(a_						)					,					deterministic=a_						)
											a__			 =						hidden_states + residual

											return self.dropout_layer(a_					,					deterministic=a_						)




class 	__snake_case (				nn.Module):







							'''simple docstring'''






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


							def _a (							self						):
											a__			 =						nn.GroupNorm(num_groups=32					,					epsilon=1E-5						)

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

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

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







							def __call__(							self					,					a_					,					a_					,					a_=True						):
											a__ ,     a__ ,     a__ ,     a__			 =						hidden_states.shape
											a__			 =						hidden_states
											a__			 =						self.norm(a_						)
											if self.use_linear_projection:
															a__			 =						hidden_states.reshape(a_					,					height * width					,					a_						)
															a__			 =						self.proj_in(a_						)
											else:
															a__			 =						self.proj_in(a_						)
															a__			 =						hidden_states.reshape(a_					,					height * width					,					a_						)

											for transformer_block in self.transformer_blocks:
															a__			 =						transformer_block(a_					,					a_					,					deterministic=a_						)

											if self.use_linear_projection:
															a__			 =						self.proj_out(a_						)
															a__			 =						hidden_states.reshape(a_					,					a_					,					a_					,					a_						)
											else:
															a__			 =						hidden_states.reshape(a_					,					a_					,					a_					,					a_						)
															a__			 =						self.proj_out(a_						)

											a__			 =						hidden_states + residual
											return self.dropout_layer(a_					,					deterministic=a_						)




class 	__snake_case (				nn.Module):







							'''simple docstring'''






							UpperCamelCase__			: int
							UpperCamelCase__			: float             =   0.0
							UpperCamelCase__			: jnp.dtype             =   jnp.floataa


							def _a (							self						):
											# 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					,					a_					,					a_=True						):
											a__			 =						self.net_a(a_					,					deterministic=a_						)
											a__			 =						self.net_a(a_						)
											return hidden_states




class 	__snake_case (				nn.Module):







							'''simple docstring'''






							UpperCamelCase__			: int
							UpperCamelCase__			: float             =   0.0
							UpperCamelCase__			: jnp.dtype             =   jnp.floataa


							def _a (							self						):
											a__			 =						self.dim * 4
											a__			 =						nn.Dense(inner_dim * 2					,					dtype=self.dtype						)
											a__			 =						nn.Dropout(rate=self.dropout						)







							def __call__(							self					,					a_					,					a_=True						):
											a__			 =						self.proj(a_						)
											a__ ,     a__			 =						jnp.split(a_					,					2					,					axis=2						)
											return self.dropout_layer(hidden_linear * nn.gelu(a_						)					,					deterministic=a_						)






import json
import os
import shutil
import tempfile
import unittest

from transformers import is_speech_available, is_torch_available
from transformers.models.speechta import SpeechTaTokenizer
from transformers.testing_utils import get_tests_dir, require_torch
from transformers.utils import FEATURE_EXTRACTOR_NAME


if is_speech_available() and is_torch_available():
	from transformers import SpeechTaFeatureExtractor, SpeechTaProcessor

	from .test_feature_extraction_speechta import floats_list


UpperCAmelCase								=   get_tests_dir("""fixtures/test_sentencepiece_bpe_char.model""")

@require_torch
class 	__snake_case (				unittest.TestCase):







							'''simple docstring'''






							def _a (							self						):
											a__			 =						tempfile.mkdtemp()

											a__			 =						SpeechTaTokenizer(a_						)
											tokenizer.save_pretrained(self.tmpdirname						)

											a__			 =						{
											    """feature_size""": 1,
											    """padding_value""": 0.0,
											    """sampling_rate""": 16_000,
											    """do_normalize""": False,
											    """num_mel_bins""": 80,
											    """hop_length""": 16,
											    """win_length""": 64,
											    """win_function""": """hann_window""",
											    """fmin""": 80,
											    """fmax""": 7_600,
											    """mel_floor""": 1E-1_0,
											    """reduction_factor""": 2,
											    """return_attention_mask""": True,
											}

											a__			 =						os.path.join(self.tmpdirname					,					a_						)
											with open(self.feature_extraction_file					,					"""w"""					,					encoding="""utf-8"""						) as fp:
															fp.write(json.dumps(a_						) + """\n"""						)


							def _a (							self					,					**a_						):
											return SpeechTaTokenizer.from_pretrained(self.tmpdirname					,					**a_						)


							def _a (							self					,					**a_						):
											return SpeechTaFeatureExtractor.from_pretrained(self.tmpdirname					,					**a_						)


							def _a (							self						):
											shutil.rmtree(self.tmpdirname						)


							def _a (							self						):
											a__			 =						self.get_tokenizer()
											a__			 =						self.get_feature_extractor()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											processor.save_pretrained(self.tmpdirname						)
											a__			 =						SpeechTaProcessor.from_pretrained(self.tmpdirname						)

											self.assertEqual(processor.tokenizer.get_vocab()					,					tokenizer.get_vocab()						)
											self.assertIsInstance(processor.tokenizer					,					a_						)

											self.assertEqual(processor.feature_extractor.to_json_string()					,					feature_extractor.to_json_string()						)
											self.assertIsInstance(processor.feature_extractor					,					a_						)


							def _a (							self						):
											a__			 =						SpeechTaProcessor(tokenizer=self.get_tokenizer()					,					feature_extractor=self.get_feature_extractor()						)
											processor.save_pretrained(self.tmpdirname						)

											a__			 =						self.get_tokenizer(bos_token="""(BOS)"""					,					eos_token="""(EOS)"""						)
											a__			 =						self.get_feature_extractor(do_normalize=a_					,					padding_value=1.0						)

											a__			 =						SpeechTaProcessor.from_pretrained(
											    self.tmpdirname					,					bos_token="""(BOS)"""					,					eos_token="""(EOS)"""					,					do_normalize=a_					,					padding_value=1.0						)

											self.assertEqual(processor.tokenizer.get_vocab()					,					tokenizer_add_kwargs.get_vocab()						)
											self.assertIsInstance(processor.tokenizer					,					a_						)

											self.assertEqual(processor.feature_extractor.to_json_string()					,					feature_extractor_add_kwargs.to_json_string()						)
											self.assertIsInstance(processor.feature_extractor					,					a_						)


							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						floats_list((3, 1_000)						)

											a__			 =						feature_extractor(audio=a_					,					return_tensors="""np"""						)
											a__			 =						processor(audio=a_					,					return_tensors="""np"""						)

											for key in input_feat_extract.keys():
															self.assertAlmostEqual(input_feat_extract[key].sum()					,					input_processor[key].sum()					,					delta=1E-2						)


							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						floats_list((3, 1_000)						)

											a__			 =						feature_extractor(audio_target=a_					,					return_tensors="""np"""						)
											a__			 =						processor(audio_target=a_					,					return_tensors="""np"""						)

											for key in input_feat_extract.keys():
															self.assertAlmostEqual(input_feat_extract[key].sum()					,					input_processor[key].sum()					,					delta=1E-2						)


							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						"""This is a test string"""

											a__			 =						processor(text=a_						)
											a__			 =						tokenizer(a_						)

											for key in encoded_tok.keys():
															self.assertListEqual(encoded_tok[key]					,					encoded_processor[key]						)


							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						"""This is a test string"""

											a__			 =						processor(text_target=a_						)
											a__			 =						tokenizer(a_						)

											for key in encoded_tok.keys():
															self.assertListEqual(encoded_tok[key]					,					encoded_processor[key]						)


							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											a__			 =						[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

											a__			 =						processor.batch_decode(a_						)
											a__			 =						tokenizer.batch_decode(a_						)

											self.assertListEqual(a_					,					a_						)







							def _a (							self						):
											a__			 =						self.get_feature_extractor()
											a__			 =						self.get_tokenizer()

											a__			 =						SpeechTaProcessor(tokenizer=a_					,					feature_extractor=a_						)

											self.assertListEqual(
											    processor.model_input_names					,					feature_extractor.model_input_names					,					msg="""`processor` and `feature_extractor` model input names do not match"""					,					)

