

"""simple docstring"""




import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp
def      __SCREAMING_SNAKE_CASE		(		lowercase__    ,							lowercase__    ,							lowercase__    ,							lowercase__    ,							lowercase__ = 4_096   ):
					"""simple docstring"""
					A    ,					A    ,					A    =      key.shape[-3:]
					A    =      value.shape[-1]
					A    =      min(lowercase__    ,							lowercase__   )
					A    =      query / jnp.sqrt(lowercase__   )

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

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

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

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

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

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

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

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

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

					return all_values / all_weights
def      __SCREAMING_SNAKE_CASE		(		lowercase__    ,							lowercase__    ,							lowercase__    ,							lowercase__=jax.lax.Precision.HIGHEST    ,							lowercase__ = 1_024    ,							lowercase__ = 4_096   ):
					"""simple docstring"""
					A    ,					A    ,					A    =      query.shape[-3:]

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

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

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

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





class     __UpperCamelCase				(     nn.Module		):
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					8
		SCREAMING_SNAKE_CASE              =					64
		SCREAMING_SNAKE_CASE              =					0.0
		SCREAMING_SNAKE_CASE              =					False
		SCREAMING_SNAKE_CASE              =					jnp.floataa


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


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


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

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

							A    =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE)
							A    =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE)
							A    =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE)

							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 % 6_4 == 0:
																	A    =      int(flatten_latent_dim / 6_4)
												elif flatten_latent_dim % 1_6 == 0:
																	A    =      int(flatten_latent_dim / 1_6)
												elif flatten_latent_dim % 4 == 0:
																	A    =      int(flatten_latent_dim / 4)
												else:
																	A    =      int(__SCREAMING_SNAKE_CASE)

												A    =      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)

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

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

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





class     __UpperCamelCase				(     nn.Module		):
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					0.0
		SCREAMING_SNAKE_CASE              =					False
		SCREAMING_SNAKE_CASE              =					jnp.floataa
		SCREAMING_SNAKE_CASE              =					False


		def       SCREAMING_SNAKE_CASE__     (self				:	Dict):
							# 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				:	Any  ,  __SCREAMING_SNAKE_CASE				:	Tuple  ,  __SCREAMING_SNAKE_CASE				:	Tuple  ,  __SCREAMING_SNAKE_CASE				:	Any=True):
							# self attention
							A    =      hidden_states
							if self.only_cross_attention:
												A    =      self.attna(self.norma(__SCREAMING_SNAKE_CASE)  ,  __SCREAMING_SNAKE_CASE  ,  deterministic=__SCREAMING_SNAKE_CASE)
							else:
												A    =      self.attna(self.norma(__SCREAMING_SNAKE_CASE)  ,  deterministic=__SCREAMING_SNAKE_CASE)
							A    =      hidden_states + residual

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

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

							return self.dropout_layer(__SCREAMING_SNAKE_CASE  ,  deterministic=__SCREAMING_SNAKE_CASE)





class     __UpperCamelCase				(     nn.Module		):
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					1
		SCREAMING_SNAKE_CASE              =					0.0
		SCREAMING_SNAKE_CASE              =					False
		SCREAMING_SNAKE_CASE              =					False
		SCREAMING_SNAKE_CASE              =					jnp.floataa
		SCREAMING_SNAKE_CASE              =					False


		def       SCREAMING_SNAKE_CASE__     (self				:	Any):
							A    =      nn.GroupNorm(num_groups=3_2  ,  epsilon=1E-5)

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

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

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


		def __call__(self				:	List[Any]  ,  __SCREAMING_SNAKE_CASE				:	int  ,  __SCREAMING_SNAKE_CASE				:	Optional[Any]  ,  __SCREAMING_SNAKE_CASE				:	Any=True):
							A    ,					A    ,					A    ,					A    =      hidden_states.shape
							A    =      hidden_states
							A    =      self.norm(__SCREAMING_SNAKE_CASE)
							if self.use_linear_projection:
												A    =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,  height * width  ,  __SCREAMING_SNAKE_CASE)
												A    =      self.proj_in(__SCREAMING_SNAKE_CASE)
							else:
												A    =      self.proj_in(__SCREAMING_SNAKE_CASE)
												A    =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,  height * width  ,  __SCREAMING_SNAKE_CASE)

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

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

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





class     __UpperCamelCase				(     nn.Module		):
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					0.0
		SCREAMING_SNAKE_CASE              =					jnp.floataa


		def       SCREAMING_SNAKE_CASE__     (self				:	Tuple):
							# 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				:	List[Any]  ,  __SCREAMING_SNAKE_CASE				:	str  ,  __SCREAMING_SNAKE_CASE				:	List[Any]=True):
							A    =      self.net_a(__SCREAMING_SNAKE_CASE  ,  deterministic=__SCREAMING_SNAKE_CASE)
							A    =      self.net_a(__SCREAMING_SNAKE_CASE)
							return hidden_states





class     __UpperCamelCase				(     nn.Module		):
		SCREAMING_SNAKE_CASE              =					42
		SCREAMING_SNAKE_CASE              =					0.0
		SCREAMING_SNAKE_CASE              =					jnp.floataa


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


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








"""simple docstring"""




def      __SCREAMING_SNAKE_CASE		(		lowercase__   ):
					"""simple docstring"""
					if not isinstance(lowercase__    ,							lowercase__   ):
										raise ValueError("double_factorial() only accepts integral values"   )
					if n < 0:
										raise ValueError("double_factorial() not defined for negative values"   )
					return 1 if n <= 1 else n * double_factorial(n - 2   )


if __name__ == "__main__":
	import doctest

	doctest.testmod()





