



"""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      _lowercase (     nn.Module						):


      def __init__(    self				:	int							,   a				:	int							,   a				:	int							,   a				:	int							,   a				:	Tuple=0.0							,   a				:	Optional[int] = None							,   a				:	str = "geglu"							,   a				:	Optional[int] = None							,   a				:	bool = False							,   a				:	bool = False							,   a				:	bool = False							,   a				:	bool = False							,   a				:	bool = True							,   a				:	str = "layer_norm"							,   a				:	bool = False							,   ):





            """simple docstring"""







            super().__init__()
            __snake_case      :    str  			=only_cross_attention

            __snake_case      :    List[Any]  			=(num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
            __snake_case      :    Tuple  			=(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:
                  __snake_case      :    Dict  			=AdaLayerNorm(a							,   a						)
            elif self.use_ada_layer_norm_zero:
                  __snake_case      :    Union[str, Any]  			=AdaLayerNormZero(a							,   a						)
            else:
                  __snake_case      :    Any  			=nn.LayerNorm(a							,   elementwise_affine=a						)
            __snake_case      :    Union[str, Any]  			=Attention(
                query_dim=a							,   heads=a							,   dim_head=a							,   dropout=a							,   bias=a							,   cross_attention_dim=cross_attention_dim if only_cross_attention else None							,   upcast_attention=a							,   )

            # 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.
                  __snake_case      :    Dict  			=(
                      AdaLayerNorm(a							,   a						)
                      if self.use_ada_layer_norm
                      else nn.LayerNorm(a							,   elementwise_affine=a						)
                  )
                  __snake_case      :    Optional[Any]  			=Attention(
                      query_dim=a							,   cross_attention_dim=cross_attention_dim if not double_self_attention else None							,   heads=a							,   dim_head=a							,   dropout=a							,   bias=a							,   upcast_attention=a							,   )  # is self-attn if encoder_hidden_states is none
            else:
                  __snake_case      :    Tuple  			=None
                  __snake_case      :    List[Any]  			=None

            # 3. Feed-forward
            __snake_case      :    str  			=nn.LayerNorm(a							,   elementwise_affine=a						)
            __snake_case      :    int  			=FeedForward(a							,   dropout=a							,   activation_fn=a							,   final_dropout=a						)

            # let chunk size default to None
            __snake_case      :    str  			=None
            __snake_case      :    Optional[int]  			=0


      def _UpperCamelCase						(    self				:	str							,   a				:	Optional[int]							,   a				:	int						):





            """simple docstring"""







            __snake_case      :    List[Any]  			=chunk_size
            __snake_case      :    Tuple  			=dim







      def _UpperCamelCase						(    self				:	List[Any]							,   a				:	torch.FloatTensor							,   a				:	Optional[torch.FloatTensor] = None							,   a				:	Optional[torch.FloatTensor] = None							,   a				:	Optional[torch.FloatTensor] = None							,   a				:	Optional[torch.LongTensor] = None							,   a				:	Dict[str, Any] = None							,   a				:	Optional[torch.LongTensor] = None							,   ):





            """simple docstring"""







            if self.use_ada_layer_norm:
                  __snake_case      :    int  			=self.norma(a							,   a						)
            elif self.use_ada_layer_norm_zero:
                  __snake_case      :    List[str]  			=self.norma(
                      a							,   a							,   a							,   hidden_dtype=hidden_states.dtype						)
            else:
                  __snake_case      :    List[str]  			=self.norma(a						)

            __snake_case      :    int  			=cross_attention_kwargs if cross_attention_kwargs is not None else {}

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

            # 2. Cross-Attention
            if self.attna is not None:
                  __snake_case      :    Tuple  			=(
                      self.norma(a							,   a						) if self.use_ada_layer_norm else self.norma(a						)
                  )

                  __snake_case      :    Optional[int]  			=self.attna(
                      a							,   encoder_hidden_states=a							,   attention_mask=a							,   **a							,   )
                  __snake_case      :    Any  			=attn_output + hidden_states

            # 3. Feed-forward
            __snake_case      :    int  			=self.norma(a						)

            if self.use_ada_layer_norm_zero:
                  __snake_case      :    Any  			=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`.'''						)

                  __snake_case      :    Dict  			=norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
                  __snake_case      :    Optional[int]  			=torch.cat(
                      [self.ff(a						) for hid_slice in norm_hidden_states.chunk(a							,   dim=self._chunk_dim						)]							,   dim=self._chunk_dim							,   )
            else:
                  __snake_case      :    str  			=self.ff(a						)

            if self.use_ada_layer_norm_zero:
                  __snake_case      :    Tuple  			=gate_mlp.unsqueeze(1						) * ff_output

            __snake_case      :    List[Any]  			=ff_output + hidden_states

            return hidden_states



class      _lowercase (     nn.Module						):


      def __init__(    self				:	List[Any]							,   a				:	int							,   a				:	Optional[int] = None							,   a				:	int = 4							,   a				:	float = 0.0							,   a				:	str = "geglu"							,   a				:	bool = False							,   ):





            """simple docstring"""







            super().__init__()
            __snake_case      :    Any  			=int(dim * mult						)
            __snake_case      :    int  			=dim_out if dim_out is not None else dim

            if activation_fn == "gelu":
                  __snake_case      :    Optional[Any]  			=GELU(a							,   a						)
            if activation_fn == "gelu-approximate":
                  __snake_case      :    Any  			=GELU(a							,   a							,   approximate='''tanh'''						)
            elif activation_fn == "geglu":
                  __snake_case      :    Optional[Any]  			=GEGLU(a							,   a						)
            elif activation_fn == "geglu-approximate":
                  __snake_case      :    str  			=ApproximateGELU(a							,   a						)

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







      def _UpperCamelCase						(    self				:	List[str]							,   a				:	Optional[Any]						):





            """simple docstring"""







            for module in self.net:
                  __snake_case      :    Dict  			=module(a						)
            return hidden_states



class      _lowercase (     nn.Module						):


      def __init__(    self				:	Any							,   a				:	int							,   a				:	int							,   a				:	str = "none"						):





            """simple docstring"""







            super().__init__()
            __snake_case      :    Any  			=nn.Linear(a							,   a						)
            __snake_case      :    Dict  			=approximate


      def _UpperCamelCase						(    self				:	Optional[int]							,   a				:	Dict						):





            """simple docstring"""







            if gate.device.type != "mps":
                  return F.gelu(a							,   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 _UpperCamelCase						(    self				:	Optional[Any]							,   a				:	Any						):





            """simple docstring"""







            __snake_case      :    Optional[int]  			=self.proj(a						)
            __snake_case      :    int  			=self.gelu(a						)
            return hidden_states



class      _lowercase (     nn.Module						):


      def __init__(    self				:	Union[str, Any]							,   a				:	int							,   a				:	int						):





            """simple docstring"""







            super().__init__()
            __snake_case      :    List[str]  			=nn.Linear(a							,   dim_out * 2						)


      def _UpperCamelCase						(    self				:	Optional[Any]							,   a				:	Dict						):





            """simple docstring"""







            if gate.device.type != "mps":
                  return F.gelu(a						)
            # mps: gelu is not implemented for float16
            return F.gelu(gate.to(dtype=torch.floataa						)						).to(dtype=gate.dtype						)







      def _UpperCamelCase						(    self				:	List[Any]							,   a				:	Optional[Any]						):





            """simple docstring"""







            __snake_case      :    Optional[int]  			=self.proj(a						).chunk(2							,   dim=-1						)
            return hidden_states * self.gelu(a						)



class      _lowercase (     nn.Module						):


      def __init__(    self				:	Union[str, Any]							,   a				:	int							,   a				:	int						):





            """simple docstring"""







            super().__init__()
            __snake_case      :    str  			=nn.Linear(a							,   a						)







      def _UpperCamelCase						(    self				:	Optional[Any]							,   a				:	List[str]						):





            """simple docstring"""







            __snake_case      :    Union[str, Any]  			=self.proj(a						)
            return x * torch.sigmoid(1.7_0_2 * x						)



class      _lowercase (     nn.Module						):


      def __init__(    self				:	Tuple							,   a				:	str							,   a				:	Any						):





            """simple docstring"""







            super().__init__()
            __snake_case      :    int  			=nn.Embedding(a							,   a						)
            __snake_case      :    List[Any]  			=nn.SiLU()
            __snake_case      :    Any  			=nn.Linear(a							,   embedding_dim * 2						)
            __snake_case      :    Any  			=nn.LayerNorm(a							,   elementwise_affine=a						)







      def _UpperCamelCase						(    self				:	Union[str, Any]							,   a				:	Optional[int]							,   a				:	str						):





            """simple docstring"""







            __snake_case      :    str  			=self.linear(self.silu(self.emb(a						)						)						)
            __snake_case      :    List[Any]  			=torch.chunk(a							,   2						)
            __snake_case      :    Tuple  			=self.norm(a						) * (1 + scale) + shift
            return x



class      _lowercase (     nn.Module						):


      def __init__(    self				:	Optional[Any]							,   a				:	str							,   a				:	Optional[int]						):





            """simple docstring"""







            super().__init__()

            __snake_case      :    Union[str, Any]  			=CombinedTimestepLabelEmbeddings(a							,   a						)

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







      def _UpperCamelCase						(    self				:	Union[str, Any]							,   a				:	List[Any]							,   a				:	int							,   a				:	Tuple							,   a				:	int=None						):





            """simple docstring"""







            __snake_case      :    Union[str, Any]  			=self.linear(self.silu(self.emb(a							,   a							,   hidden_dtype=a						)						)						)
            __snake_case      :    Union[str, Any]  			=emb.chunk(6							,   dim=1						)
            __snake_case      :    Tuple  			=self.norm(a						) * (1 + scale_msa[:, None]) + shift_msa[:, None]
            return x, gate_msa, shift_mlp, scale_mlp, gate_mlp



class      _lowercase (     nn.Module						):


      def __init__(    self				:	str							,   a				:	int							,   a				:	int							,   a				:	int							,   a				:	Optional[str] = None							,   a				:	float = 1e-5						):





            """simple docstring"""







            super().__init__()
            __snake_case      :    Union[str, Any]  			=num_groups
            __snake_case      :    Dict  			=eps

            if act_fn is None:
                  __snake_case      :    Union[str, Any]  			=None
            else:
                  __snake_case      :    str  			=get_activation(a						)

            __snake_case      :    Any  			=nn.Linear(a							,   out_dim * 2						)







      def _UpperCamelCase						(    self				:	str							,   a				:	Optional[Any]							,   a				:	List[str]						):





            """simple docstring"""







            if self.act:
                  __snake_case      :    Dict  			=self.act(a						)
            __snake_case      :    List[str]  			=self.linear(a						)
            __snake_case      :    List[Any]  			=emb[:, :, None, None]
            __snake_case      :    Dict  			=emb.chunk(2							,   dim=1						)

            __snake_case      :    Any  			=F.group_norm(a							,   self.num_groups							,   eps=self.eps						)
            __snake_case      :    str  			=x * (1 + scale) + shift
            return x










"""simple docstring"""





from __future__ import annotations

import math
from collections import Counter
from string import ascii_lowercase



def     __lowercase					(    a :				str )					->					None:
      __snake_case			,				__snake_case      :    List[Any]  			=analyze_text(a )
      __snake_case      :    Union[str, Any]  			=list(''' ''' + ascii_lowercase )
      # what is our total sum of probabilities.
      __snake_case      :    List[Any]  			=sum(single_char_strings.values() )

      # one length string
      __snake_case      :    str  			=0
      # for each alpha we go in our dict and if it is in it we calculate entropy
      for ch in my_alphas:
            if ch in single_char_strings:
                  __snake_case      :    Union[str, Any]  			=single_char_strings[ch]
                  __snake_case      :    Any  			=my_str / all_sum
                  my_fir_sum += prob * math.loga(a )  # entropy formula.

    # print entropy
      print(f'''{round(-1 * my_fir_sum ):.1f}''' )

      # two len string
      __snake_case      :    Optional[Any]  			=sum(two_char_strings.values() )
      __snake_case      :    Dict  			=0
      # for each alpha (two in size) calculate entropy.
      for cha in my_alphas:
            for cha in my_alphas:
                  __snake_case      :    Tuple  			=cha + cha
                  if sequence in two_char_strings:
                        __snake_case      :    List[Any]  			=two_char_strings[sequence]
                        __snake_case      :    Tuple  			=int(a ) / all_sum
                        my_sec_sum += prob * math.loga(a )

    # print second entropy
      print(f'''{round(-1 * my_sec_sum ):.1f}''' )

      # print the difference between them
      print(f'''{round((-1 * my_sec_sum) - (-1 * my_fir_sum) ):.1f}''' )



def     __lowercase					(    a :				str )					->					tuple[dict, dict]:
      __snake_case      :    List[str]  			=Counter()  # type: ignore
      __snake_case      :    Union[str, Any]  			=Counter()  # type: ignore
      single_char_strings[text[-1]] += 1

      # first case when we have space at start.
      two_char_strings[" " + text[0]] += 1
      for i in range(0       ,		len(a ) - 1 ):
            single_char_strings[text[i]] += 1
            two_char_strings[text[i : i + 2]] += 1
      return single_char_strings, two_char_strings



def     __lowercase					(    )					->					int:
      import doctest

      doctest.testmod()
      # text = (
      #     "Had repulsive dashwoods suspicion sincerity but advantage now him. Remark "
      #     "easily garret nor nay. Civil those mrs enjoy shy fat merry. You greatest "
      #     "jointure saw horrible. He private he on be imagine suppose. Fertile "
      #     "beloved evident through no service elderly is. Blind there if every no so "
      #     "at. Own neglected you preferred way sincerity delivered his attempted. To "
      #     "of message cottage windows do besides against uncivil.  Delightful "
      #     "unreserved impossible few estimating men favourable see entreaties. She "
      #     "propriety immediate was improving. He or entrance humoured likewise "
      #     "moderate. Much nor game son say feel. Fat make met can must form into "
      #     "gate. Me we offending prevailed discovery. "
      # )

      # calculate_prob(text)


if __name__ == "__main__":
    main()


