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