





import functools
import math

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



def _a       (			lowerCamelCase:     str      ,       lowerCamelCase:     Any      ,       lowerCamelCase:     int      ,       lowerCamelCase:     Union[str, Any]      ,       lowerCamelCase:     int = 40_96				)			->		List[str]:







      '''simple docstring'''






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

      @functools.partial(jax.checkpoint      ,       prevent_cse=lowerCamelCase				)
      def summarize_chunk(lowerCamelCase:     Dict      ,       lowerCamelCase:     List[str]      ,       lowerCamelCase:     List[Any]				):
            __A      =       jnp.einsum('''...qhd,...khd->...qhk'''      ,       lowerCamelCase      ,       lowerCamelCase      ,       precision=lowerCamelCase				)

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

            __A      =       jnp.einsum('''...vhf,...qhv->...qhf'''      ,       lowerCamelCase      ,       lowerCamelCase      ,       precision=lowerCamelCase				)
            __A      =       jnp.einsum('''...qhk->...qh'''      ,       lowerCamelCase				)

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

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

      __A     ,		__A     ,		__A      =       jax.lax.map(f=lowerCamelCase      ,       xs=jnp.arange(0      ,       lowerCamelCase      ,       lowerCamelCase				)				)

      __A      =       jnp.max(lowerCamelCase      ,       axis=0      ,       keepdims=lowerCamelCase				)
      __A      =       jnp.exp(chunk_max - global_max				)

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

      __A      =       chunk_values.sum(axis=0				)
      __A      =       jnp.expand_dims(lowerCamelCase      ,       -1				).sum(axis=0				)

      return all_values / all_weights



def _a       (			lowerCamelCase:     str      ,       lowerCamelCase:     Optional[Any]      ,       lowerCamelCase:     Optional[Any]      ,       lowerCamelCase:     Any=jax.lax.Precision.HIGHEST      ,       lowerCamelCase:     int = 10_24      ,       lowerCamelCase:     int = 40_96				)			->		Dict:







      '''simple docstring'''






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

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

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

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

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





