



'''simple docstring'''
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings






@maybe_allow_in_graph
class    __lowerCamelCase							( nn.Module   ):







	'''simple docstring'''






	def __init__(      self      ,     a__      ,     a__      ,     a__      ,     a__=0.0      ,     a__ = None      ,     a__ = "geglu"      ,     a__ = None      ,     a__ = False      ,     a__ = False      ,     a__ = False      ,     a__ = False      ,     a__ = True      ,     a__ = "layer_norm"      ,     a__ = False      ,     ):
						super().__init__()
						__SCREAMING_SNAKE_CASE				:    Optional[int]						    =							only_cross_attention

						__SCREAMING_SNAKE_CASE				:    List[str]						    =							(num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
						__SCREAMING_SNAKE_CASE				:    Any						    =							(num_embeds_ada_norm is not None) and norm_type == '''ada_norm'''

						if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
											raise ValueError(
											    f'`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'
											    f' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'					)

						# Define 3 blocks. Each block has its own normalization layer.
						# 1. Self-Attn
						if self.use_ada_layer_norm:
											__SCREAMING_SNAKE_CASE				:    Tuple						    =							AdaLayerNorm(__lowerCamelCase      ,     __lowerCamelCase					)
						elif self.use_ada_layer_norm_zero:
											__SCREAMING_SNAKE_CASE				:    List[Any]						    =							AdaLayerNormZero(__lowerCamelCase      ,     __lowerCamelCase					)
						else:
											__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							nn.LayerNorm(__lowerCamelCase      ,     elementwise_affine=__lowerCamelCase					)
						__SCREAMING_SNAKE_CASE				:    Any						    =							Attention(
						    query_dim=__lowerCamelCase      ,     heads=__lowerCamelCase      ,     dim_head=__lowerCamelCase      ,     dropout=__lowerCamelCase      ,     bias=__lowerCamelCase      ,     cross_attention_dim=cross_attention_dim if only_cross_attention else None      ,     upcast_attention=__lowerCamelCase      ,     )

						# 2. Cross-Attn
						if cross_attention_dim is not None or double_self_attention:
											# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
											# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
											# the second cross attention block.
											__SCREAMING_SNAKE_CASE				:    int						    =							(
											    AdaLayerNorm(__lowerCamelCase      ,     __lowerCamelCase					)
											    if self.use_ada_layer_norm
											    else nn.LayerNorm(__lowerCamelCase      ,     elementwise_affine=__lowerCamelCase					)
											)
											__SCREAMING_SNAKE_CASE				:    List[str]						    =							Attention(
											    query_dim=__lowerCamelCase      ,     cross_attention_dim=cross_attention_dim if not double_self_attention else None      ,     heads=__lowerCamelCase      ,     dim_head=__lowerCamelCase      ,     dropout=__lowerCamelCase      ,     bias=__lowerCamelCase      ,     upcast_attention=__lowerCamelCase      ,     )  # is self-attn if encoder_hidden_states is none
						else:
											__SCREAMING_SNAKE_CASE				:    Tuple						    =							None
											__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							None

						# 3. Feed-forward
						__SCREAMING_SNAKE_CASE				:    List[str]						    =							nn.LayerNorm(__lowerCamelCase      ,     elementwise_affine=__lowerCamelCase					)
						__SCREAMING_SNAKE_CASE				:    Any						    =							FeedForward(__lowerCamelCase      ,     dropout=__lowerCamelCase      ,     activation_fn=__lowerCamelCase      ,     final_dropout=__lowerCamelCase					)

						# let chunk size default to None
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							None
						__SCREAMING_SNAKE_CASE				:    Dict						    =							0






	def  a_					(      self      ,     a__      ,     a__					):
						# Sets chunk feed-forward
						__SCREAMING_SNAKE_CASE				:    List[str]						    =							chunk_size
						__SCREAMING_SNAKE_CASE				:    str						    =							dim






	def  a_					(      self      ,     a__      ,     a__ = None      ,     a__ = None      ,     a__ = None      ,     a__ = None      ,     a__ = None      ,     a__ = None      ,     ):
						# Notice that normalization is always applied before the real computation in the following blocks.
						# 1. Self-Attention
						if self.use_ada_layer_norm:
											__SCREAMING_SNAKE_CASE				:    Dict						    =							self.norma(__lowerCamelCase      ,     __lowerCamelCase					)
						elif self.use_ada_layer_norm_zero:
											__SCREAMING_SNAKE_CASE				:    Tuple						    =							self.norma(
											    __lowerCamelCase      ,     __lowerCamelCase      ,     __lowerCamelCase      ,     hidden_dtype=hidden_states.dtype					)
						else:
											__SCREAMING_SNAKE_CASE				:    Tuple						    =							self.norma(__lowerCamelCase					)

						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							cross_attention_kwargs if cross_attention_kwargs is not None else {}

						__SCREAMING_SNAKE_CASE				:    Any						    =							self.attna(
						    __lowerCamelCase      ,     encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None      ,     attention_mask=__lowerCamelCase      ,     **__lowerCamelCase      ,     )
						if self.use_ada_layer_norm_zero:
											__SCREAMING_SNAKE_CASE				:    List[str]						    =							gate_msa.unsqueeze(1					) * attn_output
						__SCREAMING_SNAKE_CASE				:    int						    =							attn_output + hidden_states

						# 2. Cross-Attention
						if self.attna is not None:
											__SCREAMING_SNAKE_CASE				:    Dict						    =							(
											    self.norma(__lowerCamelCase      ,     __lowerCamelCase					) if self.use_ada_layer_norm else self.norma(__lowerCamelCase					)
											)

											__SCREAMING_SNAKE_CASE				:    Any						    =							self.attna(
											    __lowerCamelCase      ,     encoder_hidden_states=__lowerCamelCase      ,     attention_mask=__lowerCamelCase      ,     **__lowerCamelCase      ,     )
											__SCREAMING_SNAKE_CASE				:    Tuple						    =							attn_output + hidden_states

						# 3. Feed-forward
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							self.norma(__lowerCamelCase					)

						if self.use_ada_layer_norm_zero:
											__SCREAMING_SNAKE_CASE				:    Dict						    =							norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

						if self._chunk_size is not None:
											# "feed_forward_chunk_size" can be used to save memory
											if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
																raise ValueError(
																    f'`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'					)

											__SCREAMING_SNAKE_CASE				:    List[str]						    =							norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
											__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							torch.cat(
											    [self.ff(__lowerCamelCase					) for hid_slice in norm_hidden_states.chunk(__lowerCamelCase      ,     dim=self._chunk_dim					)]      ,     dim=self._chunk_dim      ,     )
						else:
											__SCREAMING_SNAKE_CASE				:    List[str]						    =							self.ff(__lowerCamelCase					)

						if self.use_ada_layer_norm_zero:
											__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							gate_mlp.unsqueeze(1					) * ff_output

						__SCREAMING_SNAKE_CASE				:    List[str]						    =							ff_output + hidden_states

						return hidden_states






class    __lowerCamelCase							( nn.Module   ):







	'''simple docstring'''






	def __init__(      self      ,     a__      ,     a__ = None      ,     a__ = 4      ,     a__ = 0.0      ,     a__ = "geglu"      ,     a__ = False      ,     ):
						super().__init__()
						__SCREAMING_SNAKE_CASE				:    int						    =							int(dim * mult					)
						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							dim_out if dim_out is not None else dim

						if activation_fn == "gelu":
											__SCREAMING_SNAKE_CASE				:    List[Any]						    =							GELU(__lowerCamelCase      ,     __lowerCamelCase					)
						if activation_fn == "gelu-approximate":
											__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							GELU(__lowerCamelCase      ,     __lowerCamelCase      ,     approximate="tanh"					)
						elif activation_fn == "geglu":
											__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							GEGLU(__lowerCamelCase      ,     __lowerCamelCase					)
						elif activation_fn == "geglu-approximate":
											__SCREAMING_SNAKE_CASE				:    Any						    =							ApproximateGELU(__lowerCamelCase      ,     __lowerCamelCase					)

						__SCREAMING_SNAKE_CASE				:    List[str]						    =							nn.ModuleList([]					)
						# project in
						self.net.append(__lowerCamelCase					)
						# project dropout
						self.net.append(nn.Dropout(__lowerCamelCase					)					)
						# project out
						self.net.append(nn.Linear(__lowerCamelCase      ,     __lowerCamelCase					)					)
						# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
						if final_dropout:
											self.net.append(nn.Dropout(__lowerCamelCase					)					)






	def  a_					(      self      ,     a__					):
						for module in self.net:
											__SCREAMING_SNAKE_CASE				:    List[Any]						    =							module(__lowerCamelCase					)
						return hidden_states






class    __lowerCamelCase							( nn.Module   ):







	'''simple docstring'''






	def __init__(      self      ,     a__      ,     a__      ,     a__ = "none"					):
						super().__init__()
						__SCREAMING_SNAKE_CASE				:    Tuple						    =							nn.Linear(__lowerCamelCase      ,     __lowerCamelCase					)
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							approximate






	def  a_					(      self      ,     a__					):
						if gate.device.type != "mps":
											return F.gelu(__lowerCamelCase      ,     approximate=self.approximate					)
						# mps: gelu is not implemented for float16
						return F.gelu(gate.to(dtype=torch.floataa					)      ,     approximate=self.approximate					).to(dtype=gate.dtype					)






	def  a_					(      self      ,     a__					):
						__SCREAMING_SNAKE_CASE				:    str						    =							self.proj(__lowerCamelCase					)
						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							self.gelu(__lowerCamelCase					)
						return hidden_states






class    __lowerCamelCase							( nn.Module   ):







	'''simple docstring'''






	def __init__(      self      ,     a__      ,     a__					):
						super().__init__()
						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							nn.Linear(__lowerCamelCase      ,     dim_out * 2					)






	def  a_					(      self      ,     a__					):
						if gate.device.type != "mps":
											return F.gelu(__lowerCamelCase					)
						# mps: gelu is not implemented for float16
						return F.gelu(gate.to(dtype=torch.floataa					)					).to(dtype=gate.dtype					)






	def  a_					(      self      ,     a__					):
						__SCREAMING_SNAKE_CASE				:    int						    =							self.proj(__lowerCamelCase					).chunk(2      ,     dim=-1					)
						return hidden_states * self.gelu(__lowerCamelCase					)






class    __lowerCamelCase							( nn.Module   ):







	'''simple docstring'''






	def __init__(      self      ,     a__      ,     a__					):
						super().__init__()
						__SCREAMING_SNAKE_CASE				:    List[str]						    =							nn.Linear(__lowerCamelCase      ,     __lowerCamelCase					)






	def  a_					(      self      ,     a__					):
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							self.proj(__lowerCamelCase					)
						return x * torch.sigmoid(1.702 * x					)






class    __lowerCamelCase							( nn.Module   ):







	'''simple docstring'''






	def __init__(      self      ,     a__      ,     a__					):
						super().__init__()
						__SCREAMING_SNAKE_CASE				:    str						    =							nn.Embedding(__lowerCamelCase      ,     __lowerCamelCase					)
						__SCREAMING_SNAKE_CASE				:    Dict						    =							nn.SiLU()
						__SCREAMING_SNAKE_CASE				:    Any						    =							nn.Linear(__lowerCamelCase      ,     embedding_dim * 2					)
						__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							nn.LayerNorm(__lowerCamelCase      ,     elementwise_affine=__lowerCamelCase					)






	def  a_					(      self      ,     a__      ,     a__					):
						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							self.linear(self.silu(self.emb(__lowerCamelCase					)					)					)
						__SCREAMING_SNAKE_CASE				:    Optional[int]						    =							torch.chunk(__lowerCamelCase      ,     2					)
						__SCREAMING_SNAKE_CASE				:    Dict						    =							self.norm(__lowerCamelCase					) * (1 + scale) + shift
						return x






class    __lowerCamelCase							( nn.Module   ):







	'''simple docstring'''






	def __init__(      self      ,     a__      ,     a__					):
						super().__init__()

						__SCREAMING_SNAKE_CASE				:    Dict						    =							CombinedTimestepLabelEmbeddings(__lowerCamelCase      ,     __lowerCamelCase					)

						__SCREAMING_SNAKE_CASE				:    str						    =							nn.SiLU()
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							nn.Linear(__lowerCamelCase      ,     6 * embedding_dim      ,     bias=__lowerCamelCase					)
						__SCREAMING_SNAKE_CASE				:    Optional[int]						    =							nn.LayerNorm(__lowerCamelCase      ,     elementwise_affine=__lowerCamelCase      ,     eps=1e-6					)






	def  a_					(      self      ,     a__      ,     a__      ,     a__      ,     a__=None					):
						__SCREAMING_SNAKE_CASE				:    Optional[int]						    =							self.linear(self.silu(self.emb(__lowerCamelCase      ,     __lowerCamelCase      ,     hidden_dtype=__lowerCamelCase					)					)					)
						__SCREAMING_SNAKE_CASE				:    str						    =							emb.chunk(6      ,     dim=1					)
						__SCREAMING_SNAKE_CASE				:    Any						    =							self.norm(__lowerCamelCase					) * (1 + scale_msa[:, None]) + shift_msa[:, None]
						return x, gate_msa, shift_mlp, scale_mlp, gate_mlp






class    __lowerCamelCase							( nn.Module   ):







	'''simple docstring'''






	def __init__(      self      ,     a__      ,     a__      ,     a__      ,     a__ = None      ,     a__ = 1e-5					):
						super().__init__()
						__SCREAMING_SNAKE_CASE				:    str						    =							num_groups
						__SCREAMING_SNAKE_CASE				:    Any						    =							eps

						if act_fn is None:
											__SCREAMING_SNAKE_CASE				:    List[str]						    =							None
						else:
											__SCREAMING_SNAKE_CASE				:    Dict						    =							get_activation(__lowerCamelCase					)

						__SCREAMING_SNAKE_CASE				:    Any						    =							nn.Linear(__lowerCamelCase      ,     out_dim * 2					)






	def  a_					(      self      ,     a__      ,     a__					):
						if self.act:
											__SCREAMING_SNAKE_CASE				:    List[Any]						    =							self.act(__lowerCamelCase					)
						__SCREAMING_SNAKE_CASE				:    Dict						    =							self.linear(__lowerCamelCase					)
						__SCREAMING_SNAKE_CASE				:    Dict						    =							emb[:, :, None, None]
						__SCREAMING_SNAKE_CASE				:    int						    =							emb.chunk(2      ,     dim=1					)

						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							F.group_norm(__lowerCamelCase      ,     self.num_groups      ,     eps=self.eps					)
						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							x * (1 + scale) + shift
						return x






'''simple docstring'''





from multiprocessing import Lock, Pipe, Process

# lock used to ensure that two processes do not access a pipe at the same time
__lowerCamelCase												=						Lock()


def        a__     (					UpperCamelCase_						:      str,   UpperCamelCase_						:      Any,   UpperCamelCase_						:      Any,   UpperCamelCase_						:      int,   UpperCamelCase_						:      Union[str, Any],   UpperCamelCase_						:      List[str],   UpperCamelCase_						:      Optional[int]						):
   global process_lock

   # we perform n swaps since after n swaps we know we are sorted
   # we *could* stop early if we are sorted already, but it takes as long to
   # find out we are sorted as it does to sort the list with this algorithm
   for i in range(0,   10						):
      if (i + position) % 2 == 0 and r_send is not None:
         # send your value to your right neighbor
         process_lock.acquire()
         r_send[1].send(UpperCamelCase_						)
         process_lock.release()

         # receive your right neighbor's value
         process_lock.acquire()
         UpperCAmelCase__	:Union[str, Any]      =      rr_cv[0].recv()
         process_lock.release()

         # take the lower value since you are on the left
         UpperCAmelCase__	:List[Any]      =      min(UpperCamelCase_,   UpperCamelCase_						)
      elif (i + position) % 2 != 0 and l_send is not None:
         # send your value to your left neighbor
         process_lock.acquire()
         l_send[1].send(UpperCamelCase_						)
         process_lock.release()

         # receive your left neighbor's value
         process_lock.acquire()
         UpperCAmelCase__	:List[str]      =      lr_cv[0].recv()
         process_lock.release()

         # take the higher value since you are on the right
         UpperCAmelCase__	:List[Any]      =      max(UpperCamelCase_,   UpperCamelCase_						)
    # after all swaps are performed, send the values back to main
   result_pipe[1].send(UpperCamelCase_						)


def        a__     (					UpperCamelCase_						:      int						):
   UpperCAmelCase__	:List[str]      =      []
   UpperCAmelCase__	:Any      =      []
   # initialize the list of pipes where the values will be retrieved
   for _ in arr:
      result_pipe.append(Pipe()						)
   # creates the processes
   # the first and last process only have one neighbor so they are made outside
   # of the loop
   UpperCAmelCase__	:str      =      Pipe()
   UpperCAmelCase__	:Dict      =      Pipe()
   process_array_.append(
       Process(
           target=UpperCamelCase_,   args=(0, arr[0], None, temp_rs, None, temp_rr, result_pipe[0]),   )						)
   UpperCAmelCase__	:Optional[Any]      =      temp_rs
   UpperCAmelCase__	:Any      =      temp_rr

   for i in range(1,   len(UpperCamelCase_						) - 1						):
      UpperCAmelCase__	:Optional[int]      =      Pipe()
      UpperCAmelCase__	:List[str]      =      Pipe()
      process_array_.append(
          Process(
              target=UpperCamelCase_,   args=(i, arr[i], temp_ls, temp_rs, temp_lr, temp_rr, result_pipe[i]),   )						)
      UpperCAmelCase__	:Optional[Any]      =      temp_rs
      UpperCAmelCase__	:str      =      temp_rr

   process_array_.append(
       Process(
           target=UpperCamelCase_,   args=(
               len(UpperCamelCase_						) - 1,
               arr[len(UpperCamelCase_						) - 1],
               temp_ls,
               None,
               temp_lr,
               None,
               result_pipe[len(UpperCamelCase_						) - 1],
           ),   )						)

   # start the processes
   for p in process_array_:
      p.start()

   # wait for the processes to end and write their values to the list
   for p in range(0,   len(UpperCamelCase_						)						):
      UpperCAmelCase__	:List[str]      =      result_pipe[p][0].recv()
      process_array_[p].join()
   return arr


def        a__     (					):
   UpperCAmelCase__	:Optional[Any]      =      list(range(10,   0,   -1						)						)
   print('''Initial List'''						)
   print(*UpperCamelCase_						)
   UpperCAmelCase__	:str      =      odd_even_transposition(UpperCamelCase_						)
   print('''Sorted List\n'''						)
   print(*UpperCamelCase_						)


if __name__ == "__main__":
    main()



