
import functools
import math

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




def     a						(							snake_case__: Optional[Any]						,       snake_case__: Optional[int]						,       snake_case__: Dict						,       snake_case__: Tuple						,       snake_case__: int = 4_096 ):


				'''simple docstring'''
				lowercase_  ,     lowercase_  ,     lowercase_					  = key.shape[-3:]
				lowercase_					  = value.shape[-1]
				lowercase_					  = min(snake_case__						,       snake_case__ )
				lowercase_					  = query / jnp.sqrt(snake_case__ )

				@functools.partial(jax.checkpoint						,       prevent_cse=snake_case__ )
				def summarize_chunk(snake_case__: str						,       snake_case__: Optional[int]						,       snake_case__: Union[str, Any] ):
								lowercase_					  = jnp.einsum('''...qhd,...khd->...qhk'''						,       snake_case__						,       snake_case__						,       precision=snake_case__ )

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

								lowercase_					  = jnp.einsum('''...vhf,...qhv->...qhf'''						,       snake_case__						,       snake_case__						,       precision=snake_case__ )
								lowercase_					  = jnp.einsum('''...qhk->...qh'''						,       snake_case__ )

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

				def chunk_scanner(snake_case__: Dict ):
								# julienne key array
								lowercase_					  = 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
								lowercase_					  = 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__ )

				lowercase_  ,     lowercase_  ,     lowercase_					  = jax.lax.map(f=snake_case__						,       xs=jnp.arange(0						,       snake_case__						,       snake_case__ ) )

				lowercase_					  = jnp.max(snake_case__						,       axis=0						,       keepdims=snake_case__ )
				lowercase_					  = jnp.exp(chunk_max - global_max )

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

				lowercase_					  = chunk_values.sum(axis=0 )
				lowercase_					  = jnp.expand_dims(snake_case__						,       -1 ).sum(axis=0 )

				return all_values / all_weights




