'''simple docstring'''




import functools
import math

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

def      UpperCamelCase_							( _UpperCAmelCase  :				List[str]      ,     _UpperCAmelCase  :				List[str]      ,     _UpperCAmelCase  :				int      ,     _UpperCAmelCase  :				List[Any]      ,     _UpperCAmelCase  :				Union[str, Any] = 4_096   )   ->   str:


							"""simple docstring"""







							_UpperCAmelCase					:  Dict											=   key.shape[-3:]
							_UpperCAmelCase					:  Union[str, Any]											=   value.shape[-1]
							_UpperCAmelCase					:  Tuple											=   min(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE   )
							_UpperCAmelCase					:  List[str]											=   query / jnp.sqrt(_SCREAMING_SNAKE_CASE   )

							@functools.partial(jax.checkpoint      ,     prevent_cse=_SCREAMING_SNAKE_CASE   )
							def summarize_chunk(_UpperCAmelCase  :				int      ,     _UpperCAmelCase  :				List[str]      ,     _UpperCAmelCase  :				int   ):
														_UpperCAmelCase					:  Optional[int]											=   jnp.einsum("...qhd,...khd->...qhk"      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     precision=_SCREAMING_SNAKE_CASE   )

														_UpperCAmelCase					:  int											=   jnp.max(_SCREAMING_SNAKE_CASE      ,     axis=-1      ,     keepdims=_SCREAMING_SNAKE_CASE   )
														_UpperCAmelCase					:  List[Any]											=   jax.lax.stop_gradient(_SCREAMING_SNAKE_CASE   )
														_UpperCAmelCase					:  int											=   jnp.exp(attn_weights - max_score   )

														_UpperCAmelCase					:  Optional[Any]											=   jnp.einsum("...vhf,...qhv->...qhf"      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     precision=_SCREAMING_SNAKE_CASE   )
														_UpperCAmelCase					:  int											=   jnp.einsum("...qhk->...qh"      ,     _SCREAMING_SNAKE_CASE   )

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

							def chunk_scanner(_UpperCAmelCase  :				Optional[int]   ):
														# julienne key array
														_UpperCAmelCase					:  Optional[int]											=   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
														_UpperCAmelCase					:  Union[str, 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   )

							_UpperCAmelCase					:  Tuple											=   jax.lax.map(f=_SCREAMING_SNAKE_CASE      ,     xs=jnp.arange(0      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE   )   )

							_UpperCAmelCase					:  Tuple											=   jnp.max(_SCREAMING_SNAKE_CASE      ,     axis=0      ,     keepdims=_SCREAMING_SNAKE_CASE   )
							_UpperCAmelCase					:  Any											=   jnp.exp(chunk_max - global_max   )

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

							_UpperCAmelCase					:  Optional[Any]											=   chunk_values.sum(axis=0   )
							_UpperCAmelCase					:  Any											=   jnp.expand_dims(_SCREAMING_SNAKE_CASE      ,     -1   ).sum(axis=0   )

							return all_values / all_weights

