





import functools
import math

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


def 				SCREAMING_SNAKE_CASE__    (     __a ,			__a ,			__a ,			__a ,			__a = 40_96 ):
	snake_case_    ,snake_case_    ,snake_case_		:      Optional[int]           =						key.shape[-3:]
	snake_case_		:      int           =						value.shape[-1]
	snake_case_		:      Tuple           =						min(__a ,			__a )
	snake_case_		:      int           =						query / jnp.sqrt(__a )

	@functools.partial(jax.checkpoint ,			prevent_cse=__a )
	def summarize_chunk(__a ,			__a ,			__a ):
		snake_case_		:      Tuple           =						jnp.einsum('...qhd,...khd->...qhk' ,			__a ,			__a ,			precision=__a )

		snake_case_		:      List[Any]           =						jnp.max(__a ,			axis=-1 ,			keepdims=__a )
		snake_case_		:      Optional[Any]           =						jax.lax.stop_gradient(__a )
		snake_case_		:      Union[str, Any]           =						jnp.exp(attn_weights - max_score )

		snake_case_		:      str           =						jnp.einsum('...vhf,...qhv->...qhf' ,			__a ,			__a ,			precision=__a )
		snake_case_		:      Optional[Any]           =						jnp.einsum('...qhk->...qh' ,			__a )

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

	def chunk_scanner(__a ):
		# julienne key array
		snake_case_		:      List[Any]           =						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
		snake_case_		:      List[str]           =						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 )

	snake_case_    ,snake_case_    ,snake_case_		:      Dict           =						jax.lax.map(f=__a ,			xs=jnp.arange(0 ,			__a ,			__a ) )

	snake_case_		:      List[str]           =						jnp.max(__a ,			axis=0 ,			keepdims=__a )
	snake_case_		:      Tuple           =						jnp.exp(chunk_max - global_max )

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

	snake_case_		:      str           =						chunk_values.sum(axis=0 )
	snake_case_		:      List[str]           =						jnp.expand_dims(__a ,			-1 ).sum(axis=0 )

	return all_values / all_weights


def 				SCREAMING_SNAKE_CASE__    (     __a ,			__a ,			__a ,			__a=jax.lax.Precision.HIGHEST ,			__a = 10_24 ,			__a = 40_96 ):
	snake_case_    ,snake_case_    ,snake_case_		:      Union[str, Any]           =						query.shape[-3:]

	def chunk_scanner(__a ,			__a ):
		# julienne query array
		snake_case_		:      Union[str, Any]           =						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 ),
		)

	snake_case_    ,snake_case_		:      Tuple           =						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 SCREAMING_SNAKE_CASE_    (       nn.Module       ):
						__magic_name__:      int
						__magic_name__:      int			   =							8
						__magic_name__:      int			   =							64
						__magic_name__:      float			   =							0.0
						__magic_name__:      bool			   =							False
						__magic_name__:      jnp.dtype			   =							jnp.floataa






						def UpperCAmelCase_					(   self   :   Any						)				->     List[str]:







							"""simple docstring"""


							snake_case_		:      Union[str, Any]           =						self.dim_head * self.heads
							snake_case_		:      Tuple           =						self.dim_head**-0.5

							# Weights were exported with old names {to_q, to_k, to_v, to_out}
							snake_case_		:      str           =						nn.Dense(_A		, use_bias=_A		, dtype=self.dtype		, name='to_q'						)
							snake_case_		:      Optional[int]           =						nn.Dense(_A		, use_bias=_A		, dtype=self.dtype		, name='to_k'						)
							snake_case_		:      List[str]           =						nn.Dense(_A		, use_bias=_A		, dtype=self.dtype		, name='to_v'						)

							snake_case_		:      Union[str, Any]           =						nn.Dense(self.query_dim		, dtype=self.dtype		, name='to_out_0'						)
							snake_case_		:      Any           =						nn.Dropout(rate=self.dropout						)






						def UpperCAmelCase_					(   self   :   Union[str, Any]		, _A   :   Optional[int]						)				->     str:







							"""simple docstring"""


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






						def UpperCAmelCase_					(   self   :   Union[str, Any]		, _A   :   Any						)				->     Optional[Any]:







							"""simple docstring"""


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






						def __call__(   self   :   Optional[int]		, _A   :   Optional[int]		, _A   :   str=None		, _A   :   Dict=True						)				->     Any:







							"""simple docstring"""


							snake_case_		:      Tuple           =						hidden_states if context is None else context

							snake_case_		:      Any           =						self.query(_A						)
							snake_case_		:      Dict           =						self.key(_A						)
							snake_case_		:      int           =						self.value(_A						)

							snake_case_		:      Optional[Any]           =						self.reshape_heads_to_batch_dim(_A						)
							snake_case_		:      Union[str, Any]           =						self.reshape_heads_to_batch_dim(_A						)
							snake_case_		:      Any           =						self.reshape_heads_to_batch_dim(_A						)

							if self.use_memory_efficient_attention:
								snake_case_		:      Optional[Any]           =						query_states.transpose(1		, 0		, 2						)
								snake_case_		:      int           =						key_states.transpose(1		, 0		, 2						)
								snake_case_		:      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

								snake_case_		:      int           =						query_states.shape[-3]
								if flatten_latent_dim % 64 == 0:
									snake_case_		:      int           =						int(flatten_latent_dim / 64						)
								elif flatten_latent_dim % 16 == 0:
									snake_case_		:      Optional[int]           =						int(flatten_latent_dim / 16						)
								elif flatten_latent_dim % 4 == 0:
									snake_case_		:      List[str]           =						int(flatten_latent_dim / 4						)
								else:
									snake_case_		:      List[str]           =						int(_A						)

								snake_case_		:      Tuple           =						jax_memory_efficient_attention(
								    _A		, _A		, _A		, query_chunk_size=_A		, key_chunk_size=4096 * 4						)

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

								# attend to values
								snake_case_		:      int           =						jnp.einsum('b i j, b j d -> b i d'		, _A		, _A						)

							snake_case_		:      str           =						self.reshape_batch_dim_to_heads(_A						)
							snake_case_		:      Union[str, Any]           =						self.proj_attn(_A						)
							return self.dropout_layer(_A		, deterministic=_A						)
