

import inspect
import unittest

from transformers import SwinaSRConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch
       from torch import nn

       from transformers import SwinaSRForImageSuperResolution, SwinaSRModel
       from transformers.models.swinasr.modeling_swinasr import SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
       from PIL import Image

       from transformers import SwinaSRImageProcessor
class 					A_   :


    '''simple docstring'''






    def __init__(self     ,      lowercase__     ,      lowercase__=13     ,      lowercase__=32     ,      lowercase__=1     ,      lowercase__=3     ,      lowercase__=16     ,      lowercase__=[1, 2, 1]     ,      lowercase__=[2, 2, 4]     ,      lowercase__=2     ,      lowercase__=2.0     ,      lowercase__=True     ,      lowercase__=0.0     ,      lowercase__=0.0     ,      lowercase__=0.1     ,      lowercase__="gelu"     ,      lowercase__=False     ,      lowercase__=True     ,      lowercase__=0.02     ,      lowercase__=1E-5     ,      lowercase__=True     ,      lowercase__=None     ,      lowercase__=False     ,      lowercase__=2     ,      )     ->      Tuple:
        __UpperCAmelCase									=			parent
        __UpperCAmelCase									=			batch_size
        __UpperCAmelCase									=			image_size
        __UpperCAmelCase									=			patch_size
        __UpperCAmelCase									=			num_channels
        __UpperCAmelCase									=			embed_dim
        __UpperCAmelCase									=			depths
        __UpperCAmelCase									=			num_heads
        __UpperCAmelCase									=			window_size
        __UpperCAmelCase									=			mlp_ratio
        __UpperCAmelCase									=			qkv_bias
        __UpperCAmelCase									=			hidden_dropout_prob
        __UpperCAmelCase									=			attention_probs_dropout_prob
        __UpperCAmelCase									=			drop_path_rate
        __UpperCAmelCase									=			hidden_act
        __UpperCAmelCase									=			use_absolute_embeddings
        __UpperCAmelCase									=			patch_norm
        __UpperCAmelCase									=			layer_norm_eps
        __UpperCAmelCase									=			initializer_range
        __UpperCAmelCase									=			is_training
        __UpperCAmelCase									=			scope
        __UpperCAmelCase									=			use_labels
        __UpperCAmelCase									=			upscale

        # here we set some attributes to make tests pass
        __UpperCAmelCase									=			len(lowercase__  )
        __UpperCAmelCase									=			embed_dim
        __UpperCAmelCase									=			(image_size // patch_size) ** 2


    def      lowerCAmelCase_						(self  )     ->      Optional[Any]:
        __UpperCAmelCase									=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

        __UpperCAmelCase									=			None
        if self.use_labels:
            __UpperCAmelCase									=			ids_tensor([self.batch_size]     ,      self.type_sequence_label_size  )

        __UpperCAmelCase									=			self.get_config()

        return config, pixel_values, labels


    def      lowerCAmelCase_						(self  )     ->      Union[str, Any]:
        return SwinaSRConfig(
            image_size=self.image_size     ,      patch_size=self.patch_size     ,      num_channels=self.num_channels     ,      embed_dim=self.embed_dim     ,      depths=self.depths     ,      num_heads=self.num_heads     ,      window_size=self.window_size     ,      mlp_ratio=self.mlp_ratio     ,      qkv_bias=self.qkv_bias     ,      hidden_dropout_prob=self.hidden_dropout_prob     ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,      drop_path_rate=self.drop_path_rate     ,      hidden_act=self.hidden_act     ,      use_absolute_embeddings=self.use_absolute_embeddings     ,      path_norm=self.patch_norm     ,      layer_norm_eps=self.layer_norm_eps     ,      initializer_range=self.initializer_range     ,      upscale=self.upscale     ,      )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__  )     ->      int:
        __UpperCAmelCase									=			SwinaSRModel(config=lowercase__  )
        model.to(lowercase__  )
        model.eval()
        __UpperCAmelCase									=			model(lowercase__  )

        self.parent.assertEqual(
            result.last_hidden_state.shape     ,      (self.batch_size, self.embed_dim, self.image_size, self.image_size)  )


    def      lowerCAmelCase_						(self     ,      lowercase__     ,      lowercase__     ,      lowercase__  )     ->      List[Any]:
        __UpperCAmelCase									=			SwinaSRForImageSuperResolution(lowercase__  )
        model.to(lowercase__  )
        model.eval()
        __UpperCAmelCase									=			model(lowercase__  )

        __UpperCAmelCase									=			self.image_size * self.upscale
        self.parent.assertEqual(
            result.reconstruction.shape     ,      (self.batch_size, self.num_channels, expected_image_size, expected_image_size)  )






    def      lowerCAmelCase_						(self  )     ->      str:
        __UpperCAmelCase									=			self.prepare_config_and_inputs()
        __UpperCAmelCase					,     __UpperCAmelCase					,     __UpperCAmelCase									=			config_and_inputs
        __UpperCAmelCase									=			{'''pixel_values''': pixel_values}
        return config, inputs_dict