def      UpperCamelCase_							( _UpperCAmelCase  :				Optional[Any]      ,     _UpperCAmelCase  :				List[Any]      ,     _UpperCAmelCase  :				str      ,     _UpperCAmelCase  :				Tuple=jax.lax.Precision.HIGHEST      ,     _UpperCAmelCase  :				List[str] = 1_024      ,     _UpperCAmelCase  :				Union[str, Any] = 4_096   )   ->   str:


							"""simple docstring"""







							_UpperCAmelCase					:  int											=   query.shape[-3:]

							def chunk_scanner(_UpperCAmelCase  :				Optional[int]      ,     _UpperCAmelCase  :				Dict   ):
														# julienne query array
														_UpperCAmelCase					:  Union[str, 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   ),
														)

							_UpperCAmelCase					:  Optional[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      lowerCamelCase_     (nn.Module   ):
						'''simple docstring'''

						__UpperCamelCase:  int
						__UpperCamelCase:  int      =						8
						__UpperCamelCase:  int      =						6_4
						__UpperCamelCase:  float      =						0.0
						__UpperCamelCase:  bool      =						False
						__UpperCamelCase:  jnp.dtype      =						jnp.floataa




						def 					_A     (	self						:    List[Any]	):
													_UpperCAmelCase					:  List[str]											=   self.dim_head * self.heads
													_UpperCAmelCase					:  Any											=   self.dim_head**-0.5

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

													_UpperCAmelCase					:  List[Any]											=   nn.Dense(self.query_dim						, dtype=self.dtype						, name="to_out_0"	)
													_UpperCAmelCase					:  Union[str, Any]											=   nn.Dropout(rate=self.dropout	)




						def 					_A     (	self						:    List[Any]						, A						:    Optional[Any]	):
													_UpperCAmelCase					:  Union[str, Any]											=   tensor.shape
													_UpperCAmelCase					:  Union[str, Any]											=   self.heads
													_UpperCAmelCase					:  Tuple											=   tensor.reshape(UpperCamelCase__						, UpperCamelCase__						, UpperCamelCase__						, dim // head_size	)
													_UpperCAmelCase					:  Dict											=   jnp.transpose(UpperCamelCase__						, (0, 2, 1, 3)	)
													_UpperCAmelCase					:  Tuple											=   tensor.reshape(batch_size * head_size						, UpperCamelCase__						, dim // head_size	)
													return tensor




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





						def __call__(	self						:    Union[str, Any]						, A						:    Dict						, A						:    List[Any]=None						, A						:    int=True	):
													_UpperCAmelCase					:  int											=   hidden_states if context is None else context

													_UpperCAmelCase					:  Optional[Any]											=   self.query(UpperCamelCase__	)
													_UpperCAmelCase					:  Optional[Any]											=   self.key(UpperCamelCase__	)
													_UpperCAmelCase					:  Tuple											=   self.value(UpperCamelCase__	)

													_UpperCAmelCase					:  Optional[int]											=   self.reshape_heads_to_batch_dim(UpperCamelCase__	)
													_UpperCAmelCase					:  List[str]											=   self.reshape_heads_to_batch_dim(UpperCamelCase__	)
													_UpperCAmelCase					:  Any											=   self.reshape_heads_to_batch_dim(UpperCamelCase__	)

													if self.use_memory_efficient_attention:
																				_UpperCAmelCase					:  List[str]											=   query_states.transpose(1						, 0						, 2	)
																				_UpperCAmelCase					:  Tuple											=   key_states.transpose(1						, 0						, 2	)
																				_UpperCAmelCase					:  Optional[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

																				_UpperCAmelCase					:  Tuple											=   query_states.shape[-3]
																				if flatten_latent_dim % 64 == 0:
																											_UpperCAmelCase					:  str											=   int(flatten_latent_dim / 64	)
																				elif flatten_latent_dim % 16 == 0:
																											_UpperCAmelCase					:  Tuple											=   int(flatten_latent_dim / 16	)
																				elif flatten_latent_dim % 4 == 0:
																											_UpperCAmelCase					:  Tuple											=   int(flatten_latent_dim / 4	)
																				else:
																											_UpperCAmelCase					:  List[Any]											=   int(UpperCamelCase__	)

																				_UpperCAmelCase					:  List[Any]											=   jax_memory_efficient_attention(
																				    UpperCamelCase__						, UpperCamelCase__						, UpperCamelCase__						, query_chunk_size=UpperCamelCase__						, key_chunk_size=4096 * 4	)

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

																				# attend to values
																				_UpperCAmelCase					:  Optional[Any]											=   jnp.einsum("b i j, b j d -> b i d"						, UpperCamelCase__						, UpperCamelCase__	)

													_UpperCAmelCase					:  Tuple											=   self.reshape_batch_dim_to_heads(UpperCamelCase__	)
													_UpperCAmelCase					:  Union[str, Any]											=   self.proj_attn(UpperCamelCase__	)
													return self.dropout_layer(UpperCamelCase__						, deterministic=UpperCamelCase__	)

class      lowerCamelCase_     (nn.Module   ):
						'''simple docstring'''

						__UpperCamelCase:  int
						__UpperCamelCase:  int
						__UpperCamelCase:  int
						__UpperCamelCase:  float      =						0.0
						__UpperCamelCase:  bool      =						False
						__UpperCamelCase:  jnp.dtype      =						jnp.floataa
						__UpperCamelCase:  bool      =						False




						def 					_A     (	self						:    List[str]	):
													# self attention (or cross_attention if only_cross_attention is True)
													_UpperCAmelCase					:  Union[str, Any]											=   FlaxAttention(
													    self.dim						, self.n_heads						, self.d_head						, self.dropout						, self.use_memory_efficient_attention						, dtype=self.dtype	)
													# cross attention
													_UpperCAmelCase					:  Dict											=   FlaxAttention(
													    self.dim						, self.n_heads						, self.d_head						, self.dropout						, self.use_memory_efficient_attention						, dtype=self.dtype	)
													_UpperCAmelCase					:  List[Any]											=   FlaxFeedForward(dim=self.dim						, dropout=self.dropout						, dtype=self.dtype	)
													_UpperCAmelCase					:  Any											=   nn.LayerNorm(epsilon=1E-5						, dtype=self.dtype	)
													_UpperCAmelCase					:  Tuple											=   nn.LayerNorm(epsilon=1E-5						, dtype=self.dtype	)
													_UpperCAmelCase					:  Any											=   nn.LayerNorm(epsilon=1E-5						, dtype=self.dtype	)
													_UpperCAmelCase					:  Tuple											=   nn.Dropout(rate=self.dropout	)





						def __call__(	self						:    List[Any]						, A						:    Union[str, Any]						, A						:    Tuple						, A						:    Optional[int]=True	):
													# self attention
													_UpperCAmelCase					:  Tuple											=   hidden_states
													if self.only_cross_attention:
																				_UpperCAmelCase					:  Dict											=   self.attna(self.norma(UpperCamelCase__	)						, UpperCamelCase__						, deterministic=UpperCamelCase__	)
													else:
																				_UpperCAmelCase					:  Any											=   self.attna(self.norma(UpperCamelCase__	)						, deterministic=UpperCamelCase__	)
													_UpperCAmelCase					:  Dict											=   hidden_states + residual

													# cross attention
													_UpperCAmelCase					:  Union[str, Any]											=   hidden_states
													_UpperCAmelCase					:  Optional[Any]											=   self.attna(self.norma(UpperCamelCase__	)						, UpperCamelCase__						, deterministic=UpperCamelCase__	)
													_UpperCAmelCase					:  Tuple											=   hidden_states + residual

													# feed forward
													_UpperCAmelCase					:  Union[str, Any]											=   hidden_states
													_UpperCAmelCase					:  List[str]											=   self.ff(self.norma(UpperCamelCase__	)						, deterministic=UpperCamelCase__	)
													_UpperCAmelCase					:  Dict											=   hidden_states + residual

													return self.dropout_layer(UpperCamelCase__						, deterministic=UpperCamelCase__	)

class      lowerCamelCase_     (nn.Module   ):
						'''simple docstring'''

						__UpperCamelCase:  int
						__UpperCamelCase:  int
						__UpperCamelCase:  int
						__UpperCamelCase:  int      =						1
						__UpperCamelCase:  float      =						0.0
						__UpperCamelCase:  bool      =						False
						__UpperCamelCase:  bool      =						False
						__UpperCamelCase:  jnp.dtype      =						jnp.floataa
						__UpperCamelCase:  bool      =						False




						def 					_A     (	self						:    Optional[Any]	):
													_UpperCAmelCase					:  Optional[int]											=   nn.GroupNorm(num_groups=32						, epsilon=1E-5	)

													_UpperCAmelCase					:  Optional[Any]											=   self.n_heads * self.d_head
													if self.use_linear_projection:
																				_UpperCAmelCase					:  Any											=   nn.Dense(UpperCamelCase__						, dtype=self.dtype	)
													else:
																				_UpperCAmelCase					:  Dict											=   nn.Conv(
																				    UpperCamelCase__						, kernel_size=(1, 1)						, strides=(1, 1)						, padding="VALID"						, dtype=self.dtype						, )

													_UpperCAmelCase					:  Union[str, Any]											=   [
													    FlaxBasicTransformerBlock(
													        UpperCamelCase__						, 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:
																				_UpperCAmelCase					:  List[str]											=   nn.Dense(UpperCamelCase__						, dtype=self.dtype	)
													else:
																				_UpperCAmelCase					:  Union[str, Any]											=   nn.Conv(
																				    UpperCamelCase__						, kernel_size=(1, 1)						, strides=(1, 1)						, padding="VALID"						, dtype=self.dtype						, )

													_UpperCAmelCase					:  Union[str, Any]											=   nn.Dropout(rate=self.dropout	)





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

													for transformer_block in self.transformer_blocks:
																				_UpperCAmelCase					:  Optional[int]											=   transformer_block(UpperCamelCase__						, UpperCamelCase__						, deterministic=UpperCamelCase__	)

													if self.use_linear_projection:
																				_UpperCAmelCase					:  List[Any]											=   self.proj_out(UpperCamelCase__	)
																				_UpperCAmelCase					:  Any											=   hidden_states.reshape(UpperCamelCase__						, UpperCamelCase__						, UpperCamelCase__						, UpperCamelCase__	)
													else:
																				_UpperCAmelCase					:  Optional[int]											=   hidden_states.reshape(UpperCamelCase__						, UpperCamelCase__						, UpperCamelCase__						, UpperCamelCase__	)
																				_UpperCAmelCase					:  str											=   self.proj_out(UpperCamelCase__	)

													_UpperCAmelCase					:  Tuple											=   hidden_states + residual
													return self.dropout_layer(UpperCamelCase__						, deterministic=UpperCamelCase__	)

class      lowerCamelCase_     (nn.Module   ):
						'''simple docstring'''

						__UpperCamelCase:  int
						__UpperCamelCase:  float      =						0.0
						__UpperCamelCase:  jnp.dtype      =						jnp.floataa




						def 					_A     (	self						:    int	):
													# The second linear layer needs to be called
													# net_2 for now to match the index of the Sequential layer
													_UpperCAmelCase					:  List[Any]											=   FlaxGEGLU(self.dim						, self.dropout						, self.dtype	)
													_UpperCAmelCase					:  Dict											=   nn.Dense(self.dim						, dtype=self.dtype	)





						def __call__(	self						:    Optional[Any]						, A						:    int						, A						:    Tuple=True	):
													_UpperCAmelCase					:  List[str]											=   self.net_a(UpperCamelCase__						, deterministic=UpperCamelCase__	)
													_UpperCAmelCase					:  str											=   self.net_a(UpperCamelCase__	)
													return hidden_states

class      lowerCamelCase_     (nn.Module   ):
						'''simple docstring'''

						__UpperCamelCase:  int
						__UpperCamelCase:  float      =						0.0
						__UpperCamelCase:  jnp.dtype      =						jnp.floataa




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





						def __call__(	self						:    int						, A						:    Optional[int]						, A						:    Dict=True	):
													_UpperCAmelCase					:  Any											=   self.proj(UpperCamelCase__	)
													_UpperCAmelCase					:  Tuple											=   jnp.split(UpperCamelCase__						, 2						, axis=2	)
													return self.dropout_layer(hidden_linear * nn.gelu(UpperCamelCase__	)						, deterministic=UpperCamelCase__	)


'''simple docstring'''




def      UpperCamelCase_							( _UpperCAmelCase  :				int   )   ->   bool:


       """simple docstring"""







       _UpperCAmelCase					:  Optional[int]											=   9

       while num > 0:
              if num % 10 != digit:
                     return False
              num //= 100
              digit -= 1

       return True

def      UpperCamelCase_							( )   ->   int:


       """simple docstring"""







       _UpperCAmelCase					:  Optional[Any]											=   138_902_663

       while not is_square_form(num * num   ):
              if num % 10 == 3:
                     num -= 6  # (3 - 6) % 10 = 7
              else:
                     num -= 4  # (7 - 4) % 10 = 3

       return num * 10


if __name__ == "__main__":
  print(F'{solution() = }')