class SCREAMING_SNAKE_CASE_    (       nn.Module       ):
						__magic_name__:      int
						__magic_name__:      int
						__magic_name__:      int
						__magic_name__:      float			   =							0.0
						__magic_name__:      bool			   =							False
						__magic_name__:      jnp.dtype			   =							jnp.floataa
						__magic_name__:      bool			   =							False






						def UpperCAmelCase_					(   self   :   Dict						)				->     Any:







							"""simple docstring"""


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






						def __call__(   self   :   Optional[Any]		, _A   :   Tuple		, _A   :   List[Any]		, _A   :   Any=True						)				->     Dict:







							"""simple docstring"""


							snake_case_		:      Union[str, Any]           =						hidden_states
							if self.only_cross_attention:
								snake_case_		:      Optional[int]           =						self.attna(self.norma(_A						)		, _A		, deterministic=_A						)
							else:
								snake_case_		:      Dict           =						self.attna(self.norma(_A						)		, deterministic=_A						)
							snake_case_		:      Union[str, Any]           =						hidden_states + residual

							# cross attention
							snake_case_		:      Optional[int]           =						hidden_states
							snake_case_		:      Any           =						self.attna(self.norma(_A						)		, _A		, deterministic=_A						)
							snake_case_		:      Tuple           =						hidden_states + residual

							# feed forward
							snake_case_		:      Optional[Any]           =						hidden_states
							snake_case_		:      Optional[int]           =						self.ff(self.norma(_A						)		, deterministic=_A						)
							snake_case_		:      List[str]           =						hidden_states + residual

							return self.dropout_layer(_A		, deterministic=_A						)
class SCREAMING_SNAKE_CASE_    (       nn.Module       ):
						__magic_name__:      int
						__magic_name__:      int
						__magic_name__:      int
						__magic_name__:      int			   =							1
						__magic_name__:      float			   =							0.0
						__magic_name__:      bool			   =							False
						__magic_name__:      bool			   =							False
						__magic_name__:      jnp.dtype			   =							jnp.floataa
						__magic_name__:      bool			   =							False






						def UpperCAmelCase_					(   self   :   Union[str, Any]						)				->     int:







							"""simple docstring"""


							snake_case_		:      Union[str, Any]           =						nn.GroupNorm(num_groups=32		, epsilon=1E-5						)

							snake_case_		:      int           =						self.n_heads * self.d_head
							if self.use_linear_projection:
								snake_case_		:      Optional[int]           =						nn.Dense(_A		, dtype=self.dtype						)
							else:
								snake_case_		:      Union[str, Any]           =						nn.Conv(
								    _A		, kernel_size=(1, 1)		, strides=(1, 1)		, padding='VALID'		, dtype=self.dtype		, )

							snake_case_		:      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:
								snake_case_		:      List[str]           =						nn.Dense(_A		, dtype=self.dtype						)
							else:
								snake_case_		:      Dict           =						nn.Conv(
								    _A		, kernel_size=(1, 1)		, strides=(1, 1)		, padding='VALID'		, dtype=self.dtype		, )

							snake_case_		:      List[str]           =						nn.Dropout(rate=self.dropout						)






						def __call__(   self   :   Tuple		, _A   :   Any		, _A   :   List[str]		, _A   :   str=True						)				->     Optional[int]:







							"""simple docstring"""


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

							for transformer_block in self.transformer_blocks:
								snake_case_		:      List[str]           =						transformer_block(_A		, _A		, deterministic=_A						)

							if self.use_linear_projection:
								snake_case_		:      Tuple           =						self.proj_out(_A						)
								snake_case_		:      Optional[int]           =						hidden_states.reshape(_A		, _A		, _A		, _A						)
							else:
								snake_case_		:      str           =						hidden_states.reshape(_A		, _A		, _A		, _A						)
								snake_case_		:      Optional[Any]           =						self.proj_out(_A						)

							snake_case_		:      Optional[Any]           =						hidden_states + residual
							return self.dropout_layer(_A		, deterministic=_A						)