@require_torch
class 					A_   (      _a ,     _a ,     unittest.TestCase						):


    '''simple docstring'''






    a__     							=       (SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
    a__     							=       {"feature-extraction": SwinaSRModel} if is_torch_available() else {}

    a__     							=       False
    a__     							=       False
    a__     							=       False
    a__     							=       False
    a__     							=       False


    def      lowerCAmelCase_						(self  )     ->      List[str]:
        __UpperCAmelCase									=			SwinaSRModelTester(self  )
        __UpperCAmelCase									=			ConfigTester(self     ,      config_class=lowercase__     ,      embed_dim=37  )


    def      lowerCAmelCase_						(self  )     ->      Union[str, Any]:
        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()


    def      lowerCAmelCase_						(self  )     ->      Any:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Dict:
        __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_super_resolution(*lowercase__  )


    @unittest.skip(reason='''Got `CUDA error: misaligned address` with PyTorch 2.0.0.'''  )
    def      lowerCAmelCase_						(self  )     ->      Any:
        pass


    @unittest.skip(reason='''Swin2SR does not use inputs_embeds'''  )
    def      lowerCAmelCase_						(self  )     ->      List[Any]:
        pass


    @unittest.skip(reason='''Swin2SR does not support training yet'''  )
    def      lowerCAmelCase_						(self  )     ->      List[Any]:
        pass


    @unittest.skip(reason='''Swin2SR does not support training yet'''  )
    def      lowerCAmelCase_						(self  )     ->      List[str]:
        pass


    def      lowerCAmelCase_						(self  )     ->      int:
        __UpperCAmelCase					,     __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            __UpperCAmelCase									=			model_class(lowercase__  )
            self.assertIsInstance(model.get_input_embeddings()     ,      (nn.Module)  )
            __UpperCAmelCase									=			model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(lowercase__     ,      nn.Linear  )  )


    def      lowerCAmelCase_						(self  )     ->      int:
        __UpperCAmelCase					,     __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            __UpperCAmelCase									=			model_class(lowercase__  )
            __UpperCAmelCase									=			inspect.signature(model.forward  )
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            __UpperCAmelCase									=			[*signature.parameters.keys()]

            __UpperCAmelCase									=			['''pixel_values''']
            self.assertListEqual(arg_names[:1]     ,      lowercase__  )


    @slow
    def      lowerCAmelCase_						(self  )     ->      Union[str, Any]:
        for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCAmelCase									=			SwinaSRModel.from_pretrained(lowercase__  )
            self.assertIsNotNone(lowercase__  )


    def      lowerCAmelCase_						(self  )     ->      Tuple:
        __UpperCAmelCase					,     __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs_for_common()

        __UpperCAmelCase									=			_config_zero_init(lowercase__  )
        for model_class in self.all_model_classes:
            __UpperCAmelCase									=			model_class(config=lowercase__  )
            for name, param in model.named_parameters():
                if "logit_scale" in name:
                    continue
                if param.requires_grad:
                    self.assertIn(
                        ((param.data.mean() * 1E9).round() / 1E9).item()     ,      [0.0, 1.0]     ,      msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''     ,      )






    def      lowerCAmelCase_						(self  )     ->      List[str]:
        __UpperCAmelCase					,     __UpperCAmelCase									=			self.model_tester.prepare_config_and_inputs_for_common()
        __UpperCAmelCase									=			True

        for model_class in self.all_model_classes:
            __UpperCAmelCase									=			True
            __UpperCAmelCase									=			False
            __UpperCAmelCase									=			True
            __UpperCAmelCase									=			model_class(lowercase__  )
            model.to(lowercase__  )
            model.eval()
            with torch.no_grad():
                __UpperCAmelCase									=			model(**self._prepare_for_class(lowercase__     ,      lowercase__  )  )
            __UpperCAmelCase									=			outputs.attentions
            __UpperCAmelCase									=			len(self.model_tester.depths  )
            self.assertEqual(len(lowercase__  )     ,      lowercase__  )

            # check that output_attentions also work using config
            del inputs_dict["output_attentions"]
            __UpperCAmelCase									=			True
            __UpperCAmelCase									=			config.window_size**2
            __UpperCAmelCase									=			model_class(lowercase__  )
            model.to(lowercase__  )
            model.eval()
            with torch.no_grad():
                __UpperCAmelCase									=			model(**self._prepare_for_class(lowercase__     ,      lowercase__  )  )
            __UpperCAmelCase									=			outputs.attentions
            self.assertEqual(len(lowercase__  )     ,      lowercase__  )

            self.assertListEqual(
                list(attentions[0].shape[-3:]  )     ,      [self.model_tester.num_heads[0], window_size_squared, window_size_squared]     ,      )
            __UpperCAmelCase									=			len(lowercase__  )

            # Check attention is always last and order is fine
            __UpperCAmelCase									=			True
            __UpperCAmelCase									=			True
            __UpperCAmelCase									=			model_class(lowercase__  )
            model.to(lowercase__  )
            model.eval()
            with torch.no_grad():
                __UpperCAmelCase									=			model(**self._prepare_for_class(lowercase__     ,      lowercase__  )  )

            self.assertEqual(out_len + 1     ,      len(lowercase__  )  )

            __UpperCAmelCase									=			outputs.attentions

            self.assertEqual(len(lowercase__  )     ,      lowercase__  )

            self.assertListEqual(
                list(self_attentions[0].shape[-3:]  )     ,      [self.model_tester.num_heads[0], window_size_squared, window_size_squared]     ,      )






@require_vision
@require_torch
@slow
class 					A_   (      unittest.TestCase						):


    '''simple docstring'''






    def      lowerCAmelCase_						(self  )     ->      str:
        __UpperCAmelCase									=			SwinaSRImageProcessor()
        __UpperCAmelCase									=			SwinaSRForImageSuperResolution.from_pretrained('''caidas/swin2SR-classical-sr-x2-64'''  ).to(lowercase__  )

        __UpperCAmelCase									=			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''  )
        __UpperCAmelCase									=			processor(images=lowercase__     ,      return_tensors='''pt'''  ).to(lowercase__  )

        # forward pass
        with torch.no_grad():
            __UpperCAmelCase									=			model(**lowercase__  )

        # verify the logits
        __UpperCAmelCase									=			torch.Size([1, 3, 976, 1_296]  )
        self.assertEqual(outputs.reconstruction.shape     ,      lowercase__  )
        __UpperCAmelCase									=			torch.tensor(
            [[0.5458, 0.5546, 0.5638], [0.5526, 0.5565, 0.5651], [0.5396, 0.5426, 0.5621]]  ).to(lowercase__  )
        self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]     ,      lowercase__     ,      atol=1E-4  )  )







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__  )