def     a						(							snake_case__: List[Any]						,       snake_case__: Tuple						,       snake_case__: List[Any]						,       snake_case__: Union[str, Any]=jax.lax.Precision.HIGHEST						,       snake_case__: int = 1_024						,       snake_case__: int = 4_096 ):


				'''simple docstring'''
				lowercase_  ,     lowercase_  ,     lowercase_					  = query.shape[-3:]

				def chunk_scanner(snake_case__: int						,       snake_case__: Optional[Any] ):
								# julienne query array
								lowercase_					  = 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__ ),
								)

				lowercase_  ,     lowercase_					  = 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 			lowercase__(    nn.Module							):



	"""simple docstring"""




	a      :int
	a      :int   					=     8
	a      :int   					=     64
	a      :float   					=     0.0
	a      :bool   					=     False
	a      :jnp.dtype   					=     jnp.floataa
	def 						_lowercase  (	self  :      Any	)	->					Union[str, Any]:
					lowercase_					  = self.dim_head * self.heads
					lowercase_					  = self.dim_head**-0.5

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

					lowercase_					  = nn.Dense(self.query_dim	,							dtype=self.dtype	,							name='''to_out_0'''	)
					lowercase_					  = nn.Dropout(rate=self.dropout	)
	def 						_lowercase  (	self  :      Optional[int]	,							SCREAMING_SNAKE_CASE_  :      str	)	->					Optional[Any]:
					lowercase_  ,     lowercase_  ,     lowercase_					  = tensor.shape
					lowercase_					  = self.heads
					lowercase_					  = tensor.reshape(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							dim // head_size	)
					lowercase_					  = jnp.transpose(SCREAMING_SNAKE_CASE_	,							(0, 2, 1, 3)	)
					lowercase_					  = tensor.reshape(batch_size * head_size	,							SCREAMING_SNAKE_CASE_	,							dim // head_size	)
					return tensor
	def 						_lowercase  (	self  :      Optional[Any]	,							SCREAMING_SNAKE_CASE_  :      Any	)	->					Optional[Any]:
					lowercase_  ,     lowercase_  ,     lowercase_					  = tensor.shape
					lowercase_					  = self.heads
					lowercase_					  = tensor.reshape(batch_size // head_size	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
					lowercase_					  = jnp.transpose(SCREAMING_SNAKE_CASE_	,							(0, 2, 1, 3)	)
					lowercase_					  = tensor.reshape(batch_size // head_size	,							SCREAMING_SNAKE_CASE_	,							dim * head_size	)
					return tensor




	def __call__(	self  :      Any	,							SCREAMING_SNAKE_CASE_  :      Union[str, Any]	,							SCREAMING_SNAKE_CASE_  :      str=None	,							SCREAMING_SNAKE_CASE_  :      Tuple=True	)	->					Optional[int]:
					lowercase_					  = hidden_states if context is None else context

					lowercase_					  = self.query(SCREAMING_SNAKE_CASE_	)
					lowercase_					  = self.key(SCREAMING_SNAKE_CASE_	)
					lowercase_					  = self.value(SCREAMING_SNAKE_CASE_	)

					lowercase_					  = self.reshape_heads_to_batch_dim(SCREAMING_SNAKE_CASE_	)
					lowercase_					  = self.reshape_heads_to_batch_dim(SCREAMING_SNAKE_CASE_	)
					lowercase_					  = self.reshape_heads_to_batch_dim(SCREAMING_SNAKE_CASE_	)

					if self.use_memory_efficient_attention:
									lowercase_					  = query_states.transpose(1	,							0	,							2	)
									lowercase_					  = key_states.transpose(1	,							0	,							2	)
									lowercase_					  = 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_					  = query_states.shape[-3]
									if flatten_latent_dim % 6_4 == 0:
													lowercase_					  = int(flatten_latent_dim / 6_4	)
									elif flatten_latent_dim % 1_6 == 0:
													lowercase_					  = int(flatten_latent_dim / 1_6	)
									elif flatten_latent_dim % 4 == 0:
													lowercase_					  = int(flatten_latent_dim / 4	)
									else:
													lowercase_					  = int(SCREAMING_SNAKE_CASE_	)

									lowercase_					  = jax_memory_efficient_attention(
									    SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							query_chunk_size=SCREAMING_SNAKE_CASE_	,							key_chunk_size=4_0_9_6 * 4	)

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

									# attend to values
									lowercase_					  = jnp.einsum('''b i j, b j d -> b i d'''	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)

					lowercase_					  = self.reshape_batch_dim_to_heads(SCREAMING_SNAKE_CASE_	)
					lowercase_					  = self.proj_attn(SCREAMING_SNAKE_CASE_	)
					return self.dropout_layer(SCREAMING_SNAKE_CASE_	,							deterministic=SCREAMING_SNAKE_CASE_	)

class 			lowercase__(    nn.Module							):



	"""simple docstring"""




	a      :int
	a      :int
	a      :int
	a      :float   					=     0.0
	a      :bool   					=     False
	a      :jnp.dtype   					=     jnp.floataa
	a      :bool   					=     False
	def 						_lowercase  (	self  :      str	)	->					Optional[int]:
					# self attention (or cross_attention if only_cross_attention is True)
					lowercase_					  = FlaxAttention(
					    self.dim	,							self.n_heads	,							self.d_head	,							self.dropout	,							self.use_memory_efficient_attention	,							dtype=self.dtype	)
					# cross attention
					lowercase_					  = FlaxAttention(
					    self.dim	,							self.n_heads	,							self.d_head	,							self.dropout	,							self.use_memory_efficient_attention	,							dtype=self.dtype	)
					lowercase_					  = FlaxFeedForward(dim=self.dim	,							dropout=self.dropout	,							dtype=self.dtype	)
					lowercase_					  = nn.LayerNorm(epsilon=1e-5	,							dtype=self.dtype	)
					lowercase_					  = nn.LayerNorm(epsilon=1e-5	,							dtype=self.dtype	)
					lowercase_					  = nn.LayerNorm(epsilon=1e-5	,							dtype=self.dtype	)
					lowercase_					  = nn.Dropout(rate=self.dropout	)




	def __call__(	self  :      List[str]	,							SCREAMING_SNAKE_CASE_  :      Dict	,							SCREAMING_SNAKE_CASE_  :      int	,							SCREAMING_SNAKE_CASE_  :      Dict=True	)	->					List[str]:
					# self attention
					lowercase_					  = hidden_states
					if self.only_cross_attention:
									lowercase_					  = self.attna(self.norma(SCREAMING_SNAKE_CASE_	)	,							SCREAMING_SNAKE_CASE_	,							deterministic=SCREAMING_SNAKE_CASE_	)
					else:
									lowercase_					  = self.attna(self.norma(SCREAMING_SNAKE_CASE_	)	,							deterministic=SCREAMING_SNAKE_CASE_	)
					lowercase_					  = hidden_states + residual

					# cross attention
					lowercase_					  = hidden_states
					lowercase_					  = self.attna(self.norma(SCREAMING_SNAKE_CASE_	)	,							SCREAMING_SNAKE_CASE_	,							deterministic=SCREAMING_SNAKE_CASE_	)
					lowercase_					  = hidden_states + residual

					# feed forward
					lowercase_					  = hidden_states
					lowercase_					  = self.ff(self.norma(SCREAMING_SNAKE_CASE_	)	,							deterministic=SCREAMING_SNAKE_CASE_	)
					lowercase_					  = hidden_states + residual

					return self.dropout_layer(SCREAMING_SNAKE_CASE_	,							deterministic=SCREAMING_SNAKE_CASE_	)

class 			lowercase__(    nn.Module							):



	"""simple docstring"""




	a      :int
	a      :int
	a      :int
	a      :int   					=     1
	a      :float   					=     0.0
	a      :bool   					=     False
	a      :bool   					=     False
	a      :jnp.dtype   					=     jnp.floataa
	a      :bool   					=     False
	def 						_lowercase  (	self  :      Any	)	->					List[str]:
					lowercase_					  = nn.GroupNorm(num_groups=3_2	,							epsilon=1e-5	)

					lowercase_					  = self.n_heads * self.d_head
					if self.use_linear_projection:
									lowercase_					  = nn.Dense(SCREAMING_SNAKE_CASE_	,							dtype=self.dtype	)
					else:
									lowercase_					  = nn.Conv(
									    SCREAMING_SNAKE_CASE_	,							kernel_size=(1, 1)	,							strides=(1, 1)	,							padding='''VALID'''	,							dtype=self.dtype	,							)

					lowercase_					  = [
					    FlaxBasicTransformerBlock(
					        SCREAMING_SNAKE_CASE_	,							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_					  = nn.Dense(SCREAMING_SNAKE_CASE_	,							dtype=self.dtype	)
					else:
									lowercase_					  = nn.Conv(
									    SCREAMING_SNAKE_CASE_	,							kernel_size=(1, 1)	,							strides=(1, 1)	,							padding='''VALID'''	,							dtype=self.dtype	,							)

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




	def __call__(	self  :      Any	,							SCREAMING_SNAKE_CASE_  :      Dict	,							SCREAMING_SNAKE_CASE_  :      Optional[Any]	,							SCREAMING_SNAKE_CASE_  :      Optional[int]=True	)	->					Union[str, Any]:
					lowercase_  ,     lowercase_  ,     lowercase_  ,     lowercase_					  = hidden_states.shape
					lowercase_					  = hidden_states
					lowercase_					  = self.norm(SCREAMING_SNAKE_CASE_	)
					if self.use_linear_projection:
									lowercase_					  = hidden_states.reshape(SCREAMING_SNAKE_CASE_	,							height * width	,							SCREAMING_SNAKE_CASE_	)
									lowercase_					  = self.proj_in(SCREAMING_SNAKE_CASE_	)
					else:
									lowercase_					  = self.proj_in(SCREAMING_SNAKE_CASE_	)
									lowercase_					  = hidden_states.reshape(SCREAMING_SNAKE_CASE_	,							height * width	,							SCREAMING_SNAKE_CASE_	)

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

					if self.use_linear_projection:
									lowercase_					  = self.proj_out(SCREAMING_SNAKE_CASE_	)
									lowercase_					  = hidden_states.reshape(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
					else:
									lowercase_					  = hidden_states.reshape(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
									lowercase_					  = self.proj_out(SCREAMING_SNAKE_CASE_	)

					lowercase_					  = hidden_states + residual
					return self.dropout_layer(SCREAMING_SNAKE_CASE_	,							deterministic=SCREAMING_SNAKE_CASE_	)

class 			lowercase__(    nn.Module							):



	"""simple docstring"""




	a      :int
	a      :float   					=     0.0
	a      :jnp.dtype   					=     jnp.floataa
	def 						_lowercase  (	self  :      Optional[Any]	)	->					Union[str, Any]:
					# The second linear layer needs to be called
					# net_2 for now to match the index of the Sequential layer
					lowercase_					  = FlaxGEGLU(self.dim	,							self.dropout	,							self.dtype	)
					lowercase_					  = nn.Dense(self.dim	,							dtype=self.dtype	)




	def __call__(	self  :      Tuple	,							SCREAMING_SNAKE_CASE_  :      Any	,							SCREAMING_SNAKE_CASE_  :      int=True	)	->					Optional[Any]:
					lowercase_					  = self.net_a(SCREAMING_SNAKE_CASE_	,							deterministic=SCREAMING_SNAKE_CASE_	)
					lowercase_					  = self.net_a(SCREAMING_SNAKE_CASE_	)
					return hidden_states

class 			lowercase__(    nn.Module							):



	"""simple docstring"""




	a      :int
	a      :float   					=     0.0
	a      :jnp.dtype   					=     jnp.floataa
	def 						_lowercase  (	self  :      Tuple	)	->					Optional[Any]:
					lowercase_					  = self.dim * 4
					lowercase_					  = nn.Dense(inner_dim * 2	,							dtype=self.dtype	)
					lowercase_					  = nn.Dropout(rate=self.dropout	)




	def __call__(	self  :      Optional[int]	,							SCREAMING_SNAKE_CASE_  :      List[Any]	,							SCREAMING_SNAKE_CASE_  :      List[str]=True	)	->					Tuple:
					lowercase_					  = self.proj(SCREAMING_SNAKE_CASE_	)
					lowercase_  ,     lowercase_					  = jnp.split(SCREAMING_SNAKE_CASE_	,							2	,							axis=2	)
					return self.dropout_layer(hidden_linear * nn.gelu(SCREAMING_SNAKE_CASE_	)	,							deterministic=SCREAMING_SNAKE_CASE_	)






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__a        = {
    'configuration_nystromformer': ['NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'NystromformerConfig'],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a        = [
						    'NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'NystromformerForMaskedLM',
						    'NystromformerForMultipleChoice',
						    'NystromformerForQuestionAnswering',
						    'NystromformerForSequenceClassification',
						    'NystromformerForTokenClassification',
						    'NystromformerLayer',
						    'NystromformerModel',
						    'NystromformerPreTrainedModel',
						]


if TYPE_CHECKING:
						from .configuration_nystromformer import NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, NystromformerConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_nystromformer import (
												    NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    NystromformerForMaskedLM,
												    NystromformerForMultipleChoice,
												    NystromformerForQuestionAnswering,
												    NystromformerForSequenceClassification,
												    NystromformerForTokenClassification,
												    NystromformerLayer,
												    NystromformerModel,
												    NystromformerPreTrainedModel,
												)


else:
						import sys

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