class SCREAMING_SNAKE_CASE_    (       nn.Module       ):
						__magic_name__:      int
						__magic_name__:      float			   =							0.0
						__magic_name__:      jnp.dtype			   =							jnp.floataa






						def UpperCAmelCase_					(   self   :   List[str]						)				->     Dict:







							"""simple docstring"""


							snake_case_		:      Optional[int]           =						FlaxGEGLU(self.dim		, self.dropout		, self.dtype						)
							snake_case_		:      Union[str, Any]           =						nn.Dense(self.dim		, dtype=self.dtype						)






						def __call__(   self   :   str		, _A   :   int		, _A   :   Dict=True						)				->     Tuple:







							"""simple docstring"""


							snake_case_		:      Optional[int]           =						self.net_a(_A		, deterministic=_A						)
							snake_case_		:      Optional[int]           =						self.net_a(_A						)
							return hidden_states
class SCREAMING_SNAKE_CASE_    (       nn.Module       ):
						__magic_name__:      int
						__magic_name__:      float			   =							0.0
						__magic_name__:      jnp.dtype			   =							jnp.floataa






						def UpperCAmelCase_					(   self   :   Optional[int]						)				->     Tuple:







							"""simple docstring"""


							snake_case_		:      Tuple           =						self.dim * 4
							snake_case_		:      int           =						nn.Dense(inner_dim * 2		, dtype=self.dtype						)
							snake_case_		:      Any           =						nn.Dropout(rate=self.dropout						)






						def __call__(   self   :   Union[str, Any]		, _A   :   List[str]		, _A   :   str=True						)				->     Dict:







							"""simple docstring"""


							snake_case_		:      List[str]           =						self.proj(_A						)
							snake_case_    ,snake_case_		:      int           =						jnp.split(_A		, 2		, axis=2						)
							return self.dropout_layer(hidden_linear * nn.gelu(_A						)		, deterministic=_A						)





'''simple docstring'''



import unittest
import warnings
from dataclasses import dataclass

from transformers.convert_slow_tokenizer import SpmConverter
from transformers.testing_utils import get_tests_dir



@dataclass
class       a							:


     """simple docstring"""






     SCREAMING_SNAKE_CASE			:					str



class       a							(						unittest.TestCase			):


     """simple docstring"""






     def 					lowerCamelCase__					(							self	:						str     )		->				Dict:
          __UpperCAmelCase      :					str    				=     get_tests_dir('''fixtures/test_sentencepiece.model'''     )
          __UpperCAmelCase      :					List[Any]    				=     get_tests_dir('''fixtures/test_sentencepiece_with_bytefallback.model'''     )

          __UpperCAmelCase      :					Tuple    				=     FakeOriginalTokenizer(vocab_file=snake_case     )

          with warnings.catch_warnings(record=snake_case     ) as w:
               __UpperCAmelCase      :					Optional[int]    				=     SpmConverter(snake_case     )
          self.assertEqual(len(snake_case     )				,					0     )

          __UpperCAmelCase      :					str    				=     FakeOriginalTokenizer(vocab_file=snake_case     )

          with warnings.catch_warnings(record=snake_case     ) as w:
               __UpperCAmelCase      :					Optional[int]    				=     SpmConverter(snake_case     )
          self.assertEqual(len(snake_case     )				,					1     )

          self.assertIn(
              '''The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option'''
              ''' which is not implemented in the fast tokenizers.'''				,					str(w[0].message     )				,					)