

"""simple docstring"""






import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp
def 		A_	(   __UpperCamelCase     :				str       ,	__UpperCamelCase     :				List[Any]       ,	__UpperCamelCase     :				str       ,	__UpperCamelCase     :				Optional[Any]       ,	__UpperCamelCase     :				int = 40_96				):
		lowercase					,  lowercase					,  lowercase										=							key.shape[-3:]
		lowercase										=							value.shape[-1]
		lowercase										=							min(__UpperCamelCase       ,	__UpperCamelCase				)
		lowercase										=							query / jnp.sqrt(__UpperCamelCase				)

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

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

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

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

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

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

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

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

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

		return all_values / all_weights
def 		A_	(   __UpperCamelCase     :				Optional[Any]       ,	__UpperCamelCase     :				Dict       ,	__UpperCamelCase     :				Any       ,	__UpperCamelCase     :				Tuple=jax.lax.Precision.HIGHEST       ,	__UpperCamelCase     :				int = 10_24       ,	__UpperCamelCase     :				int = 40_96				):
		lowercase					,  lowercase					,  lowercase										=							query.shape[-3:]

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

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

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

		return jnp.concatenate(__UpperCamelCase       ,	axis=-3				)  # fuse the chunked result back
class 	_lowerCAmelCase    (      nn.Module       ):
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    int     =  8
							__lowerCAmelCase					:    int     =  64
							__lowerCAmelCase					:    float     =  0.0
							__lowerCAmelCase					:    bool     =  False
							__lowerCAmelCase					:    jnp.dtype     =  jnp.floataa

							def 					_lowerCAmelCase  ( self       :    str    )   ->    Optional[int]:

									"""simple docstring"""
									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(a ,		use_bias=a ,		dtype=self.dtype ,		name='''to_q'''    )
									lowercase										=							nn.Dense(a ,		use_bias=a ,		dtype=self.dtype ,		name='''to_k'''    )
									lowercase										=							nn.Dense(a ,		use_bias=a ,		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 					_lowerCAmelCase  ( self       :    int ,		a       :    str    )   ->    List[str]:

									"""simple docstring"""
									lowercase					,  lowercase					,  lowercase										=							tensor.shape
									lowercase										=							self.heads
									lowercase										=							tensor.reshape(a ,		a ,		a ,		dim // head_size    )
									lowercase										=							jnp.transpose(a ,		(0, 2, 1, 3)    )
									lowercase										=							tensor.reshape(batch_size * head_size ,		a ,		dim // head_size    )
									return tensor

							def 					_lowerCAmelCase  ( self       :    Union[str, Any] ,		a       :    Union[str, Any]    )   ->    Tuple:

									"""simple docstring"""
									lowercase					,  lowercase					,  lowercase										=							tensor.shape
									lowercase										=							self.heads
									lowercase										=							tensor.reshape(batch_size // head_size ,		a ,		a ,		a    )
									lowercase										=							jnp.transpose(a ,		(0, 2, 1, 3)    )
									lowercase										=							tensor.reshape(batch_size // head_size ,		a ,		dim * head_size    )
									return tensor

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

									"""simple docstring"""
									lowercase										=							hidden_states if context is None else context

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

									lowercase										=							self.reshape_heads_to_batch_dim(a    )
									lowercase										=							self.reshape_heads_to_batch_dim(a    )
									lowercase										=							self.reshape_heads_to_batch_dim(a    )

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

											lowercase										=							jax_memory_efficient_attention(
											    a ,		a ,		a ,		query_chunk_size=a ,		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''' ,		a ,		a    )
											lowercase										=							attention_scores * self.scale
											lowercase										=							nn.softmax(a ,		axis=2    )

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

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




class 	_lowerCAmelCase    (      nn.Module       ):
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    float     =  0.0
							__lowerCAmelCase					:    bool     =  False
							__lowerCAmelCase					:    jnp.dtype     =  jnp.floataa
							__lowerCAmelCase					:    bool     =  False

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

									"""simple docstring"""
									# 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       :    Optional[Any] ,		a       :    Dict ,		a       :    Optional[int] ,		a       :    Optional[int]=True    )   ->    str:

									"""simple docstring"""
									# self attention
									lowercase										=							hidden_states
									if self.only_cross_attention:
											lowercase										=							self.attna(self.norma(a    ) ,		a ,		deterministic=a    )
									else:
											lowercase										=							self.attna(self.norma(a    ) ,		deterministic=a    )
									lowercase										=							hidden_states + residual

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

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

									return self.dropout_layer(a ,		deterministic=a    )




class 	_lowerCAmelCase    (      nn.Module       ):
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    int     =  1
							__lowerCAmelCase					:    float     =  0.0
							__lowerCAmelCase					:    bool     =  False
							__lowerCAmelCase					:    bool     =  False
							__lowerCAmelCase					:    jnp.dtype     =  jnp.floataa
							__lowerCAmelCase					:    bool     =  False

							def 					_lowerCAmelCase  ( self       :    List[Any]    )   ->    List[Any]:

									"""simple docstring"""
									lowercase										=							nn.GroupNorm(num_groups=32 ,		epsilon=1E-5    )

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

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

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

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

									"""simple docstring"""
									lowercase					,  lowercase					,  lowercase					,  lowercase										=							hidden_states.shape
									lowercase										=							hidden_states
									lowercase										=							self.norm(a    )
									if self.use_linear_projection:
											lowercase										=							hidden_states.reshape(a ,		height * width ,		a    )
											lowercase										=							self.proj_in(a    )
									else:
											lowercase										=							self.proj_in(a    )
											lowercase										=							hidden_states.reshape(a ,		height * width ,		a    )

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

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

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




class 	_lowerCAmelCase    (      nn.Module       ):
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    float     =  0.0
							__lowerCAmelCase					:    jnp.dtype     =  jnp.floataa

							def 					_lowerCAmelCase  ( self       :    Dict    )   ->    Tuple:

									"""simple docstring"""
									# 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       :    Union[str, Any] ,		a       :    List[str] ,		a       :    List[str]=True    )   ->    int:

									"""simple docstring"""
									lowercase										=							self.net_a(a ,		deterministic=a    )
									lowercase										=							self.net_a(a    )
									return hidden_states




class 	_lowerCAmelCase    (      nn.Module       ):
							__lowerCAmelCase					:    int
							__lowerCAmelCase					:    float     =  0.0
							__lowerCAmelCase					:    jnp.dtype     =  jnp.floataa

							def 					_lowerCAmelCase  ( self       :    List[Any]    )   ->    Tuple:

									"""simple docstring"""
									lowercase										=							self.dim * 4
									lowercase										=							nn.Dense(inner_dim * 2 ,		dtype=self.dtype    )
									lowercase										=							nn.Dropout(rate=self.dropout    )

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

									"""simple docstring"""
									lowercase										=							self.proj(a    )
									lowercase					,  lowercase										=							jnp.split(a ,		2 ,		axis=2    )
									return self.dropout_layer(hidden_linear * nn.gelu(a    ) ,		deterministic=a    )


"""simple docstring"""






from ..utils import DummyObject, requires_backends
class 	_lowerCAmelCase    (      metaclass=__snake_case       ):
							__lowerCAmelCase					:    Optional[Any]     =  ['''tensorflow_text''']

							def __init__( self       :    Optional[int] ,		*a       :    int ,		**a       :    str    )   ->    str:

									"""simple docstring"""
									requires_backends(self ,		['''tensorflow_text''']    )