import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp
def    __magic_name__(	__UpperCAmelCase			,    __UpperCAmelCase			,    __UpperCAmelCase			,    __UpperCAmelCase			,    __UpperCAmelCase = 4096		)    ->     Optional[Any]:



      '''simple docstring'''

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

      @functools.partial(jax.checkpoint			,    prevent_cse=__UpperCAmelCase		)
      def summarize_chunk(__UpperCAmelCase			,    __UpperCAmelCase			,    __UpperCAmelCase		):
            _lowerCamelCase           =					jnp.einsum('''...qhd,...khd->...qhk'''			,    __UpperCAmelCase			,    __UpperCAmelCase			,    precision=__UpperCAmelCase		)

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

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

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

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

      _lowerCamelCase     ,       _lowerCamelCase     ,       _lowerCamelCase           =					jax.lax.map(f=__UpperCAmelCase			,    xs=jnp.arange(0			,    __UpperCAmelCase			,    __UpperCAmelCase		)		)

      _lowerCamelCase           =					jnp.max(__UpperCAmelCase			,    axis=0			,    keepdims=__UpperCAmelCase		)
      _lowerCamelCase           =					jnp.exp(chunk_max - global_max		)

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

      _lowerCamelCase           =					chunk_values.sum(axis=0		)
      _lowerCamelCase           =					jnp.expand_dims(__UpperCAmelCase			,    -1		).sum(axis=0		)

      return all_values / all_weights
def    __magic_name__(	__UpperCAmelCase			,    __UpperCAmelCase			,    __UpperCAmelCase			,    __UpperCAmelCase=jax.lax.Precision.HIGHEST			,    __UpperCAmelCase = 1024			,    __UpperCAmelCase = 4096		)    ->     Optional[int]:



      '''simple docstring'''

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

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

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

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

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



