


'''simple docstring'''



import functools
import math

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






def 						lowerCamelCase							(_SCREAMING_SNAKE_CASE   :     Optional[Any]		,   _SCREAMING_SNAKE_CASE   :     Union[str, Any]		,   _SCREAMING_SNAKE_CASE   :     int		,   _SCREAMING_SNAKE_CASE   :     str		,   _SCREAMING_SNAKE_CASE   :     int = 4_096     ):
					__a    : str      					=						key.shape[-3:]
					__a    : str      					=						value.shape[-1]
					__a    : Dict      					=						min(_SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE     )
					__a    : Optional[int]      					=						query / jnp.sqrt(_SCREAMING_SNAKE_CASE     )

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

										__a    : Dict      					=						jnp.max(_SCREAMING_SNAKE_CASE		,   axis=-1		,   keepdims=_SCREAMING_SNAKE_CASE     )
										__a    : int      					=						jax.lax.stop_gradient(_SCREAMING_SNAKE_CASE     )
										__a    : str      					=						jnp.exp(attn_weights - max_score     )

										__a    : Dict      					=						jnp.einsum('...vhf,...qhv->...qhf'		,   _SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE		,   precision=_SCREAMING_SNAKE_CASE     )
										__a    : Tuple      					=						jnp.einsum('...qhk->...qh'		,   _SCREAMING_SNAKE_CASE     )

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

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

					__a    : List[Any]      					=						jax.lax.map(f=_SCREAMING_SNAKE_CASE		,   xs=jnp.arange(0		,   _SCREAMING_SNAKE_CASE		,   _SCREAMING_SNAKE_CASE     )     )

					__a    : Dict      					=						jnp.max(_SCREAMING_SNAKE_CASE		,   axis=0		,   keepdims=_SCREAMING_SNAKE_CASE     )
					__a    : str      					=						jnp.exp(chunk_max - global_max     )

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

					__a    : List[Any]      					=						chunk_values.sum(axis=0     )
					__a    : Tuple      					=						jnp.expand_dims(_SCREAMING_SNAKE_CASE		,   -1     ).sum(axis=0     )

					return all_values / all_weights






