


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       __lowerCAmelCase       (	a   ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE						      =    42






class       __lowerCAmelCase       (	a	,					a   ):






  """simple docstring"""





  @register_to_config
  def __init__(       self    :    List[Any]      ,			_lowerCAmelCase    :    int = 1_6      ,			_lowerCAmelCase    :    int = 8_8      ,			_lowerCAmelCase    :    Optional[int] = None      ,			_lowerCAmelCase    :    Optional[int] = None      ,			_lowerCAmelCase    :    int = 1      ,			_lowerCAmelCase    :    float = 0.0      ,			_lowerCAmelCase    :    int = 3_2      ,			_lowerCAmelCase    :    Optional[int] = None      ,			_lowerCAmelCase    :    bool = False      ,			_lowerCAmelCase    :    Optional[int] = None      ,			_lowerCAmelCase    :    Optional[int] = None      ,			_lowerCAmelCase    :    Optional[int] = None      ,			_lowerCAmelCase    :    str = "geglu"      ,			_lowerCAmelCase    :    Optional[int] = None      ,			_lowerCAmelCase    :    bool = False      ,			_lowerCAmelCase    :    bool = False      ,			_lowerCAmelCase    :    bool = False      ,			_lowerCAmelCase    :    str = "layer_norm"      ,			_lowerCAmelCase    :    bool = True      ,			)				->	Any:




     """simple docstring"""
     super().__init__()
     snake_case_         =   use_linear_projection
     snake_case_         =   num_attention_heads
     snake_case_         =   attention_head_dim
     snake_case_         =   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
     snake_case_         =   (in_channels is not None) and (patch_size is None)
     snake_case_         =   num_vector_embeds is not None
     snake_case_         =   in_channels is not None and patch_size is not None

     if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
        snake_case_         =   (
            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"      ,			_lowerCAmelCase      ,			standard_warn=_lowerCAmelCase					)
        snake_case_         =   "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:
        snake_case_         =   in_channels

        snake_case_         =   torch.nn.GroupNorm(num_groups=_lowerCAmelCase      ,			num_channels=_lowerCAmelCase      ,			eps=1e-6      ,			affine=_lowerCAmelCase					)
        if use_linear_projection:
           snake_case_         =   nn.Linear(_lowerCAmelCase      ,			_lowerCAmelCase					)
        else:
           snake_case_         =   nn.Convad(_lowerCAmelCase      ,			_lowerCAmelCase      ,			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"

        snake_case_         =   sample_size
        snake_case_         =   sample_size
        snake_case_         =   num_vector_embeds
        snake_case_         =   self.height * self.width

        snake_case_         =   ImagePositionalEmbeddings(
            num_embed=_lowerCAmelCase      ,			embed_dim=_lowerCAmelCase      ,			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"

        snake_case_         =   sample_size
        snake_case_         =   sample_size

        snake_case_         =   patch_size
        snake_case_         =   PatchEmbed(
            height=_lowerCAmelCase      ,			width=_lowerCAmelCase      ,			patch_size=_lowerCAmelCase      ,			in_channels=_lowerCAmelCase      ,			embed_dim=_lowerCAmelCase      ,			)

     # 3. Define transformers blocks
     snake_case_         =   nn.ModuleList(
         [
             BasicTransformerBlock(
                 _lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			dropout=_lowerCAmelCase      ,			cross_attention_dim=_lowerCAmelCase      ,			activation_fn=_lowerCAmelCase      ,			num_embeds_ada_norm=_lowerCAmelCase      ,			attention_bias=_lowerCAmelCase      ,			only_cross_attention=_lowerCAmelCase      ,			upcast_attention=_lowerCAmelCase      ,			norm_type=_lowerCAmelCase      ,			norm_elementwise_affine=_lowerCAmelCase      ,			)
             for d in range(_lowerCAmelCase					)
         ]					)

     # 4. Define output layers
     snake_case_         =   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:
           snake_case_         =   nn.Linear(_lowerCAmelCase      ,			_lowerCAmelCase					)
        else:
           snake_case_         =   nn.Convad(_lowerCAmelCase      ,			_lowerCAmelCase      ,			kernel_size=1      ,			stride=1      ,			padding=0					)
     elif self.is_input_vectorized:
        snake_case_         =   nn.LayerNorm(_lowerCAmelCase					)
        snake_case_         =   nn.Linear(_lowerCAmelCase      ,			self.num_vector_embeds - 1					)
     elif self.is_input_patches:
        snake_case_         =   nn.LayerNorm(_lowerCAmelCase      ,			elementwise_affine=_lowerCAmelCase      ,			eps=1e-6					)
        snake_case_         =   nn.Linear(_lowerCAmelCase      ,			2 * inner_dim					)
        snake_case_         =   nn.Linear(_lowerCAmelCase      ,			patch_size * patch_size * self.out_channels					)



  def    lowerCAmelCase__					(       self    :    str      ,			_lowerCAmelCase    :    torch.Tensor      ,			_lowerCAmelCase    :    Optional[torch.Tensor] = None      ,			_lowerCAmelCase    :    Optional[torch.LongTensor] = None      ,			_lowerCAmelCase    :    Optional[torch.LongTensor] = None      ,			_lowerCAmelCase    :    Dict[str, Any] = None      ,			_lowerCAmelCase    :    Optional[torch.Tensor] = None      ,			_lowerCAmelCase    :    Optional[torch.Tensor] = None      ,			_lowerCAmelCase    :    bool = True      ,			)				->	str:




     """simple docstring"""
     # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
     #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
     #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
     # expects mask of shape:
     #   [batch, key_tokens]
     # adds singleton query_tokens dimension:
     #   [batch,                    1, key_tokens]
     # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
     #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
     #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
     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)
        snake_case_         =   (1 - attention_mask.to(hidden_states.dtype					)) * -10_000.0
        snake_case_         =   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:
        snake_case_         =   (1 - encoder_attention_mask.to(hidden_states.dtype					)) * -10_000.0
        snake_case_         =   encoder_attention_mask.unsqueeze(1					)

     # 1. Input
     if self.is_input_continuous:
        snake_case_      ,      snake_case_      ,      snake_case_      ,      snake_case_         =   hidden_states.shape
        snake_case_         =   hidden_states

        snake_case_         =   self.norm(_lowerCAmelCase					)
        if not self.use_linear_projection:
           snake_case_         =   self.proj_in(_lowerCAmelCase					)
           snake_case_         =   hidden_states.shape[1]
           snake_case_         =   hidden_states.permute(0      ,			2      ,			3      ,			1					).reshape(_lowerCAmelCase      ,			height * width      ,			_lowerCAmelCase					)
        else:
           snake_case_         =   hidden_states.shape[1]
           snake_case_         =   hidden_states.permute(0      ,			2      ,			3      ,			1					).reshape(_lowerCAmelCase      ,			height * width      ,			_lowerCAmelCase					)
           snake_case_         =   self.proj_in(_lowerCAmelCase					)
     elif self.is_input_vectorized:
        snake_case_         =   self.latent_image_embedding(_lowerCAmelCase					)
     elif self.is_input_patches:
        snake_case_         =   self.pos_embed(_lowerCAmelCase					)

     # 2. Blocks
     for block in self.transformer_blocks:
        snake_case_         =   block(
            _lowerCAmelCase      ,			attention_mask=_lowerCAmelCase      ,			encoder_hidden_states=_lowerCAmelCase      ,			encoder_attention_mask=_lowerCAmelCase      ,			timestep=_lowerCAmelCase      ,			cross_attention_kwargs=_lowerCAmelCase      ,			class_labels=_lowerCAmelCase      ,			)

     # 3. Output
     if self.is_input_continuous:
        if not self.use_linear_projection:
           snake_case_         =   hidden_states.reshape(_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					).permute(0      ,			3      ,			1      ,			2					).contiguous()
           snake_case_         =   self.proj_out(_lowerCAmelCase					)
        else:
           snake_case_         =   self.proj_out(_lowerCAmelCase					)
           snake_case_         =   hidden_states.reshape(_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					).permute(0      ,			3      ,			1      ,			2					).contiguous()

        snake_case_         =   hidden_states + residual
     elif self.is_input_vectorized:
        snake_case_         =   self.norm_out(_lowerCAmelCase					)
        snake_case_         =   self.out(_lowerCAmelCase					)
        # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
        snake_case_         =   logits.permute(0      ,			2      ,			1					)

        # log(p(x_0))
        snake_case_         =   F.log_softmax(logits.double()      ,			dim=1					).float()
     elif self.is_input_patches:
        # TODO: cleanup!
        snake_case_         =   self.transformer_blocks[0].norma.emb(
            _lowerCAmelCase      ,			_lowerCAmelCase      ,			hidden_dtype=hidden_states.dtype					)
        snake_case_      ,      snake_case_         =   self.proj_out_a(F.silu(_lowerCAmelCase					)					).chunk(2      ,			dim=1					)
        snake_case_         =   self.norm_out(_lowerCAmelCase					) * (1 + scale[:, None]) + shift[:, None]
        snake_case_         =   self.proj_out_a(_lowerCAmelCase					)

        # unpatchify
        snake_case_         =   snake_case_         =   int(hidden_states.shape[1] ** 0.5					)
        snake_case_         =   hidden_states.reshape(
            shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)					)
        snake_case_         =   torch.einsum("nhwpqc->nchpwq"      ,			_lowerCAmelCase					)
        snake_case_         =   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=_lowerCAmelCase					)








from typing import Any

import numpy as np



def      _lowerCAmelCase  (							lowerCAmelCase_							:np.ndarray							)->bool:
   '''simple docstring'''






   assert np.shape(lowerCAmelCase_							)[0] == np.shape(lowerCAmelCase_							)[1]

   # If matrix not symmetric, exit right away.
   if np.allclose(lowerCAmelCase_						,  matrix.T							) is False:
      return False

   # Get eigenvalues and eignevectors for a symmetric matrix.
   snake_case_      ,      snake_case_         =   np.linalg.eigh(lowerCAmelCase_							)

   # Check sign of all eigenvalues.
   # np.all returns a value of type np.bool_
   return bool(np.all(eigen_values > 0							)							)



def      _lowerCAmelCase  (							lowerCAmelCase_							:int							)->Any:
   '''simple docstring'''






   snake_case_         =   np.random.randn(lowerCAmelCase_						,  lowerCAmelCase_							)
   snake_case_         =   np.dot(lowerCAmelCase_						,  random_matrix.T							)
   assert _is_matrix_spd(lowerCAmelCase_							)
   return spd_matrix



def      _lowerCAmelCase  (							lowerCAmelCase_							:np.ndarray						,  lowerCAmelCase_							:np.ndarray						,  lowerCAmelCase_							:int = 1_000						,  lowerCAmelCase_							:float = 1e-8						,  )->Any:
   '''simple docstring'''






   assert np.shape(lowerCAmelCase_							)[0] == np.shape(lowerCAmelCase_							)[1]
   assert np.shape(lowerCAmelCase_							)[0] == np.shape(lowerCAmelCase_							)[0]
   assert _is_matrix_spd(lowerCAmelCase_							)

   # Initialize solution guess, residual, search direction.
   snake_case_         =   np.zeros((np.shape(lowerCAmelCase_							)[0], 1)							)
   snake_case_         =   np.copy(lowerCAmelCase_							)
   snake_case_         =   np.copy(lowerCAmelCase_							)

   # Set initial errors in solution guess and residual.
   snake_case_         =   1e9
   snake_case_         =   1e9
   snake_case_         =   1e9

   # Set iteration counter to threshold number of iterations.
   snake_case_         =   0

   while error > tol:
      # Save this value so we only calculate the matrix-vector product once.
      snake_case_         =   np.dot(lowerCAmelCase_						,  lowerCAmelCase_							)

      # The main algorithm.

      # Update search direction magnitude.
      snake_case_         =   np.dot(ra.T						,  lowerCAmelCase_							) / np.dot(pa.T						,  lowerCAmelCase_							)
      # Update solution guess.
      snake_case_         =   xa + alpha * pa
      # Calculate new residual.
      snake_case_         =   ra - alpha * w
      # Calculate new Krylov subspace scale.
      snake_case_         =   np.dot(r.T						,  lowerCAmelCase_							) / np.dot(ra.T						,  lowerCAmelCase_							)
      # Calculate new A conjuage search direction.
      snake_case_         =   r + beta * pa

      # Calculate errors.
      snake_case_         =   np.linalg.norm(r - ra							)
      snake_case_         =   np.linalg.norm(x - xa							)
      snake_case_         =   np.maximum(lowerCAmelCase_						,  lowerCAmelCase_							)

      # Update variables.
      snake_case_         =   np.copy(lowerCAmelCase_							)
      snake_case_         =   np.copy(lowerCAmelCase_							)
      snake_case_         =   np.copy(lowerCAmelCase_							)

      # Update number of iterations.
      iterations += 1
      if iterations > max_iterations:
         break

   return x



def      _lowerCAmelCase  (							)->None:
   '''simple docstring'''






   snake_case_         =   3
   snake_case_         =   _create_spd_matrix(lowerCAmelCase_							)
   snake_case_         =   np.random.randn(lowerCAmelCase_						,  1							)
   snake_case_         =   np.dot(lowerCAmelCase_						,  lowerCAmelCase_							)

   # Numpy solution.
   snake_case_         =   np.linalg.solve(lowerCAmelCase_						,  lowerCAmelCase_							)

   # Our implementation.
   snake_case_         =   conjugate_gradient(lowerCAmelCase_						,  lowerCAmelCase_							)

   # Ensure both solutions are close to x_true (and therefore one another).
   assert np.linalg.norm(x_numpy - x_true							) <= 1e-6
   assert np.linalg.norm(x_conjugate_gradient - x_true							) <= 1e-6


if __name__ == "__main__":
   import doctest

   doctest.testmod()
   test_conjugate_gradient()