class A_      (    nn.Module     ):
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				8
   lowerCAmelCase__       =				64
   lowerCAmelCase__       =				0.0
   lowerCAmelCase__       =				False
   lowerCAmelCase__       =				jnp.floataa




   def _lowerCAmelCase   (self			:int				)->      Optional[Any]:
         __A      =       self.dim_head * self.heads
         __A      =       self.dim_head**-0.5

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

         __A      =       nn.Dense(self.query_dim					,     dtype=self.dtype					,     name='''to_out_0'''				)
         __A      =       nn.Dropout(rate=self.dropout				)




   def _lowerCAmelCase   (self			:List[Any]					,     _UpperCamelCase			:int				)->      Optional[int]:
         __A     ,		__A     ,		__A      =       tensor.shape
         __A      =       self.heads
         __A      =       tensor.reshape(_UpperCamelCase					,     _UpperCamelCase					,     _UpperCamelCase					,     dim // head_size				)
         __A      =       jnp.transpose(_UpperCamelCase					,     (0, 2, 1, 3)				)
         __A      =       tensor.reshape(batch_size * head_size					,     _UpperCamelCase					,     dim // head_size				)
         return tensor




   def _lowerCAmelCase   (self			:List[str]					,     _UpperCamelCase			:Tuple				)->      List[Any]:
         __A     ,		__A     ,		__A      =       tensor.shape
         __A      =       self.heads
         __A      =       tensor.reshape(batch_size // head_size					,     _UpperCamelCase					,     _UpperCamelCase					,     _UpperCamelCase				)
         __A      =       jnp.transpose(_UpperCamelCase					,     (0, 2, 1, 3)				)
         __A      =       tensor.reshape(batch_size // head_size					,     _UpperCamelCase					,     dim * head_size				)
         return tensor






   def __call__(self			:Optional[int]					,     _UpperCamelCase			:List[str]					,     _UpperCamelCase			:str=None					,     _UpperCamelCase			:Tuple=True				)->      List[str]:
         __A      =       hidden_states if context is None else context

         __A      =       self.query(_UpperCamelCase				)
         __A      =       self.key(_UpperCamelCase				)
         __A      =       self.value(_UpperCamelCase				)

         __A      =       self.reshape_heads_to_batch_dim(_UpperCamelCase				)
         __A      =       self.reshape_heads_to_batch_dim(_UpperCamelCase				)
         __A      =       self.reshape_heads_to_batch_dim(_UpperCamelCase				)

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

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

               __A      =       jax_memory_efficient_attention(
                   _UpperCamelCase					,     _UpperCamelCase					,     _UpperCamelCase					,     query_chunk_size=_UpperCamelCase					,     key_chunk_size=4096 * 4				)

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

               # attend to values
               __A      =       jnp.einsum('''b i j, b j d -> b i d'''					,     _UpperCamelCase					,     _UpperCamelCase				)

         __A      =       self.reshape_batch_dim_to_heads(_UpperCamelCase				)
         __A      =       self.proj_attn(_UpperCamelCase				)
         return self.dropout_layer(_UpperCamelCase					,     deterministic=_UpperCamelCase				)





class A_      (    nn.Module     ):
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				0.0
   lowerCAmelCase__       =				False
   lowerCAmelCase__       =				jnp.floataa
   lowerCAmelCase__       =				False




   def _lowerCAmelCase   (self			:List[str]				)->      Optional[Any]:
         # self attention (or cross_attention if only_cross_attention is True)
         __A      =       FlaxAttention(
             self.dim					,     self.n_heads					,     self.d_head					,     self.dropout					,     self.use_memory_efficient_attention					,     dtype=self.dtype				)
         # cross attention
         __A      =       FlaxAttention(
             self.dim					,     self.n_heads					,     self.d_head					,     self.dropout					,     self.use_memory_efficient_attention					,     dtype=self.dtype				)
         __A      =       FlaxFeedForward(dim=self.dim					,     dropout=self.dropout					,     dtype=self.dtype				)
         __A      =       nn.LayerNorm(epsilon=1e-5					,     dtype=self.dtype				)
         __A      =       nn.LayerNorm(epsilon=1e-5					,     dtype=self.dtype				)
         __A      =       nn.LayerNorm(epsilon=1e-5					,     dtype=self.dtype				)
         __A      =       nn.Dropout(rate=self.dropout				)






   def __call__(self			:str					,     _UpperCamelCase			:Optional[Any]					,     _UpperCamelCase			:Union[str, Any]					,     _UpperCamelCase			:Any=True				)->      Tuple:
         # self attention
         __A      =       hidden_states
         if self.only_cross_attention:
               __A      =       self.attna(self.norma(_UpperCamelCase				)					,     _UpperCamelCase					,     deterministic=_UpperCamelCase				)
         else:
               __A      =       self.attna(self.norma(_UpperCamelCase				)					,     deterministic=_UpperCamelCase				)
         __A      =       hidden_states + residual

         # cross attention
         __A      =       hidden_states
         __A      =       self.attna(self.norma(_UpperCamelCase				)					,     _UpperCamelCase					,     deterministic=_UpperCamelCase				)
         __A      =       hidden_states + residual

         # feed forward
         __A      =       hidden_states
         __A      =       self.ff(self.norma(_UpperCamelCase				)					,     deterministic=_UpperCamelCase				)
         __A      =       hidden_states + residual

         return self.dropout_layer(_UpperCamelCase					,     deterministic=_UpperCamelCase				)





class A_      (    nn.Module     ):
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				1
   lowerCAmelCase__       =				0.0
   lowerCAmelCase__       =				False
   lowerCAmelCase__       =				False
   lowerCAmelCase__       =				jnp.floataa
   lowerCAmelCase__       =				False




   def _lowerCAmelCase   (self			:Any				)->      Any:
         __A      =       nn.GroupNorm(num_groups=32					,     epsilon=1e-5				)

         __A      =       self.n_heads * self.d_head
         if self.use_linear_projection:
               __A      =       nn.Dense(_UpperCamelCase					,     dtype=self.dtype				)
         else:
               __A      =       nn.Conv(
                   _UpperCamelCase					,     kernel_size=(1, 1)					,     strides=(1, 1)					,     padding='''VALID'''					,     dtype=self.dtype					,     )

         __A      =       [
             FlaxBasicTransformerBlock(
                 _UpperCamelCase					,     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:
               __A      =       nn.Dense(_UpperCamelCase					,     dtype=self.dtype				)
         else:
               __A      =       nn.Conv(
                   _UpperCamelCase					,     kernel_size=(1, 1)					,     strides=(1, 1)					,     padding='''VALID'''					,     dtype=self.dtype					,     )

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






   def __call__(self			:List[str]					,     _UpperCamelCase			:int					,     _UpperCamelCase			:Optional[Any]					,     _UpperCamelCase			:Tuple=True				)->      List[str]:
         __A     ,		__A     ,		__A     ,		__A      =       hidden_states.shape
         __A      =       hidden_states
         __A      =       self.norm(_UpperCamelCase				)
         if self.use_linear_projection:
               __A      =       hidden_states.reshape(_UpperCamelCase					,     height * width					,     _UpperCamelCase				)
               __A      =       self.proj_in(_UpperCamelCase				)
         else:
               __A      =       self.proj_in(_UpperCamelCase				)
               __A      =       hidden_states.reshape(_UpperCamelCase					,     height * width					,     _UpperCamelCase				)

         for transformer_block in self.transformer_blocks:
               __A      =       transformer_block(_UpperCamelCase					,     _UpperCamelCase					,     deterministic=_UpperCamelCase				)

         if self.use_linear_projection:
               __A      =       self.proj_out(_UpperCamelCase				)
               __A      =       hidden_states.reshape(_UpperCamelCase					,     _UpperCamelCase					,     _UpperCamelCase					,     _UpperCamelCase				)
         else:
               __A      =       hidden_states.reshape(_UpperCamelCase					,     _UpperCamelCase					,     _UpperCamelCase					,     _UpperCamelCase				)
               __A      =       self.proj_out(_UpperCamelCase				)

         __A      =       hidden_states + residual
         return self.dropout_layer(_UpperCamelCase					,     deterministic=_UpperCamelCase				)





class A_      (    nn.Module     ):
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				0.0
   lowerCAmelCase__       =				jnp.floataa




   def _lowerCAmelCase   (self			:Dict				)->      str:
         # The second linear layer needs to be called
         # net_2 for now to match the index of the Sequential layer
         __A      =       FlaxGEGLU(self.dim					,     self.dropout					,     self.dtype				)
         __A      =       nn.Dense(self.dim					,     dtype=self.dtype				)






   def __call__(self			:Any					,     _UpperCamelCase			:int					,     _UpperCamelCase			:int=True				)->      str:
         __A      =       self.net_a(_UpperCamelCase					,     deterministic=_UpperCamelCase				)
         __A      =       self.net_a(_UpperCamelCase				)
         return hidden_states





class A_      (    nn.Module     ):
   lowerCAmelCase__       =				42
   lowerCAmelCase__       =				0.0
   lowerCAmelCase__       =				jnp.floataa




   def _lowerCAmelCase   (self			:List[str]				)->      str:
         __A      =       self.dim * 4
         __A      =       nn.Dense(inner_dim * 2					,     dtype=self.dtype				)
         __A      =       nn.Dropout(rate=self.dropout				)






   def __call__(self			:Any					,     _UpperCamelCase			:Optional[int]					,     _UpperCamelCase			:Tuple=True				)->      Optional[int]:
         __A      =       self.proj(_UpperCamelCase				)
         __A     ,		__A      =       jnp.split(_UpperCamelCase					,     2					,     axis=2				)
         return self.dropout_layer(hidden_linear * nn.gelu(_UpperCamelCase				)					,     deterministic=_UpperCamelCase				)












from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


snake_case__      :							Optional[int]									=  {'configuration_cvt': ['CVT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'CvtConfig']}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       snake_case__      :							Union[str, Any]									=  [
           'CVT_PRETRAINED_MODEL_ARCHIVE_LIST',
           'CvtForImageClassification',
           'CvtModel',
           'CvtPreTrainedModel',
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       snake_case__      :							List[str]									=  [
           'TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST',
           'TFCvtForImageClassification',
           'TFCvtModel',
           'TFCvtPreTrainedModel',
       ]

if TYPE_CHECKING:
       from .configuration_cvt import CVT_PRETRAINED_CONFIG_ARCHIVE_MAP, CvtConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_cvt import (
                  CVT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  CvtForImageClassification,
                  CvtModel,
                  CvtPreTrainedModel,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_cvt import (
                  TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFCvtForImageClassification,
                  TFCvtModel,
                  TFCvtPreTrainedModel,
              )


else:
       import sys

       snake_case__      :							Dict									=  _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





