






import functools
import math

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

def 			__lowerCamelCase     (						snake_case__			,snake_case__			,snake_case__			,snake_case__			,snake_case__ = 40_96						)				->    Tuple:






						"""simple docstring"""





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

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

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

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

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

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

						_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    jax.lax.map(f=snake_case__			,xs=jnp.arange(0			,snake_case__			,snake_case__						)						)

						_SCREAMING_SNAKE_CASE       =    jnp.max(snake_case__			,axis=0			,keepdims=snake_case__						)
						_SCREAMING_SNAKE_CASE       =    jnp.exp(chunk_max - global_max						)

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

						_SCREAMING_SNAKE_CASE       =    chunk_values.sum(axis=0						)
						_SCREAMING_SNAKE_CASE       =    jnp.expand_dims(snake_case__			,-1						).sum(axis=0						)

						return all_values / all_weights

def 			__lowerCamelCase     (						snake_case__			,snake_case__			,snake_case__			,snake_case__=jax.lax.Precision.HIGHEST			,snake_case__ = 10_24			,snake_case__ = 40_96						)				->    List[str]:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    query.shape[-3:]

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

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

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

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



class  __UpperCAmelCase						(nn.Module       ):
			__snake_case							: List[Any]							=				42
			__snake_case							: Tuple							=				8
			__snake_case							: Dict							=				64
			__snake_case							: Optional[Any]							=				0.0
			__snake_case							: Optional[Any]							=				False
			__snake_case							: Optional[Any]							=				jnp.floataa







			def 				UpperCamelCase (      self:    Union[str, Any]	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.dim_head * self.heads
									_SCREAMING_SNAKE_CASE       =    self.dim_head**-0.5

									# Weights were exported with old names {to_q, to_k, to_v, to_out}
									_SCREAMING_SNAKE_CASE       =    nn.Dense(_snake_case ,      use_bias=_snake_case ,      dtype=self.dtype ,      name="""to_q"""	)
									_SCREAMING_SNAKE_CASE       =    nn.Dense(_snake_case ,      use_bias=_snake_case ,      dtype=self.dtype ,      name="""to_k"""	)
									_SCREAMING_SNAKE_CASE       =    nn.Dense(_snake_case ,      use_bias=_snake_case ,      dtype=self.dtype ,      name="""to_v"""	)

									_SCREAMING_SNAKE_CASE       =    nn.Dense(self.query_dim ,      dtype=self.dtype ,      name="""to_out_0"""	)
									_SCREAMING_SNAKE_CASE       =    nn.Dropout(rate=self.dropout	)







			def 				UpperCamelCase (      self:    List[Any] ,      UpperCAmelCase_:    Any	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    tensor.shape
									_SCREAMING_SNAKE_CASE       =    self.heads
									_SCREAMING_SNAKE_CASE       =    tensor.reshape(_snake_case ,      _snake_case ,      _snake_case ,      dim // head_size	)
									_SCREAMING_SNAKE_CASE       =    jnp.transpose(_snake_case ,      (0, 2, 1, 3)	)
									_SCREAMING_SNAKE_CASE       =    tensor.reshape(batch_size * head_size ,      _snake_case ,      dim // head_size	)
									return tensor







			def 				UpperCamelCase (      self:    Union[str, Any] ,      UpperCAmelCase_:    int	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    tensor.shape
									_SCREAMING_SNAKE_CASE       =    self.heads
									_SCREAMING_SNAKE_CASE       =    tensor.reshape(batch_size // head_size ,      _snake_case ,      _snake_case ,      _snake_case	)
									_SCREAMING_SNAKE_CASE       =    jnp.transpose(_snake_case ,      (0, 2, 1, 3)	)
									_SCREAMING_SNAKE_CASE       =    tensor.reshape(batch_size // head_size ,      _snake_case ,      dim * head_size	)
									return tensor







			def __call__(      self:    Optional[int] ,      UpperCAmelCase_:    Union[str, Any] ,      UpperCAmelCase_:    List[str]=None ,      UpperCAmelCase_:    Union[str, Any]=True	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    hidden_states if context is None else context

									_SCREAMING_SNAKE_CASE       =    self.query(_snake_case	)
									_SCREAMING_SNAKE_CASE       =    self.key(_snake_case	)
									_SCREAMING_SNAKE_CASE       =    self.value(_snake_case	)

									_SCREAMING_SNAKE_CASE       =    self.reshape_heads_to_batch_dim(_snake_case	)
									_SCREAMING_SNAKE_CASE       =    self.reshape_heads_to_batch_dim(_snake_case	)
									_SCREAMING_SNAKE_CASE       =    self.reshape_heads_to_batch_dim(_snake_case	)

									if self.use_memory_efficient_attention:
															_SCREAMING_SNAKE_CASE       =    query_states.transpose(1 ,      0 ,      2	)
															_SCREAMING_SNAKE_CASE       =    key_states.transpose(1 ,      0 ,      2	)
															_SCREAMING_SNAKE_CASE       =    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

															_SCREAMING_SNAKE_CASE       =    query_states.shape[-3]
															if flatten_latent_dim % 64 == 0:
																					_SCREAMING_SNAKE_CASE       =    int(flatten_latent_dim / 64	)
															elif flatten_latent_dim % 16 == 0:
																					_SCREAMING_SNAKE_CASE       =    int(flatten_latent_dim / 16	)
															elif flatten_latent_dim % 4 == 0:
																					_SCREAMING_SNAKE_CASE       =    int(flatten_latent_dim / 4	)
															else:
																					_SCREAMING_SNAKE_CASE       =    int(_snake_case	)

															_SCREAMING_SNAKE_CASE       =    jax_memory_efficient_attention(
															    _snake_case ,      _snake_case ,      _snake_case ,      query_chunk_size=_snake_case ,      key_chunk_size=4_096 * 4	)

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

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

									_SCREAMING_SNAKE_CASE       =    self.reshape_batch_dim_to_heads(_snake_case	)
									_SCREAMING_SNAKE_CASE       =    self.proj_attn(_snake_case	)
									return self.dropout_layer(_snake_case ,      deterministic=_snake_case	)



class  __UpperCAmelCase						(nn.Module       ):
			__snake_case							: Optional[int]							=				42
			__snake_case							: List[Any]							=				42
			__snake_case							: Dict							=				42
			__snake_case							: Tuple							=				0.0
			__snake_case							: int							=				False
			__snake_case							: Tuple							=				jnp.floataa
			__snake_case							: int							=				False







			def 				UpperCamelCase (      self:    Optional[int]	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    FlaxAttention(
									    self.dim ,      self.n_heads ,      self.d_head ,      self.dropout ,      self.use_memory_efficient_attention ,      dtype=self.dtype	)
									# cross attention
									_SCREAMING_SNAKE_CASE       =    FlaxAttention(
									    self.dim ,      self.n_heads ,      self.d_head ,      self.dropout ,      self.use_memory_efficient_attention ,      dtype=self.dtype	)
									_SCREAMING_SNAKE_CASE       =    FlaxFeedForward(dim=self.dim ,      dropout=self.dropout ,      dtype=self.dtype	)
									_SCREAMING_SNAKE_CASE       =    nn.LayerNorm(epsilon=1E-5 ,      dtype=self.dtype	)
									_SCREAMING_SNAKE_CASE       =    nn.LayerNorm(epsilon=1E-5 ,      dtype=self.dtype	)
									_SCREAMING_SNAKE_CASE       =    nn.LayerNorm(epsilon=1E-5 ,      dtype=self.dtype	)
									_SCREAMING_SNAKE_CASE       =    nn.Dropout(rate=self.dropout	)







			def __call__(      self:    Any ,      UpperCAmelCase_:    Tuple ,      UpperCAmelCase_:    int ,      UpperCAmelCase_:    Optional[Any]=True	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    hidden_states
									if self.only_cross_attention:
															_SCREAMING_SNAKE_CASE       =    self.attna(self.norma(_snake_case	) ,      _snake_case ,      deterministic=_snake_case	)
									else:
															_SCREAMING_SNAKE_CASE       =    self.attna(self.norma(_snake_case	) ,      deterministic=_snake_case	)
									_SCREAMING_SNAKE_CASE       =    hidden_states + residual

									# cross attention
									_SCREAMING_SNAKE_CASE       =    hidden_states
									_SCREAMING_SNAKE_CASE       =    self.attna(self.norma(_snake_case	) ,      _snake_case ,      deterministic=_snake_case	)
									_SCREAMING_SNAKE_CASE       =    hidden_states + residual

									# feed forward
									_SCREAMING_SNAKE_CASE       =    hidden_states
									_SCREAMING_SNAKE_CASE       =    self.ff(self.norma(_snake_case	) ,      deterministic=_snake_case	)
									_SCREAMING_SNAKE_CASE       =    hidden_states + residual

									return self.dropout_layer(_snake_case ,      deterministic=_snake_case	)



class  __UpperCAmelCase						(nn.Module       ):
			__snake_case							: Optional[Any]							=				42
			__snake_case							: Dict							=				42
			__snake_case							: Any							=				42
			__snake_case							: Any							=				1
			__snake_case							: Dict							=				0.0
			__snake_case							: List[str]							=				False
			__snake_case							: Tuple							=				False
			__snake_case							: str							=				jnp.floataa
			__snake_case							: Any							=				False







			def 				UpperCamelCase (      self:    Tuple	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    nn.GroupNorm(num_groups=32 ,      epsilon=1E-5	)

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

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

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







			def __call__(      self:    List[str] ,      UpperCAmelCase_:    Tuple ,      UpperCAmelCase_:    int ,      UpperCAmelCase_:    int=True	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    hidden_states.shape
									_SCREAMING_SNAKE_CASE       =    hidden_states
									_SCREAMING_SNAKE_CASE       =    self.norm(_snake_case	)
									if self.use_linear_projection:
															_SCREAMING_SNAKE_CASE       =    hidden_states.reshape(_snake_case ,      height * width ,      _snake_case	)
															_SCREAMING_SNAKE_CASE       =    self.proj_in(_snake_case	)
									else:
															_SCREAMING_SNAKE_CASE       =    self.proj_in(_snake_case	)
															_SCREAMING_SNAKE_CASE       =    hidden_states.reshape(_snake_case ,      height * width ,      _snake_case	)

									for transformer_block in self.transformer_blocks:
															_SCREAMING_SNAKE_CASE       =    transformer_block(_snake_case ,      _snake_case ,      deterministic=_snake_case	)

									if self.use_linear_projection:
															_SCREAMING_SNAKE_CASE       =    self.proj_out(_snake_case	)
															_SCREAMING_SNAKE_CASE       =    hidden_states.reshape(_snake_case ,      _snake_case ,      _snake_case ,      _snake_case	)
									else:
															_SCREAMING_SNAKE_CASE       =    hidden_states.reshape(_snake_case ,      _snake_case ,      _snake_case ,      _snake_case	)
															_SCREAMING_SNAKE_CASE       =    self.proj_out(_snake_case	)

									_SCREAMING_SNAKE_CASE       =    hidden_states + residual
									return self.dropout_layer(_snake_case ,      deterministic=_snake_case	)



class  __UpperCAmelCase						(nn.Module       ):
			__snake_case							: Dict							=				42
			__snake_case							: Any							=				0.0
			__snake_case							: List[str]							=				jnp.floataa







			def 				UpperCamelCase (      self:    Optional[Any]	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    FlaxGEGLU(self.dim ,      self.dropout ,      self.dtype	)
									_SCREAMING_SNAKE_CASE       =    nn.Dense(self.dim ,      dtype=self.dtype	)







			def __call__(      self:    Tuple ,      UpperCAmelCase_:    Dict ,      UpperCAmelCase_:    Tuple=True	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.net_a(_snake_case ,      deterministic=_snake_case	)
									_SCREAMING_SNAKE_CASE       =    self.net_a(_snake_case	)
									return hidden_states



class  __UpperCAmelCase						(nn.Module       ):
			__snake_case							: int							=				42
			__snake_case							: Union[str, Any]							=				0.0
			__snake_case							: Union[str, Any]							=				jnp.floataa







			def 				UpperCamelCase (      self:    Dict	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.dim * 4
									_SCREAMING_SNAKE_CASE       =    nn.Dense(inner_dim * 2 ,      dtype=self.dtype	)
									_SCREAMING_SNAKE_CASE       =    nn.Dropout(rate=self.dropout	)







			def __call__(      self:    Dict ,      UpperCAmelCase_:    List[Any] ,      UpperCAmelCase_:    List[str]=True	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.proj(_snake_case	)
									_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE       =    jnp.split(_snake_case ,      2 ,      axis=2	)
									return self.dropout_layer(hidden_linear * nn.gelu(_snake_case	) ,      deterministic=_snake_case	)






# A Bipartite Graph is a graph whose vertices can be divided into two independent sets,
# U and V such that every edge (u, v) either connects a vertex from U to V or a vertex
# from V to U. In other words, for every edge (u, v), either u belongs to U and v to V,
# or u belongs to V and v to U. We can also say that there is no edge that connects
# vertices of same set.
from queue import Queue


def 						_UpperCAmelCase							( snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       Queue()
					_lowerCAmelCase         =       [False] * len(snake_case     )
					_lowerCAmelCase         =       [-1] * len(snake_case     )

					def bfs():
										while not queue.empty():
															_lowerCAmelCase         =       queue.get()
															_lowerCAmelCase         =       True

															for neighbour in graph[u]:
																				if neighbour == u:
																									return False

																				if color[neighbour] == -1:
																									_lowerCAmelCase         =       1 - color[u]
																									queue.put(snake_case     )

																				elif color[neighbour] == color[u]:
																									return False

										return True

					for i in range(len(snake_case     )     ):
										if not visited[i]:
															queue.put(snake_case     )
															_lowerCAmelCase         =       0
															if bfs() is False:
																				return False

					return True


if __name__ == "__main__":
	# Adjacency List of graph
	print(check_bipartite({0: [1, 3], 1: [0, 2], 2: [1, 3], 3: [0, 2]}))