def 						lowerCamelCase							(_SCREAMING_SNAKE_CASE   :     int		,   _SCREAMING_SNAKE_CASE   :     str		,   _SCREAMING_SNAKE_CASE   :     Union[str, Any]		,   _SCREAMING_SNAKE_CASE   :     Optional[Any]=jax.lax.Precision.HIGHEST		,   _SCREAMING_SNAKE_CASE   :     int = 1_024		,   _SCREAMING_SNAKE_CASE   :     int = 4_096     ):
					__a    : Dict      					=						query.shape[-3:]

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

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

					__a    : Union[str, Any]      					=						jax.lax.scan(
					    f=_SCREAMING_SNAKE_CASE		,   init=0		,   xs=_SCREAMING_SNAKE_CASE		,   length=math.ceil(num_q / query_chunk_size     )     )

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





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




	def    __UpperCAmelCase							(     self							):

						'''simple docstring'''





						__a    : Dict      					=						self.dim_head * self.heads
						__a    : Any      					=						self.dim_head**-0.5

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

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




	def    __UpperCAmelCase							(     self		,       __a							):

						'''simple docstring'''





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




	def    __UpperCAmelCase							(     self		,       __a							):

						'''simple docstring'''





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





	def __call__(     self		,       __a		,       __a=None		,       __a=True							):

						'''simple docstring'''





						__a    : str      					=						hidden_states if context is None else context

						__a    : List[str]      					=						self.query(__a							)
						__a    : Union[str, Any]      					=						self.key(__a							)
						__a    : List[Any]      					=						self.value(__a							)

						__a    : Union[str, Any]      					=						self.reshape_heads_to_batch_dim(__a							)
						__a    : List[str]      					=						self.reshape_heads_to_batch_dim(__a							)
						__a    : Tuple      					=						self.reshape_heads_to_batch_dim(__a							)

						if self.use_memory_efficient_attention:
											__a    : Union[str, Any]      					=						query_states.transpose(1		,       0		,       2							)
											__a    : Dict      					=						key_states.transpose(1		,       0		,       2							)
											__a    : Union[str, Any]      					=						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    : Optional[Any]      					=						query_states.shape[-3]
											if flatten_latent_dim % 64 == 0:
																__a    : Optional[Any]      					=						int(flatten_latent_dim / 64							)
											elif flatten_latent_dim % 16 == 0:
																__a    : Optional[Any]      					=						int(flatten_latent_dim / 16							)
											elif flatten_latent_dim % 4 == 0:
																__a    : Tuple      					=						int(flatten_latent_dim / 4							)
											else:
																__a    : Any      					=						int(__a							)

											__a    : int      					=						jax_memory_efficient_attention(
											    __a		,       __a		,       __a		,       query_chunk_size=__a		,       key_chunk_size=4096 * 4							)

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

											# attend to values
											__a    : List[str]      					=						jnp.einsum('b i j, b j d -> b i d'		,       __a		,       __a							)

						__a    : List[str]      					=						self.reshape_batch_dim_to_heads(__a							)
						__a    : Union[str, Any]      					=						self.proj_attn(__a							)
						return self.dropout_layer(__a		,       deterministic=__a							)





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




	def    __UpperCAmelCase							(     self							):

						'''simple docstring'''





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





	def __call__(     self		,       __a		,       __a		,       __a=True							):

						'''simple docstring'''





						__a    : Optional[Any]      					=						hidden_states
						if self.only_cross_attention:
											__a    : List[str]      					=						self.attna(self.norma(__a							)		,       __a		,       deterministic=__a							)
						else:
											__a    : str      					=						self.attna(self.norma(__a							)		,       deterministic=__a							)
						__a    : Tuple      					=						hidden_states + residual

						# cross attention
						__a    : Dict      					=						hidden_states
						__a    : List[str]      					=						self.attna(self.norma(__a							)		,       __a		,       deterministic=__a							)
						__a    : Union[str, Any]      					=						hidden_states + residual

						# feed forward
						__a    : str      					=						hidden_states
						__a    : Any      					=						self.ff(self.norma(__a							)		,       deterministic=__a							)
						__a    : str      					=						hidden_states + residual

						return self.dropout_layer(__a		,       deterministic=__a							)





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




	def    __UpperCAmelCase							(     self							):

						'''simple docstring'''





						__a    : int      					=						nn.GroupNorm(num_groups=32		,       epsilon=1E-5							)

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

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

						__a    : Tuple      					=						nn.Dropout(rate=self.dropout							)





	def __call__(     self		,       __a		,       __a		,       __a=True							):

						'''simple docstring'''





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

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

						if self.use_linear_projection:
											__a    : int      					=						self.proj_out(__a							)
											__a    : Any      					=						hidden_states.reshape(__a		,       __a		,       __a		,       __a							)
						else:
											__a    : str      					=						hidden_states.reshape(__a		,       __a		,       __a		,       __a							)
											__a    : str      					=						self.proj_out(__a							)

						__a    : List[Any]      					=						hidden_states + residual
						return self.dropout_layer(__a		,       deterministic=__a							)





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




	def    __UpperCAmelCase							(     self							):

						'''simple docstring'''





						__a    : int      					=						FlaxGEGLU(self.dim		,       self.dropout		,       self.dtype							)
						__a    : int      					=						nn.Dense(self.dim		,       dtype=self.dtype							)





	def __call__(     self		,       __a		,       __a=True							):

						'''simple docstring'''





						__a    : Union[str, Any]      					=						self.net_a(__a		,       deterministic=__a							)
						__a    : List[Any]      					=						self.net_a(__a							)
						return hidden_states





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




	def    __UpperCAmelCase							(     self							):

						'''simple docstring'''





						__a    : int      					=						self.dim * 4
						__a    : Optional[Any]      					=						nn.Dense(inner_dim * 2		,       dtype=self.dtype							)
						__a    : Any      					=						nn.Dropout(rate=self.dropout							)





	def __call__(     self		,       __a		,       __a=True							):

						'''simple docstring'''





						__a    : Optional[Any]      					=						self.proj(__a							)
						__a    : List[Any]      					=						jnp.split(__a		,       2		,       axis=2							)
						return self.dropout_layer(hidden_linear * nn.gelu(__a							)		,       deterministic=__a							)







'''simple docstring'''



import itertools






def 						lowerCamelCase							(_SCREAMING_SNAKE_CASE   :     int = 1_000_000     ):
					__a    : str      					=						[1]

					for i in itertools.count(len(_SCREAMING_SNAKE_CASE     )     ):
										__a    : Dict      					=						0
										for j in itertools.count(1     ):
															__a    : int      					=						-1 if j % 2 == 0 else +1
															__a    : Optional[int]      					=						(j * j * 3 - j) // 2
															if index > i:
																				break
															item += partitions[i - index] * sign
															item %= number
															index += j
															if index > i:
																				break
															item += partitions[i - index] * sign
															item %= number

										if item == 0:
															return i
										partitions.append(_SCREAMING_SNAKE_CASE     )

					return 0


if __name__ == "__main__":
	import doctest

	doctest.testmod()

	print(f'''{solution() = }''')



