code: 
  
  from typing import TYPE_CHECKING
  
  from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
  
  
  A_		:			str										=   {'configuration_dinat': ['DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'DinatConfig']}
  
  
  try:
         if not is_torch_available():
                raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
         pass
  else:
         A_		:			str										=   [
             'DINAT_PRETRAINED_MODEL_ARCHIVE_LIST',
             'DinatForImageClassification',
             'DinatModel',
             'DinatPreTrainedModel',
             'DinatBackbone',
         ]
  
  if TYPE_CHECKING:
         from .configuration_dinat import DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP, DinatConfig
  
         try:
                if not is_torch_available():
                       raise OptionalDependencyNotAvailable()
         except OptionalDependencyNotAvailable:
                pass
         else:
                from .modeling_dinat import (
                    DINAT_PRETRAINED_MODEL_ARCHIVE_LIST,
                    DinatBackbone,
                    DinatForImageClassification,
                    DinatModel,
                    DinatPreTrainedModel,
                )
  
  else:
         import sys
  
         A_		:			List[Any]										=   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
  
  
  
  
  
style_context: 
  
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  
  
  def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE = 4_0_9_6      )      ->     Any:
  
  
  
  
  
  
      '''simple docstring'''
      __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			key.shape[-3:]
      __UpperCAmelCase									=			value.shape[-1]
      __UpperCAmelCase									=			min(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )
      __UpperCAmelCase									=			query / jnp.sqrt(SCREAMING_SNAKE_CASE      )
  
      @functools.partial(jax.checkpoint       ,   prevent_cse=SCREAMING_SNAKE_CASE      )
      def summarize_chunk(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      ):
          __UpperCAmelCase									=			jnp.einsum('''...qhd,...khd->...qhk'''       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   precision=SCREAMING_SNAKE_CASE      )
  
          __UpperCAmelCase									=			jnp.max(SCREAMING_SNAKE_CASE       ,   axis=-1       ,   keepdims=SCREAMING_SNAKE_CASE      )
          __UpperCAmelCase									=			jax.lax.stop_gradient(SCREAMING_SNAKE_CASE      )
          __UpperCAmelCase									=			jnp.exp(attn_weights - max_score      )
  
          __UpperCAmelCase									=			jnp.einsum('''...vhf,...qhv->...qhf'''       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   precision=SCREAMING_SNAKE_CASE      )
          __UpperCAmelCase									=			jnp.einsum('''...qhk->...qh'''       ,   SCREAMING_SNAKE_CASE      )
  
          return (exp_values, exp_weights.sum(axis=-1      ), max_score)
  
      def chunk_scanner(SCREAMING_SNAKE_CASE      ):
          # julienne key array
          __UpperCAmelCase									=			jax.lax.dynamic_slice(
              operand=SCREAMING_SNAKE_CASE       ,   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
          __UpperCAmelCase									=			jax.lax.dynamic_slice(
              operand=SCREAMING_SNAKE_CASE       ,   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(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )
  
      __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			jax.lax.map(f=SCREAMING_SNAKE_CASE       ,   xs=jnp.arange(0       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      )      )
  
      __UpperCAmelCase									=			jnp.max(SCREAMING_SNAKE_CASE       ,   axis=0       ,   keepdims=SCREAMING_SNAKE_CASE      )
      __UpperCAmelCase									=			jnp.exp(chunk_max - global_max      )
  
      chunk_values *= jnp.expand_dims(SCREAMING_SNAKE_CASE       ,   axis=-1      )
      chunk_weights *= max_diffs
  
      __UpperCAmelCase									=			chunk_values.sum(axis=0      )
      __UpperCAmelCase									=			jnp.expand_dims(SCREAMING_SNAKE_CASE       ,   -1      ).sum(axis=0      )
  
      return all_values / all_weights
  
  
  def       __a       (   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE=jax.lax.Precision.HIGHEST       ,   SCREAMING_SNAKE_CASE = 1_0_2_4       ,   SCREAMING_SNAKE_CASE = 4_0_9_6      )      ->     int:
  
  
  
  
  
  
      '''simple docstring'''
      __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			query.shape[-3:]
  
      def chunk_scanner(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      ):
          # julienne query array
          __UpperCAmelCase									=			jax.lax.dynamic_slice(
              operand=SCREAMING_SNAKE_CASE       ,   start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]       ,   slice_sizes=list(query.shape[:-3]      ) + [min(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE      ), num_heads, q_features]       ,   )
  
          return (
              chunk_idx + query_chunk_size,  # unused ignore it
              _query_chunk_attention(
                  query=SCREAMING_SNAKE_CASE       ,   key=SCREAMING_SNAKE_CASE       ,   value=SCREAMING_SNAKE_CASE       ,   precision=SCREAMING_SNAKE_CASE       ,   key_chunk_size=SCREAMING_SNAKE_CASE      ),
          )
  
      __UpperCAmelCase					,     __UpperCAmelCase									=			jax.lax.scan(
          f=SCREAMING_SNAKE_CASE       ,   init=0       ,   xs=SCREAMING_SNAKE_CASE       ,   length=math.ceil(num_q / query_chunk_size      )      )
  
      return jnp.concatenate(SCREAMING_SNAKE_CASE       ,   axis=-3      )  # fuse the chunked result back
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       8
      a__     							=       64
      a__     							=       0.0
      a__     							=       False
      a__     							=       jnp.floataa
  
  
      def      lowerCAmelCase_						(self  )     ->      Tuple:
          __UpperCAmelCase									=			self.dim_head * self.heads
          __UpperCAmelCase									=			self.dim_head**-0.5
  
          # Weights were exported with old names {to_q, to_k, to_v, to_out}
          __UpperCAmelCase									=			nn.Dense(lowercase__     ,      use_bias=lowercase__     ,      dtype=self.dtype     ,      name='''to_q'''  )
          __UpperCAmelCase									=			nn.Dense(lowercase__     ,      use_bias=lowercase__     ,      dtype=self.dtype     ,      name='''to_k'''  )
          __UpperCAmelCase									=			nn.Dense(lowercase__     ,      use_bias=lowercase__     ,      dtype=self.dtype     ,      name='''to_v'''  )
  
          __UpperCAmelCase									=			nn.Dense(self.query_dim     ,      dtype=self.dtype     ,      name='''to_out_0'''  )
          __UpperCAmelCase									=			nn.Dropout(rate=self.dropout  )
  
  
      def      lowerCAmelCase_						(self     ,      lowercase__  )     ->      Optional[int]:
          __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			tensor.shape
          __UpperCAmelCase									=			self.heads
          __UpperCAmelCase									=			tensor.reshape(lowercase__     ,      lowercase__     ,      lowercase__     ,      dim // head_size  )
          __UpperCAmelCase									=			jnp.transpose(lowercase__     ,      (0, 2, 1, 3)  )
          __UpperCAmelCase									=			tensor.reshape(batch_size * head_size     ,      lowercase__     ,      dim // head_size  )
          return tensor
  
  
      def      lowerCAmelCase_						(self     ,      lowercase__  )     ->      Any:
          __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			tensor.shape
          __UpperCAmelCase									=			self.heads
          __UpperCAmelCase									=			tensor.reshape(batch_size // head_size     ,      lowercase__     ,      lowercase__     ,      lowercase__  )
          __UpperCAmelCase									=			jnp.transpose(lowercase__     ,      (0, 2, 1, 3)  )
          __UpperCAmelCase									=			tensor.reshape(batch_size // head_size     ,      lowercase__     ,      dim * head_size  )
          return tensor
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__=None     ,      lowercase__=True  )     ->      Optional[Any]:
          __UpperCAmelCase									=			hidden_states if context is None else context
  
          __UpperCAmelCase									=			self.query(lowercase__  )
          __UpperCAmelCase									=			self.key(lowercase__  )
          __UpperCAmelCase									=			self.value(lowercase__  )
  
          __UpperCAmelCase									=			self.reshape_heads_to_batch_dim(lowercase__  )
          __UpperCAmelCase									=			self.reshape_heads_to_batch_dim(lowercase__  )
          __UpperCAmelCase									=			self.reshape_heads_to_batch_dim(lowercase__  )
  
          if self.use_memory_efficient_attention:
              __UpperCAmelCase									=			query_states.transpose(1     ,      0     ,      2  )
              __UpperCAmelCase									=			key_states.transpose(1     ,      0     ,      2  )
              __UpperCAmelCase									=			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
  
              __UpperCAmelCase									=			query_states.shape[-3]
              if flatten_latent_dim % 64 == 0:
                  __UpperCAmelCase									=			int(flatten_latent_dim / 64  )
              elif flatten_latent_dim % 16 == 0:
                  __UpperCAmelCase									=			int(flatten_latent_dim / 16  )
              elif flatten_latent_dim % 4 == 0:
                  __UpperCAmelCase									=			int(flatten_latent_dim / 4  )
              else:
                  __UpperCAmelCase									=			int(lowercase__  )
  
              __UpperCAmelCase									=			jax_memory_efficient_attention(
                  lowercase__     ,      lowercase__     ,      lowercase__     ,      query_chunk_size=lowercase__     ,      key_chunk_size=4_096 * 4  )
  
              __UpperCAmelCase									=			hidden_states.transpose(1     ,      0     ,      2  )
          else:
              # compute attentions
              __UpperCAmelCase									=			jnp.einsum('''b i d, b j d->b i j'''     ,      lowercase__     ,      lowercase__  )
              __UpperCAmelCase									=			attention_scores * self.scale
              __UpperCAmelCase									=			nn.softmax(lowercase__     ,      axis=2  )
  
              # attend to values
              __UpperCAmelCase									=			jnp.einsum('''b i j, b j d -> b i d'''     ,      lowercase__     ,      lowercase__  )
  
          __UpperCAmelCase									=			self.reshape_batch_dim_to_heads(lowercase__  )
          __UpperCAmelCase									=			self.proj_attn(lowercase__  )
          return self.dropout_layer(lowercase__     ,      deterministic=lowercase__  )
  
  
  
  
  
  
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       42
      a__     							=       42
      a__     							=       0.0
      a__     							=       False
      a__     							=       jnp.floataa
      a__     							=       False
  
  
      def      lowerCAmelCase_						(self  )     ->      Union[str, Any]:
          # self attention (or cross_attention if only_cross_attention is True)
          __UpperCAmelCase									=			FlaxAttention(
              self.dim     ,      self.n_heads     ,      self.d_head     ,      self.dropout     ,      self.use_memory_efficient_attention     ,      dtype=self.dtype  )
          # cross attention
          __UpperCAmelCase									=			FlaxAttention(
              self.dim     ,      self.n_heads     ,      self.d_head     ,      self.dropout     ,      self.use_memory_efficient_attention     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			FlaxFeedForward(dim=self.dim     ,      dropout=self.dropout     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.LayerNorm(epsilon=1E-5     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.LayerNorm(epsilon=1E-5     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.LayerNorm(epsilon=1E-5     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.Dropout(rate=self.dropout  )
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__     ,      lowercase__=True  )     ->      Optional[int]:
          # self attention
          __UpperCAmelCase									=			hidden_states
          if self.only_cross_attention:
              __UpperCAmelCase									=			self.attna(self.norma(lowercase__  )     ,      lowercase__     ,      deterministic=lowercase__  )
          else:
              __UpperCAmelCase									=			self.attna(self.norma(lowercase__  )     ,      deterministic=lowercase__  )
          __UpperCAmelCase									=			hidden_states + residual
  
          # cross attention
          __UpperCAmelCase									=			hidden_states
          __UpperCAmelCase									=			self.attna(self.norma(lowercase__  )     ,      lowercase__     ,      deterministic=lowercase__  )
          __UpperCAmelCase									=			hidden_states + residual
  
          # feed forward
          __UpperCAmelCase									=			hidden_states
          __UpperCAmelCase									=			self.ff(self.norma(lowercase__  )     ,      deterministic=lowercase__  )
          __UpperCAmelCase									=			hidden_states + residual
  
          return self.dropout_layer(lowercase__     ,      deterministic=lowercase__  )
  
  
  
  
  
  
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       42
      a__     							=       42
      a__     							=       1
      a__     							=       0.0
      a__     							=       False
      a__     							=       False
      a__     							=       jnp.floataa
      a__     							=       False
  
  
      def      lowerCAmelCase_						(self  )     ->      int:
          __UpperCAmelCase									=			nn.GroupNorm(num_groups=32     ,      epsilon=1E-5  )
  
          __UpperCAmelCase									=			self.n_heads * self.d_head
          if self.use_linear_projection:
              __UpperCAmelCase									=			nn.Dense(lowercase__     ,      dtype=self.dtype  )
          else:
              __UpperCAmelCase									=			nn.Conv(
                  lowercase__     ,      kernel_size=(1, 1)     ,      strides=(1, 1)     ,      padding='''VALID'''     ,      dtype=self.dtype     ,      )
  
          __UpperCAmelCase									=			[
              FlaxBasicTransformerBlock(
                  lowercase__     ,      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:
              __UpperCAmelCase									=			nn.Dense(lowercase__     ,      dtype=self.dtype  )
          else:
              __UpperCAmelCase									=			nn.Conv(
                  lowercase__     ,      kernel_size=(1, 1)     ,      strides=(1, 1)     ,      padding='''VALID'''     ,      dtype=self.dtype     ,      )
  
          __UpperCAmelCase									=			nn.Dropout(rate=self.dropout  )
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__     ,      lowercase__=True  )     ->      Dict:
          __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			hidden_states.shape
          __UpperCAmelCase									=			hidden_states
          __UpperCAmelCase									=			self.norm(lowercase__  )
          if self.use_linear_projection:
              __UpperCAmelCase									=			hidden_states.reshape(lowercase__     ,      height * width     ,      lowercase__  )
              __UpperCAmelCase									=			self.proj_in(lowercase__  )
          else:
              __UpperCAmelCase									=			self.proj_in(lowercase__  )
              __UpperCAmelCase									=			hidden_states.reshape(lowercase__     ,      height * width     ,      lowercase__  )
  
          for transformer_block in self.transformer_blocks:
              __UpperCAmelCase									=			transformer_block(lowercase__     ,      lowercase__     ,      deterministic=lowercase__  )
  
          if self.use_linear_projection:
              __UpperCAmelCase									=			self.proj_out(lowercase__  )
              __UpperCAmelCase									=			hidden_states.reshape(lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__  )
          else:
              __UpperCAmelCase									=			hidden_states.reshape(lowercase__     ,      lowercase__     ,      lowercase__     ,      lowercase__  )
              __UpperCAmelCase									=			self.proj_out(lowercase__  )
  
          __UpperCAmelCase									=			hidden_states + residual
          return self.dropout_layer(lowercase__     ,      deterministic=lowercase__  )
  
  
  
  
  
  
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       0.0
      a__     							=       jnp.floataa
  
  
      def      lowerCAmelCase_						(self  )     ->      int:
          # The second linear layer needs to be called
          # net_2 for now to match the index of the Sequential layer
          __UpperCAmelCase									=			FlaxGEGLU(self.dim     ,      self.dropout     ,      self.dtype  )
          __UpperCAmelCase									=			nn.Dense(self.dim     ,      dtype=self.dtype  )
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__=True  )     ->      Dict:
          __UpperCAmelCase									=			self.net_a(lowercase__     ,      deterministic=lowercase__  )
          __UpperCAmelCase									=			self.net_a(lowercase__  )
          return hidden_states
  
  
  
  
  
  
  class 					A_   (      nn.Module						):
  
  
      '''simple docstring'''
  
  
  
  
  
  
      a__     							=       42
      a__     							=       0.0
      a__     							=       jnp.floataa
  
  
      def      lowerCAmelCase_						(self  )     ->      Optional[int]:
          __UpperCAmelCase									=			self.dim * 4
          __UpperCAmelCase									=			nn.Dense(inner_dim * 2     ,      dtype=self.dtype  )
          __UpperCAmelCase									=			nn.Dropout(rate=self.dropout  )
  
  
  
  
  
  
      def __call__(self     ,      lowercase__     ,      lowercase__=True  )     ->      Union[str, Any]:
          __UpperCAmelCase									=			self.proj(lowercase__  )
          __UpperCAmelCase					,     __UpperCAmelCase									=			jnp.split(lowercase__     ,      2     ,      axis=2  )
          return self.dropout_layer(hidden_linear * nn.gelu(lowercase__  )     ,      deterministic=lowercase__  )
  
  
  
  
  