"""simple docstring"""


import functools
import math

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


def   __UpperCAmelCase     (	_snake_case		:       Optional[int],				_snake_case		:       Tuple,				_snake_case		:       Union[str, Any],				_snake_case		:       Dict,				_snake_case		:       int = 4_0_9_6       ):
   _lowercase     ,  _lowercase     ,  _lowercase							     =       key.shape[-3:]
   _lowercase							     =       value.shape[-1]
   _lowercase							     =       min(_snake_case,				_snake_case       )
   _lowercase							     =       query / jnp.sqrt(_snake_case       )

   @functools.partial(jax.checkpoint,				prevent_cse=_snake_case       )
   def summarize_chunk(_snake_case		:       Dict,				_snake_case		:       Tuple,				_snake_case		:       List[str]       ):
      _lowercase							     =       jnp.einsum("...qhd,...khd->...qhk",				_snake_case,				_snake_case,				precision=_snake_case       )

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

      _lowercase							     =       jnp.einsum("...vhf,...qhv->...qhf",				_snake_case,				_snake_case,				precision=_snake_case       )
      _lowercase							     =       jnp.einsum("...qhk->...qh",				_snake_case       )

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

   def chunk_scanner(_snake_case		:       Any       ):
      # julienne key array
      _lowercase							     =       jax.lax.dynamic_slice(
          operand=_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
      _lowercase							     =       jax.lax.dynamic_slice(
          operand=_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(_snake_case,				_snake_case,				_snake_case       )

   _lowercase     ,  _lowercase     ,  _lowercase							     =       jax.lax.map(f=_snake_case,				xs=jnp.arange(0,				_snake_case,				_snake_case       )       )

   _lowercase							     =       jnp.max(_snake_case,				axis=0,				keepdims=_snake_case       )
   _lowercase							     =       jnp.exp(chunk_max - global_max       )

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

   _lowercase							     =       chunk_values.sum(axis=0       )
   _lowercase							     =       jnp.expand_dims(_snake_case,				-1       ).sum(axis=0       )

   return all_values / all_weights


def   __UpperCAmelCase     (	_snake_case		:       Optional[Any],				_snake_case		:       Optional[int],				_snake_case		:       Tuple,				_snake_case		:       Tuple=jax.lax.Precision.HIGHEST,				_snake_case		:       int = 1_0_2_4,				_snake_case		:       int = 4_0_9_6       ):
   _lowercase     ,  _lowercase     ,  _lowercase							     =       query.shape[-3:]

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

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

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

   return jnp.concatenate(_snake_case,				axis=-3       )  # fuse the chunked result back
class   UpperCAmelCase_	(				nn.Module   ):
    snake_case_											=      42
    snake_case_											=      8
    snake_case_											=      64
    snake_case_											=      0.0
    snake_case_											=      False
    snake_case_											=      jnp.floataa
    def 		_lowerCamelCase      (	self			:       Optional[int]		)						->    List[Any]:
       _lowercase							     =       self.dim_head * self.heads
       _lowercase							     =       self.dim_head**-0.5

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

       _lowercase							     =       nn.Dense(self.query_dim   ,						dtype=self.dtype   ,						name="to_out_0"		)
       _lowercase							     =       nn.Dropout(rate=self.dropout		)
    def 		_lowerCamelCase      (	self			:       Union[str, Any]   ,						_lowercase			:       List[Any]		)						->    Optional[Any]:
       _lowercase     ,  _lowercase     ,  _lowercase							     =       tensor.shape
       _lowercase							     =       self.heads
       _lowercase							     =       tensor.reshape(_lowercase   ,						_lowercase   ,						_lowercase   ,						dim // head_size		)
       _lowercase							     =       jnp.transpose(_lowercase   ,						(0, 2, 1, 3)		)
       _lowercase							     =       tensor.reshape(batch_size * head_size   ,						_lowercase   ,						dim // head_size		)
       return tensor
    def 		_lowerCamelCase      (	self			:       str   ,						_lowercase			:       Any		)						->    List[Any]:
       _lowercase     ,  _lowercase     ,  _lowercase							     =       tensor.shape
       _lowercase							     =       self.heads
       _lowercase							     =       tensor.reshape(batch_size // head_size   ,						_lowercase   ,						_lowercase   ,						_lowercase		)
       _lowercase							     =       jnp.transpose(_lowercase   ,						(0, 2, 1, 3)		)
       _lowercase							     =       tensor.reshape(batch_size // head_size   ,						_lowercase   ,						dim * head_size		)
       return tensor




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

       _lowercase							     =       self.query(_lowercase		)
       _lowercase							     =       self.key(_lowercase		)
       _lowercase							     =       self.value(_lowercase		)

       _lowercase							     =       self.reshape_heads_to_batch_dim(_lowercase		)
       _lowercase							     =       self.reshape_heads_to_batch_dim(_lowercase		)
       _lowercase							     =       self.reshape_heads_to_batch_dim(_lowercase		)

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

          _lowercase							     =       query_states.shape[-3]
          if flatten_latent_dim % 6_4 == 0:
             _lowercase							     =       int(flatten_latent_dim / 6_4		)
          elif flatten_latent_dim % 1_6 == 0:
             _lowercase							     =       int(flatten_latent_dim / 1_6		)
          elif flatten_latent_dim % 4 == 0:
             _lowercase							     =       int(flatten_latent_dim / 4		)
          else:
             _lowercase							     =       int(_lowercase		)

          _lowercase							     =       jax_memory_efficient_attention(
              _lowercase   ,						_lowercase   ,						_lowercase   ,						query_chunk_size=_lowercase   ,						key_chunk_size=4_0_9_6 * 4		)

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

          # attend to values
          _lowercase							     =       jnp.einsum("b i j, b j d -> b i d"   ,						_lowercase   ,						_lowercase		)

       _lowercase							     =       self.reshape_batch_dim_to_heads(_lowercase		)
       _lowercase							     =       self.proj_attn(_lowercase		)
       return self.dropout_layer(_lowercase   ,						deterministic=_lowercase		)
class   UpperCAmelCase_	(				nn.Module   ):
    snake_case_											=      42
    snake_case_											=      42
    snake_case_											=      42
    snake_case_											=      0.0
    snake_case_											=      False
    snake_case_											=      jnp.floataa
    snake_case_											=      False
    def 		_lowerCamelCase      (	self			:       int		)						->    str:
       # self attention (or cross_attention if only_cross_attention is True)
       _lowercase							     =       FlaxAttention(
           self.dim   ,						self.n_heads   ,						self.d_head   ,						self.dropout   ,						self.use_memory_efficient_attention   ,						dtype=self.dtype		)
       # cross attention
       _lowercase							     =       FlaxAttention(
           self.dim   ,						self.n_heads   ,						self.d_head   ,						self.dropout   ,						self.use_memory_efficient_attention   ,						dtype=self.dtype		)
       _lowercase							     =       FlaxFeedForward(dim=self.dim   ,						dropout=self.dropout   ,						dtype=self.dtype		)
       _lowercase							     =       nn.LayerNorm(epsilon=1e-5   ,						dtype=self.dtype		)
       _lowercase							     =       nn.LayerNorm(epsilon=1e-5   ,						dtype=self.dtype		)
       _lowercase							     =       nn.LayerNorm(epsilon=1e-5   ,						dtype=self.dtype		)
       _lowercase							     =       nn.Dropout(rate=self.dropout		)




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

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

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

       return self.dropout_layer(_lowercase   ,						deterministic=_lowercase		)
class   UpperCAmelCase_	(				nn.Module   ):
    snake_case_											=      42
    snake_case_											=      42
    snake_case_											=      42
    snake_case_											=      1
    snake_case_											=      0.0
    snake_case_											=      False
    snake_case_											=      False
    snake_case_											=      jnp.floataa
    snake_case_											=      False
    def 		_lowerCamelCase      (	self			:       Union[str, Any]		)						->    Union[str, Any]:
       _lowercase							     =       nn.GroupNorm(num_groups=3_2   ,						epsilon=1e-5		)

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

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

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




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

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

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

       _lowercase							     =       hidden_states + residual
       return self.dropout_layer(_lowercase   ,						deterministic=_lowercase		)
class   UpperCAmelCase_	(				nn.Module   ):
    snake_case_											=      42
    snake_case_											=      0.0
    snake_case_											=      jnp.floataa
    def 		_lowerCamelCase      (	self			:       Union[str, Any]		)						->    Optional[int]:
       # The second linear layer needs to be called
       # net_2 for now to match the index of the Sequential layer
       _lowercase							     =       FlaxGEGLU(self.dim   ,						self.dropout   ,						self.dtype		)
       _lowercase							     =       nn.Dense(self.dim   ,						dtype=self.dtype		)




    def __call__(	self			:       Dict   ,						_lowercase			:       int   ,						_lowercase			:       str=True		)						->    str:
       _lowercase							     =       self.net_a(_lowercase   ,						deterministic=_lowercase		)
       _lowercase							     =       self.net_a(_lowercase		)
       return hidden_states
class   UpperCAmelCase_	(				nn.Module   ):
    snake_case_											=      42
    snake_case_											=      0.0
    snake_case_											=      jnp.floataa
    def 		_lowerCamelCase      (	self			:       List[Any]		)						->    Union[str, Any]:
       _lowercase							     =       self.dim * 4
       _lowercase							     =       nn.Dense(inner_dim * 2   ,						dtype=self.dtype		)
       _lowercase							     =       nn.Dropout(rate=self.dropout		)




    def __call__(	self			:       Dict   ,						_lowercase			:       List[str]   ,						_lowercase			:       Optional[int]=True		)						->    Any:
       _lowercase							     =       self.proj(_lowercase		)
       _lowercase     ,  _lowercase							     =       jnp.split(_lowercase   ,						2   ,						axis=2		)
       return self.dropout_layer(hidden_linear * nn.gelu(_lowercase		)   ,						deterministic=_lowercase		)
"""simple docstring"""


import warnings

from ...utils import logging
from .image_processing_convnext import ConvNextImageProcessor


__UpperCamelCase       :    str         =		logging.get_logger(__name__)
class   UpperCAmelCase_	(				lowercase__   ):
    def __init__(	self			:       List[Any]   ,						*_lowercase			:       Optional[Any]   ,						**_lowercase			:       str		)						->    None:
       warnings.warn(
           "The class ConvNextFeatureExtractor is deprecated and will be removed in version 5 of Transformers."
           " Please use ConvNextImageProcessor instead."   ,						_lowercase   ,						)
       super().__init__(*_lowercase   ,						**_lowercase		)