class 			UpperCamelCase					(			nn.Module			):

     '''simple docstring'''

     A_					    =	42
     A_					    =	8
     A_					    =	64
     A_					    =	0.0
     A_					    =	False
     A_					    =	jnp.floataa



     def 		UpperCamelCase_    (       self    )			->					List[str]:






           """simple docstring"""
           _lowerCamelCase           =					self.dim_head * self.heads
           _lowerCamelCase           =					self.dim_head**-0.5

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

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



     def 		UpperCamelCase_    (       self					,					A_    )			->					List[Any]:






           """simple docstring"""
           _lowerCamelCase     ,       _lowerCamelCase     ,       _lowerCamelCase           =					tensor.shape
           _lowerCamelCase           =					self.heads
           _lowerCamelCase           =					tensor.reshape(A_					,					A_					,					A_					,					dim // head_size    )
           _lowerCamelCase           =					jnp.transpose(A_					,					(0, 2, 1, 3)    )
           _lowerCamelCase           =					tensor.reshape(batch_size * head_size					,					A_					,					dim // head_size    )
           return tensor



     def 		UpperCamelCase_    (       self					,					A_    )			->					List[str]:






           """simple docstring"""
           _lowerCamelCase     ,       _lowerCamelCase     ,       _lowerCamelCase           =					tensor.shape
           _lowerCamelCase           =					self.heads
           _lowerCamelCase           =					tensor.reshape(batch_size // head_size					,					A_					,					A_					,					A_    )
           _lowerCamelCase           =					jnp.transpose(A_					,					(0, 2, 1, 3)    )
           _lowerCamelCase           =					tensor.reshape(batch_size // head_size					,					A_					,					dim * head_size    )
           return tensor





     def __call__(       self					,					A_					,					A_=None					,					A_=True    )			->					Optional[Any]:






           """simple docstring"""
           _lowerCamelCase           =					hidden_states if context is None else context

           _lowerCamelCase           =					self.query(A_    )
           _lowerCamelCase           =					self.key(A_    )
           _lowerCamelCase           =					self.value(A_    )

           _lowerCamelCase           =					self.reshape_heads_to_batch_dim(A_    )
           _lowerCamelCase           =					self.reshape_heads_to_batch_dim(A_    )
           _lowerCamelCase           =					self.reshape_heads_to_batch_dim(A_    )

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

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

                 _lowerCamelCase           =					jax_memory_efficient_attention(
                     A_					,					A_					,					A_					,					query_chunk_size=A_					,					key_chunk_size=40_96 * 4    )

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

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

           _lowerCamelCase           =					self.reshape_batch_dim_to_heads(A_    )
           _lowerCamelCase           =					self.proj_attn(A_    )
           return self.dropout_layer(A_					,					deterministic=A_    )





class 			UpperCamelCase					(			nn.Module			):

     '''simple docstring'''

     A_					    =	42
     A_					    =	42
     A_					    =	42
     A_					    =	0.0
     A_					    =	False
     A_					    =	jnp.floataa
     A_					    =	False



     def 		UpperCamelCase_    (       self    )			->					str:






           """simple docstring"""
           # self attention (or cross_attention if only_cross_attention is True)
           _lowerCamelCase           =					FlaxAttention(
               self.dim					,					self.n_heads					,					self.d_head					,					self.dropout					,					self.use_memory_efficient_attention					,					dtype=self.dtype    )
           # cross attention
           _lowerCamelCase           =					FlaxAttention(
               self.dim					,					self.n_heads					,					self.d_head					,					self.dropout					,					self.use_memory_efficient_attention					,					dtype=self.dtype    )
           _lowerCamelCase           =					FlaxFeedForward(dim=self.dim					,					dropout=self.dropout					,					dtype=self.dtype    )
           _lowerCamelCase           =					nn.LayerNorm(epsilon=1E-5					,					dtype=self.dtype    )
           _lowerCamelCase           =					nn.LayerNorm(epsilon=1E-5					,					dtype=self.dtype    )
           _lowerCamelCase           =					nn.LayerNorm(epsilon=1E-5					,					dtype=self.dtype    )
           _lowerCamelCase           =					nn.Dropout(rate=self.dropout    )





     def __call__(       self					,					A_					,					A_					,					A_=True    )			->					Dict:






           """simple docstring"""
           # self attention
           _lowerCamelCase           =					hidden_states
           if self.only_cross_attention:
                 _lowerCamelCase           =					self.attna(self.norma(A_    )					,					A_					,					deterministic=A_    )
           else:
                 _lowerCamelCase           =					self.attna(self.norma(A_    )					,					deterministic=A_    )
           _lowerCamelCase           =					hidden_states + residual

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

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

           return self.dropout_layer(A_					,					deterministic=A_    )





class 			UpperCamelCase					(			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 		UpperCamelCase_    (       self    )			->					int:






           """simple docstring"""
           _lowerCamelCase           =					nn.GroupNorm(num_groups=32					,					epsilon=1E-5    )

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

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

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





     def __call__(       self					,					A_					,					A_					,					A_=True    )			->					Optional[int]:






           """simple docstring"""
           _lowerCamelCase     ,       _lowerCamelCase     ,       _lowerCamelCase     ,       _lowerCamelCase           =					hidden_states.shape
           _lowerCamelCase           =					hidden_states
           _lowerCamelCase           =					self.norm(A_    )
           if self.use_linear_projection:
                 _lowerCamelCase           =					hidden_states.reshape(A_					,					height * width					,					A_    )
                 _lowerCamelCase           =					self.proj_in(A_    )
           else:
                 _lowerCamelCase           =					self.proj_in(A_    )
                 _lowerCamelCase           =					hidden_states.reshape(A_					,					height * width					,					A_    )

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

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

           _lowerCamelCase           =					hidden_states + residual
           return self.dropout_layer(A_					,					deterministic=A_    )





class 			UpperCamelCase					(			nn.Module			):

     '''simple docstring'''

     A_					    =	42
     A_					    =	0.0
     A_					    =	jnp.floataa



     def 		UpperCamelCase_    (       self    )			->					List[Any]:






           """simple docstring"""
           # The second linear layer needs to be called
           # net_2 for now to match the index of the Sequential layer
           _lowerCamelCase           =					FlaxGEGLU(self.dim					,					self.dropout					,					self.dtype    )
           _lowerCamelCase           =					nn.Dense(self.dim					,					dtype=self.dtype    )





     def __call__(       self					,					A_					,					A_=True    )			->					Union[str, Any]:






           """simple docstring"""
           _lowerCamelCase           =					self.net_a(A_					,					deterministic=A_    )
           _lowerCamelCase           =					self.net_a(A_    )
           return hidden_states





class 			UpperCamelCase					(			nn.Module			):

     '''simple docstring'''

     A_					    =	42
     A_					    =	0.0
     A_					    =	jnp.floataa



     def 		UpperCamelCase_    (       self    )			->					Tuple:






           """simple docstring"""
           _lowerCamelCase           =					self.dim * 4
           _lowerCamelCase           =					nn.Dense(inner_dim * 2					,					dtype=self.dtype    )
           _lowerCamelCase           =					nn.Dropout(rate=self.dropout    )





     def __call__(       self					,					A_					,					A_=True    )			->					List[Any]:






           """simple docstring"""
           _lowerCamelCase           =					self.proj(A_    )
           _lowerCamelCase     ,       _lowerCamelCase           =					jnp.split(A_					,					2					,					axis=2    )
           return self.dropout_layer(hidden_linear * nn.gelu(A_    )					,					deterministic=A_    )
# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import sys
import warnings
from os.path import abspath, dirname, join


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
snake_case__						      =			abspath(join(dirname(dirname(dirname(__file__))), 'src'))
sys.path.insert(1, git_repo_path)


# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action='ignore', category=FutureWarning)
def    __magic_name__(	__UpperCAmelCase		)    ->     int:



      '''simple docstring'''

      from transformers.testing_utils import pytest_addoption_shared

      pytest_addoption_shared(__UpperCAmelCase		)




def    __magic_name__(	__UpperCAmelCase		)    ->     List[str]:



      '''simple docstring'''

      from transformers.testing_utils import pytest_terminal_summary_main

      _lowerCamelCase           =					terminalreporter.config.getoption('''--make-reports'''		)
      if make_reports:
            pytest_terminal_summary_main(__UpperCAmelCase			,    id=__UpperCAmelCase		)