code: 
  
  
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  from dataclasses import dataclass
  from typing import Any, Dict, Optional
  
  import torch
  import torch.nn.functional as F
  from torch import nn
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..models.embeddings import ImagePositionalEmbeddings
  from ..utils import BaseOutput, deprecate
  from .attention import BasicTransformerBlock
  from .embeddings import PatchEmbed
  from .modeling_utils import ModelMixin
  
  @dataclass
  class 					snake_case_							(       UpperCamelCase__      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
     __lowerCAmelCase      =						42
  
  
  
  
  
  
  class 					snake_case_							(       UpperCamelCase__      ,			UpperCamelCase__      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
     @register_to_config
     def __init__(    self:	List[str]   ,       snake_case__:	Dict = 16   ,       snake_case__:	Optional[int] = 88   ,       snake_case__:	Any = None   ,       snake_case__:	List[Any] = None   ,       snake_case__:	int = 1   ,       snake_case__:	Optional[int] = 0.0   ,       snake_case__:	Optional[int] = 32   ,       snake_case__:	Dict = None   ,       snake_case__:	Union[str, Any] = False   ,       snake_case__:	Tuple = None   ,       snake_case__:	Optional[Any] = None   ,       snake_case__:	Tuple = None   ,       snake_case__:	Any = "geglu"   ,       snake_case__:	Tuple = None   ,       snake_case__:	Tuple = False   ,       snake_case__:	Any = False   ,       snake_case__:	Optional[Any] = False   ,       snake_case__:	str = "layer_norm"   ,       snake_case__:	Optional[Any] = True   ,       )  ->		Tuple:
  
  
         """simple docstring"""
  
  
  
         super().__init__()
         _lowercase     =	use_linear_projection
         _lowercase     =	num_attention_heads
         _lowercase     =	attention_head_dim
         _lowercase     =	num_attention_heads * attention_head_dim
  
         # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
         # Define whether input is continuous or discrete depending on configuration
         _lowercase     =	(in_channels is not None) and (patch_size is None)
         _lowercase     =	num_vector_embeds is not None
         _lowercase     =	in_channels is not None and patch_size is not None
  
         if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
             _lowercase     =	(
                 f'''The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or'''
                 ''' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'''
                 ''' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'''
                 ''' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'''
                 ''' would be very nice if you could open a Pull request for the `transformer/config.json` file'''
             )
             deprecate('norm_type!=num_embeds_ada_norm'   ,       '1.0.0'   ,       snake_case_   ,       standard_warn=snake_case_	)
             _lowercase     =	'''ada_norm'''
  
         if self.is_input_continuous and self.is_input_vectorized:
             raise ValueError(
                 f'''Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make'''
                 ' sure that either `in_channels` or `num_vector_embeds` is None.'	)
         elif self.is_input_vectorized and self.is_input_patches:
             raise ValueError(
                 f'''Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make'''
                 ' sure that either `num_vector_embeds` or `num_patches` is None.'	)
         elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
             raise ValueError(
                 f'''Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:'''
                 f''' {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None.'''	)
  
         # 2. Define input layers
         if self.is_input_continuous:
             _lowercase     =	in_channels
  
             _lowercase     =	torch.nn.GroupNorm(num_groups=snake_case_   ,       num_channels=snake_case_   ,       eps=1E-6   ,       affine=snake_case_	)
             if use_linear_projection:
                 _lowercase     =	nn.Linear(snake_case_   ,       snake_case_	)
             else:
                 _lowercase     =	nn.Convad(snake_case_   ,       snake_case_   ,       kernel_size=1   ,       stride=1   ,       padding=0	)
         elif self.is_input_vectorized:
             assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
             assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"
  
             _lowercase     =	sample_size
             _lowercase     =	sample_size
             _lowercase     =	num_vector_embeds
             _lowercase     =	self.height * self.width
  
             _lowercase     =	ImagePositionalEmbeddings(
                 num_embed=snake_case_   ,       embed_dim=snake_case_   ,       height=self.height   ,       width=self.width	)
         elif self.is_input_patches:
             assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"
  
             _lowercase     =	sample_size
             _lowercase     =	sample_size
  
             _lowercase     =	patch_size
             _lowercase     =	PatchEmbed(
                 height=snake_case_   ,       width=snake_case_   ,       patch_size=snake_case_   ,       in_channels=snake_case_   ,       embed_dim=snake_case_   ,       )
  
         # 3. Define transformers blocks
         _lowercase     =	nn.ModuleList(
             [
                 BasicTransformerBlock(
                     snake_case_   ,       snake_case_   ,       snake_case_   ,       dropout=snake_case_   ,       cross_attention_dim=snake_case_   ,       activation_fn=snake_case_   ,       num_embeds_ada_norm=snake_case_   ,       attention_bias=snake_case_   ,       only_cross_attention=snake_case_   ,       upcast_attention=snake_case_   ,       norm_type=snake_case_   ,       norm_elementwise_affine=snake_case_   ,       )
                 for d in range(snake_case_	)
             ]	)
  
         # 4. Define output layers
         _lowercase     =	in_channels if out_channels is None else out_channels
         if self.is_input_continuous:
             # TODO: should use out_channels for continuous projections
             if use_linear_projection:
                 _lowercase     =	nn.Linear(snake_case_   ,       snake_case_	)
             else:
                 _lowercase     =	nn.Convad(snake_case_   ,       snake_case_   ,       kernel_size=1   ,       stride=1   ,       padding=0	)
         elif self.is_input_vectorized:
             _lowercase     =	nn.LayerNorm(snake_case_	)
             _lowercase     =	nn.Linear(snake_case_   ,       self.num_vector_embeds - 1	)
         elif self.is_input_patches:
             _lowercase     =	nn.LayerNorm(snake_case_   ,       elementwise_affine=snake_case_   ,       eps=1E-6	)
             _lowercase     =	nn.Linear(snake_case_   ,       2 * inner_dim	)
             _lowercase     =	nn.Linear(snake_case_   ,       patch_size * patch_size * self.out_channels	)
  
     def 			SCREAMING_SNAKE_CASE__				(    self:	int   ,       snake_case__:	str   ,       snake_case__:	Union[str, Any] = None   ,       snake_case__:	List[str] = None   ,       snake_case__:	List[str] = None   ,       snake_case__:	List[Any] = None   ,       snake_case__:	str = None   ,       snake_case__:	Optional[int] = None   ,       snake_case__:	str = True   ,       )  ->		Tuple:
  
  
         """simple docstring"""
  
  
  
         if attention_mask is not None and attention_mask.ndim == 2:
             # assume that mask is expressed as:
             #   (1 = keep,      0 = discard)
             # convert mask into a bias that can be added to attention scores:
             #       (keep = +0,     discard = -10000.0)
             _lowercase     =	(1 - attention_mask.to(hidden_states.dtype	)) * -10_000.0
             _lowercase     =	attention_mask.unsqueeze(1	)
  
         # convert encoder_attention_mask to a bias the same way we do for attention_mask
         if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
             _lowercase     =	(1 - encoder_attention_mask.to(hidden_states.dtype	)) * -10_000.0
             _lowercase     =	encoder_attention_mask.unsqueeze(1	)
  
         # 1. Input
         if self.is_input_continuous:
             _lowercase     =	hidden_states.shape
             _lowercase     =	hidden_states
  
             _lowercase     =	self.norm(snake_case_	)
             if not self.use_linear_projection:
                 _lowercase     =	self.proj_in(snake_case_	)
                 _lowercase     =	hidden_states.shape[1]
                 _lowercase     =	hidden_states.permute(0   ,       2   ,       3   ,       1	).reshape(snake_case_   ,       height * width   ,       snake_case_	)
             else:
                 _lowercase     =	hidden_states.shape[1]
                 _lowercase     =	hidden_states.permute(0   ,       2   ,       3   ,       1	).reshape(snake_case_   ,       height * width   ,       snake_case_	)
                 _lowercase     =	self.proj_in(snake_case_	)
         elif self.is_input_vectorized:
             _lowercase     =	self.latent_image_embedding(snake_case_	)
         elif self.is_input_patches:
             _lowercase     =	self.pos_embed(snake_case_	)
  
         # 2. Blocks
         for block in self.transformer_blocks:
             _lowercase     =	block(
                 snake_case_   ,       attention_mask=snake_case_   ,       encoder_hidden_states=snake_case_   ,       encoder_attention_mask=snake_case_   ,       timestep=snake_case_   ,       cross_attention_kwargs=snake_case_   ,       class_labels=snake_case_   ,       )
  
         # 3. Output
         if self.is_input_continuous:
             if not self.use_linear_projection:
                 _lowercase     =	hidden_states.reshape(snake_case_   ,       snake_case_   ,       snake_case_   ,       snake_case_	).permute(0   ,       3   ,       1   ,       2	).contiguous()
                 _lowercase     =	self.proj_out(snake_case_	)
             else:
                 _lowercase     =	self.proj_out(snake_case_	)
                 _lowercase     =	hidden_states.reshape(snake_case_   ,       snake_case_   ,       snake_case_   ,       snake_case_	).permute(0   ,       3   ,       1   ,       2	).contiguous()
  
             _lowercase     =	hidden_states + residual
         elif self.is_input_vectorized:
             _lowercase     =	self.norm_out(snake_case_	)
             _lowercase     =	self.out(snake_case_	)
             # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
             _lowercase     =	logits.permute(0   ,       2   ,       1	)
  
             # log(p(x_0))
             _lowercase     =	F.log_softmax(logits.double()   ,       dim=1	).float()
         elif self.is_input_patches:
             # TODO: cleanup!
             _lowercase     =	self.transformer_blocks[0].norma.emb(
                 snake_case_   ,       snake_case_   ,       hidden_dtype=hidden_states.dtype	)
             _lowercase     =	self.proj_out_a(F.silu(snake_case_	)	).chunk(2   ,       dim=1	)
             _lowercase     =	self.norm_out(snake_case_	) * (1 + scale[:, None]) + shift[:, None]
             _lowercase     =	self.proj_out_a(snake_case_	)
  
             # unpatchify
             _lowercase     =	int(hidden_states.shape[1] ** 0.5	)
             _lowercase     =	hidden_states.reshape(
                 shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)	)
             _lowercase     =	torch.einsum('nhwpqc->nchpwq'   ,       snake_case_	)
             _lowercase     =	hidden_states.reshape(
                 shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)	)
  
         if not return_dict:
             return (output,)
  
         return TransformeraDModelOutput(sample=snake_case_	)
  
