
'''simple docstring'''





import functools
import math

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

def    __snake_case       (  __magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__ = 4096	):







				'''simple docstring'''


				lowercase						    =			key.shape[-3:]
				lowercase						    =			value.shape[-1]
				lowercase						    =			min(__magic_name__		,						__magic_name__	)
				lowercase						    =			query / jnp.sqrt(__magic_name__	)

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

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

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

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

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

				lowercase						    =			jax.lax.map(f=__magic_name__		,						xs=jnp.arange(0		,						__magic_name__		,						__magic_name__	)	)

				lowercase						    =			jnp.max(__magic_name__		,						axis=0		,						keepdims=__magic_name__	)
				lowercase						    =			jnp.exp(chunk_max - global_max	)

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

				lowercase						    =			chunk_values.sum(axis=0	)
				lowercase						    =			jnp.expand_dims(__magic_name__		,						-1	).sum(axis=0	)

				return all_values / all_weights

def    __snake_case       (  __magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__=jax.lax.Precision.HIGHEST		,						__magic_name__ = 1024		,						__magic_name__ = 4096	):







				'''simple docstring'''


				lowercase						    =			query.shape[-3:]

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

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

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

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





class      UpperCamelCase_			(			nn.Module		):







						'''simple docstring'''



						UpperCamelCase			:     int							      =				42
						UpperCamelCase			:     Dict							      =				8
						UpperCamelCase			:     List[Any]							      =				64
						UpperCamelCase			:     str							      =				0.0
						UpperCamelCase			:     List[str]							      =				False
						UpperCamelCase			:     Optional[Any]							      =				jnp.floataa





						def 					SCREAMING_SNAKE_CASE( self       :Optional[Any]    )     ->int:
										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(lowerCamelCase__		,						use_bias=lowerCamelCase__		,						dtype=self.dtype		,						name="to_q"    )
										lowercase						    =			nn.Dense(lowerCamelCase__		,						use_bias=lowerCamelCase__		,						dtype=self.dtype		,						name="to_k"    )
										lowercase						    =			nn.Dense(lowerCamelCase__		,						use_bias=lowerCamelCase__		,						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 					SCREAMING_SNAKE_CASE( self       :int		,						lowerCAmelCase__       :Optional[int]    )     ->Dict:
										lowercase						    =			tensor.shape
										lowercase						    =			self.heads
										lowercase						    =			tensor.reshape(lowerCamelCase__		,						lowerCamelCase__		,						lowerCamelCase__		,						dim // head_size    )
										lowercase						    =			jnp.transpose(lowerCamelCase__		,						(0, 2, 1, 3)    )
										lowercase						    =			tensor.reshape(batch_size * head_size		,						lowerCamelCase__		,						dim // head_size    )
										return tensor





						def 					SCREAMING_SNAKE_CASE( self       :Optional[int]		,						lowerCAmelCase__       :Union[str, Any]    )     ->Any:
										lowercase						    =			tensor.shape
										lowercase						    =			self.heads
										lowercase						    =			tensor.reshape(batch_size // head_size		,						lowerCamelCase__		,						lowerCamelCase__		,						lowerCamelCase__    )
										lowercase						    =			jnp.transpose(lowerCamelCase__		,						(0, 2, 1, 3)    )
										lowercase						    =			tensor.reshape(batch_size // head_size		,						lowerCamelCase__		,						dim * head_size    )
										return tensor





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

										lowercase						    =			self.query(lowerCamelCase__    )
										lowercase						    =			self.key(lowerCamelCase__    )
										lowercase						    =			self.value(lowerCamelCase__    )

										lowercase						    =			self.reshape_heads_to_batch_dim(lowerCamelCase__    )
										lowercase						    =			self.reshape_heads_to_batch_dim(lowerCamelCase__    )
										lowercase						    =			self.reshape_heads_to_batch_dim(lowerCamelCase__    )

										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 % 64 == 0:
																		lowercase						    =			int(flatten_latent_dim / 64    )
														elif flatten_latent_dim % 16 == 0:
																		lowercase						    =			int(flatten_latent_dim / 16    )
														elif flatten_latent_dim % 4 == 0:
																		lowercase						    =			int(flatten_latent_dim / 4    )
														else:
																		lowercase						    =			int(lowerCamelCase__    )

														lowercase						    =			jax_memory_efficient_attention(
														    lowerCamelCase__		,						lowerCamelCase__		,						lowerCamelCase__		,						query_chunk_size=lowerCamelCase__		,						key_chunk_size=4096 * 4    )

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

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

										lowercase						    =			self.reshape_batch_dim_to_heads(lowerCamelCase__    )
										lowercase						    =			self.proj_attn(lowerCamelCase__    )
										return self.dropout_layer(lowerCamelCase__		,						deterministic=lowerCamelCase__    )





class      UpperCamelCase_			(			nn.Module		):







						'''simple docstring'''



						UpperCamelCase			:     Optional[Any]							      =				42
						UpperCamelCase			:     Optional[Any]							      =				42
						UpperCamelCase			:     List[Any]							      =				42
						UpperCamelCase			:     List[str]							      =				0.0
						UpperCamelCase			:     List[str]							      =				False
						UpperCamelCase			:     Optional[Any]							      =				jnp.floataa
						UpperCamelCase			:     Union[str, Any]							      =				False





						def 					SCREAMING_SNAKE_CASE( self       :List[str]    )     ->Dict:
										# 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[Any]		,						lowerCAmelCase__       :Dict		,						lowerCAmelCase__       :Union[str, Any]		,						lowerCAmelCase__       :Tuple=True    )     ->Any:
										# self attention
										lowercase						    =			hidden_states
										if self.only_cross_attention:
														lowercase						    =			self.attna(self.norma(lowerCamelCase__    )		,						lowerCamelCase__		,						deterministic=lowerCamelCase__    )
										else:
														lowercase						    =			self.attna(self.norma(lowerCamelCase__    )		,						deterministic=lowerCamelCase__    )
										lowercase						    =			hidden_states + residual

										# cross attention
										lowercase						    =			hidden_states
										lowercase						    =			self.attna(self.norma(lowerCamelCase__    )		,						lowerCamelCase__		,						deterministic=lowerCamelCase__    )
										lowercase						    =			hidden_states + residual

										# feed forward
										lowercase						    =			hidden_states
										lowercase						    =			self.ff(self.norma(lowerCamelCase__    )		,						deterministic=lowerCamelCase__    )
										lowercase						    =			hidden_states + residual

										return self.dropout_layer(lowerCamelCase__		,						deterministic=lowerCamelCase__    )





class      UpperCamelCase_			(			nn.Module		):







						'''simple docstring'''



						UpperCamelCase			:     Any							      =				42
						UpperCamelCase			:     Tuple							      =				42
						UpperCamelCase			:     List[Any]							      =				42
						UpperCamelCase			:     Union[str, Any]							      =				1
						UpperCamelCase			:     List[Any]							      =				0.0
						UpperCamelCase			:     List[str]							      =				False
						UpperCamelCase			:     Optional[Any]							      =				False
						UpperCamelCase			:     Optional[Any]							      =				jnp.floataa
						UpperCamelCase			:     Optional[int]							      =				False





						def 					SCREAMING_SNAKE_CASE( self       :Tuple    )     ->Union[str, Any]:
										lowercase						    =			nn.GroupNorm(num_groups=32		,						epsilon=1E-5    )

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

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

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





						def __call__( self       :List[Any]		,						lowerCAmelCase__       :Union[str, Any]		,						lowerCAmelCase__       :List[str]		,						lowerCAmelCase__       :Any=True    )     ->Optional[Any]:
										lowercase						    =			hidden_states.shape
										lowercase						    =			hidden_states
										lowercase						    =			self.norm(lowerCamelCase__    )
										if self.use_linear_projection:
														lowercase						    =			hidden_states.reshape(lowerCamelCase__		,						height * width		,						lowerCamelCase__    )
														lowercase						    =			self.proj_in(lowerCamelCase__    )
										else:
														lowercase						    =			self.proj_in(lowerCamelCase__    )
														lowercase						    =			hidden_states.reshape(lowerCamelCase__		,						height * width		,						lowerCamelCase__    )

										for transformer_block in self.transformer_blocks:
														lowercase						    =			transformer_block(lowerCamelCase__		,						lowerCamelCase__		,						deterministic=lowerCamelCase__    )

										if self.use_linear_projection:
														lowercase						    =			self.proj_out(lowerCamelCase__    )
														lowercase						    =			hidden_states.reshape(lowerCamelCase__		,						lowerCamelCase__		,						lowerCamelCase__		,						lowerCamelCase__    )
										else:
														lowercase						    =			hidden_states.reshape(lowerCamelCase__		,						lowerCamelCase__		,						lowerCamelCase__		,						lowerCamelCase__    )
														lowercase						    =			self.proj_out(lowerCamelCase__    )

										lowercase						    =			hidden_states + residual
										return self.dropout_layer(lowerCamelCase__		,						deterministic=lowerCamelCase__    )





class      UpperCamelCase_			(			nn.Module		):







						'''simple docstring'''



						UpperCamelCase			:     Any							      =				42
						UpperCamelCase			:     Union[str, Any]							      =				0.0
						UpperCamelCase			:     Tuple							      =				jnp.floataa





						def 					SCREAMING_SNAKE_CASE( self       :int    )     ->Tuple:
										# 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       :List[Any]		,						lowerCAmelCase__       :Union[str, Any]		,						lowerCAmelCase__       :Any=True    )     ->Optional[int]:
										lowercase						    =			self.net_a(lowerCamelCase__		,						deterministic=lowerCamelCase__    )
										lowercase						    =			self.net_a(lowerCamelCase__    )
										return hidden_states





class      UpperCamelCase_			(			nn.Module		):







						'''simple docstring'''



						UpperCamelCase			:     Union[str, Any]							      =				42
						UpperCamelCase			:     str							      =				0.0
						UpperCamelCase			:     int							      =				jnp.floataa





						def 					SCREAMING_SNAKE_CASE( self       :int    )     ->Optional[Any]:
										lowercase						    =			self.dim * 4
										lowercase						    =			nn.Dense(inner_dim * 2		,						dtype=self.dtype    )
										lowercase						    =			nn.Dropout(rate=self.dropout    )





						def __call__( self       :int		,						lowerCAmelCase__       :Union[str, Any]		,						lowerCAmelCase__       :Optional[Any]=True    )     ->int:
										lowercase						    =			self.proj(lowerCamelCase__    )
										lowercase						    =			jnp.split(lowerCamelCase__		,						2		,						axis=2    )
										return self.dropout_layer(hidden_linear * nn.gelu(lowerCamelCase__    )		,						deterministic=lowerCamelCase__    )


from transformers import PretrainedConfig





class      UpperCamelCase_			(			__a		):







						'''simple docstring'''



						UpperCamelCase			:     List[str]							      =				'''custom'''





						def __init__( self       :Tuple		,						lowerCAmelCase__       :int=1		,						**lowerCAmelCase__       :List[Any]    )     ->int:
										lowercase						    =			attribute
										super().__init__(**lowerCAmelCase__    )





class      UpperCamelCase_			(			__a		):







						'''simple docstring'''



						UpperCamelCase			:     Union[str, Any]							      =				'''custom'''





						def __init__( self       :Union[str, Any]		,						lowerCAmelCase__       :Dict=1		,						**lowerCAmelCase__       :Dict    )     ->List[str]:
										lowercase						    =			attribute