style_context: 
  
  
  
  
  
  
  from __future__ import annotations
  def        UpperCAmelCase      (   _lowercase,      _lowercase      ):
  
  
      '''simple docstring'''
  
  
  
      _lowercase     =	[None] * (len(_lowercase      ) + len(_lowercase      ))
  
      _lowercase     =	0  # pointer to current index for left Half
      _lowercase     =	0  # pointer to current index for the right Half
      _lowercase     =	0  # pointer to current index for the sorted array Half
  
      while pointera < len(_lowercase      ) and pointera < len(_lowercase      ):
          if left_half[pointera] < right_half[pointera]:
              _lowercase     =	left_half[pointera]
              pointera += 1
              index += 1
          else:
              _lowercase     =	right_half[pointera]
              pointera += 1
              index += 1
      while pointera < len(_lowercase      ):
          _lowercase     =	left_half[pointera]
          pointera += 1
          index += 1
  
      while pointera < len(_lowercase      ):
          _lowercase     =	right_half[pointera]
          pointera += 1
          index += 1
  
      return sorted_array
  def        UpperCAmelCase      (   _lowercase      ):
  
  
      '''simple docstring'''
  
  
  
      if len(_lowercase      ) <= 1:
          return array
      # the actual formula to calculate the middle element = left + (right - left) // 2
      # this avoids integer overflow in case of large N
      _lowercase     =	0 + (len(_lowercase      ) - 0) // 2
  
      # Split the array into halves till the array length becomes equal to One
      # merge the arrays of single length returned by mergeSort function and
      # pass them into the merge arrays function which merges the array
      _lowercase     =	array[:middle]
      _lowercase     =	array[middle:]
  
      return merge(merge_sort(_lowercase      ),      merge_sort(_lowercase      )      )
  
  
  if __name__ == "__main__":
     import doctest
  
     doctest.testmod()
  