
'''simple docstring'''





import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
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 DinatBackbone, DinatForImageClassification, DinatModel
     from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor



class  lowerCAmelCase__   :



  '''simple docstring'''






  def __init__(self  : int	, lowerCAmelCase_  : Any	, lowerCAmelCase_  : Tuple=13	, lowerCAmelCase_  : str=64	, lowerCAmelCase_  : Dict=4	, lowerCAmelCase_  : Tuple=3	, lowerCAmelCase_  : Optional[int]=16	, lowerCAmelCase_  : Optional[int]=[1, 2, 1]	, lowerCAmelCase_  : Tuple=[2, 4, 8]	, lowerCAmelCase_  : Optional[Any]=3	, lowerCAmelCase_  : List[Any]=[[3], [1, 2], [1]]	, lowerCAmelCase_  : Optional[Any]=2.0	, lowerCAmelCase_  : Optional[Any]=True	, lowerCAmelCase_  : Tuple=0.0	, lowerCAmelCase_  : int=0.0	, lowerCAmelCase_  : Optional[int]=0.1	, lowerCAmelCase_  : Optional[Any]="gelu"	, lowerCAmelCase_  : List[Any]=True	, lowerCAmelCase_  : List[str]=0.02	, lowerCAmelCase_  : Union[str, Any]=1E-5	, lowerCAmelCase_  : List[Any]=True	, lowerCAmelCase_  : List[Any]=None	, lowerCAmelCase_  : List[str]=True	, lowerCAmelCase_  : List[str]=10	, lowerCAmelCase_  : List[Any]=["stage1", "stage2"]	, lowerCAmelCase_  : str=[1, 2]	, )   ->       Tuple:
        snake_case_						       =							parent
        snake_case_						       =							batch_size
        snake_case_						       =							image_size
        snake_case_						       =							patch_size
        snake_case_						       =							num_channels
        snake_case_						       =							embed_dim
        snake_case_						       =							depths
        snake_case_						       =							num_heads
        snake_case_						       =							kernel_size
        snake_case_						       =							dilations
        snake_case_						       =							mlp_ratio
        snake_case_						       =							qkv_bias
        snake_case_						       =							hidden_dropout_prob
        snake_case_						       =							attention_probs_dropout_prob
        snake_case_						       =							drop_path_rate
        snake_case_						       =							hidden_act
        snake_case_						       =							patch_norm
        snake_case_						       =							layer_norm_eps
        snake_case_						       =							initializer_range
        snake_case_						       =							is_training
        snake_case_						       =							scope
        snake_case_						       =							use_labels
        snake_case_						       =							num_labels
        snake_case_						       =							out_features
        snake_case_						       =							out_indices

  def 		_lowerCamelCase	(self  : List[Any]    )   ->       List[Any]:
        snake_case_						       =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

        snake_case_						       =							None
        if self.use_labels:
              snake_case_						       =							ids_tensor([self.batch_size]	, self.num_labels    )

        snake_case_						       =							self.get_config()

        return config, pixel_values, labels

  def 		_lowerCamelCase	(self  : List[str]    )   ->       Optional[Any]:
        return DinatConfig(
            num_labels=self.num_labels	, 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	, kernel_size=self.kernel_size	, dilations=self.dilations	, 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	, patch_norm=self.patch_norm	, layer_norm_eps=self.layer_norm_eps	, initializer_range=self.initializer_range	, out_features=self.out_features	, out_indices=self.out_indices	, )

  def 		_lowerCamelCase	(self  : Dict	, lowerCAmelCase_  : str	, lowerCAmelCase_  : List[Any]	, lowerCAmelCase_  : List[str]    )   ->       List[str]:
        snake_case_						       =							DinatModel(config=lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()
        snake_case_						       =							model(lowerCAmelCase_    )

        snake_case_						       =							snake_case_						       =							(config.image_size // config.patch_size) // (2 ** (len(config.depths    ) - 1))
        snake_case_						       =							int(config.embed_dim * 2 ** (len(config.depths    ) - 1)    )

        self.parent.assertEqual(
            result.last_hidden_state.shape	, (self.batch_size, expected_height, expected_width, expected_dim)    )

  def 		_lowerCamelCase	(self  : int	, lowerCAmelCase_  : Dict	, lowerCAmelCase_  : Optional[Any]	, lowerCAmelCase_  : Any    )   ->       str:
        snake_case_						       =							DinatForImageClassification(lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()
        snake_case_						       =							model(lowerCAmelCase_	, labels=lowerCAmelCase_    )
        self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.num_labels)    )

        # test greyscale images
        snake_case_						       =							1
        snake_case_						       =							DinatForImageClassification(lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()

        snake_case_						       =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
        snake_case_						       =							model(lowerCAmelCase_    )
        self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.num_labels)    )

  def 		_lowerCamelCase	(self  : List[Any]	, lowerCAmelCase_  : Dict	, lowerCAmelCase_  : List[str]	, lowerCAmelCase_  : List[Any]    )   ->       Union[str, Any]:
        snake_case_						       =							DinatBackbone(config=lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()
        snake_case_						       =							model(lowerCAmelCase_    )

        # verify hidden states
        self.parent.assertEqual(len(result.feature_maps    )	, len(config.out_features    )    )
        self.parent.assertListEqual(list(result.feature_maps[0].shape    )	, [self.batch_size, model.channels[0], 16, 16]    )

        # verify channels
        self.parent.assertEqual(len(model.channels    )	, len(config.out_features    )    )

        # verify backbone works with out_features=None
        snake_case_						       =							None
        snake_case_						       =							DinatBackbone(config=lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()
        snake_case_						       =							model(lowerCAmelCase_    )

        # verify feature maps
        self.parent.assertEqual(len(result.feature_maps    )	, 1    )
        self.parent.assertListEqual(list(result.feature_maps[0].shape    )	, [self.batch_size, model.channels[-1], 4, 4]    )

        # verify channels
        self.parent.assertEqual(len(model.channels    )	, 1    )


  def 		_lowerCamelCase	(self  : Union[str, Any]    )   ->       List[str]:
        snake_case_						       =							self.prepare_config_and_inputs()
        snake_case_       ,	snake_case_       ,	snake_case_						       =							config_and_inputs
        snake_case_						       =							{"pixel_values": pixel_values}
        return config, inputs_dict



@require_natten
@require_torch
class  lowerCAmelCase__   (       UpperCamelCase					,							UpperCamelCase					,							unittest.TestCase       ):



  '''simple docstring'''






  _UpperCamelCase						= (
      (
          DinatModel,
          DinatForImageClassification,
          DinatBackbone,
      )
      if is_torch_available()
      else ()
  )
  _UpperCamelCase						= (
      {'feature-extraction': DinatModel, 'image-classification': DinatForImageClassification}
      if is_torch_available()
      else {}
  )
  _UpperCamelCase						= False

  _UpperCamelCase						= False
  _UpperCamelCase						= False
  _UpperCamelCase						= False
  _UpperCamelCase						= False

  def 		_lowerCamelCase	(self  : Tuple    )   ->       Tuple:
        snake_case_						       =							DinatModelTester(self    )
        snake_case_						       =							ConfigTester(self	, config_class=lowerCAmelCase_	, embed_dim=37    )

  def 		_lowerCamelCase	(self  : Tuple    )   ->       Dict:
        self.create_and_test_config_common_properties()
        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  : List[str]    )   ->       Optional[int]:
        return

  def 		_lowerCamelCase	(self  : Dict    )   ->       Union[str, Any]:
        snake_case_						       =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase_    )

  def 		_lowerCamelCase	(self  : Dict    )   ->       Union[str, Any]:
        snake_case_						       =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase_    )

  def 		_lowerCamelCase	(self  : Optional[int]    )   ->       str:
        snake_case_						       =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_backbone(*lowerCAmelCase_    )

  @unittest.skip(reason="Dinat does not use inputs_embeds"    )
  def 		_lowerCamelCase	(self  : Optional[Any]    )   ->       Optional[Any]:
        pass

  @unittest.skip(reason="Dinat does not use feedforward chunking"    )
  def 		_lowerCamelCase	(self  : Optional[int]    )   ->       Dict:
        pass

  def 		_lowerCamelCase	(self  : Tuple    )   ->       Union[str, Any]:
        snake_case_       ,	snake_case_						       =							self.model_tester.prepare_config_and_inputs_for_common()

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

  def 		_lowerCamelCase	(self  : Any    )   ->       List[str]:
        snake_case_       ,	snake_case_						       =							self.model_tester.prepare_config_and_inputs_for_common()

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

              snake_case_						       =							["pixel_values"]
              self.assertListEqual(arg_names[:1]	, lowerCAmelCase_    )

  def 		_lowerCamelCase	(self  : List[Any]    )   ->       Optional[int]:
        self.skipTest("Dinat's attention operation is handled entirely by NATTEN."    )

  def 		_lowerCamelCase	(self  : str	, lowerCAmelCase_  : Union[str, Any]	, lowerCAmelCase_  : Tuple	, lowerCAmelCase_  : int	, lowerCAmelCase_  : List[Any]    )   ->       str:
        snake_case_						       =							model_class(lowerCAmelCase_    )
        model.to(lowerCAmelCase_    )
        model.eval()

        with torch.no_grad():
              snake_case_						       =							model(**self._prepare_for_class(lowerCAmelCase_	, lowerCAmelCase_    )    )

        snake_case_						       =							outputs.hidden_states

        snake_case_						       =							getattr(
            self.model_tester	, "expected_num_hidden_layers"	, len(self.model_tester.depths    ) + 1    )
        self.assertEqual(len(lowerCAmelCase_    )	, lowerCAmelCase_    )

        # Dinat has a different seq_length
        snake_case_						       =							(
            config.patch_size
            if isinstance(config.patch_size	, collections.abc.Iterable    )
            else (config.patch_size, config.patch_size)
        )

        snake_case_						       =							image_size[0] // patch_size[0]
        snake_case_						       =							image_size[1] // patch_size[1]

        self.assertListEqual(
            list(hidden_states[0].shape[-3:]    )	, [height, width, self.model_tester.embed_dim]	, )

        if model_class.__name__ != "DinatBackbone":
              snake_case_						       =							outputs.reshaped_hidden_states
              self.assertEqual(len(lowerCAmelCase_    )	, lowerCAmelCase_    )

              snake_case_       ,	snake_case_       ,	snake_case_       ,	snake_case_						       =							reshaped_hidden_states[0].shape
              snake_case_						       =							(
                  reshaped_hidden_states[0].view(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    ).permute(0	, 2	, 3	, 1    )
              )
              self.assertListEqual(
                  list(reshaped_hidden_states.shape[-3:]    )	, [height, width, self.model_tester.embed_dim]	, )

  def 		_lowerCamelCase	(self  : Dict    )   ->       Tuple:
        snake_case_       ,	snake_case_						       =							self.model_tester.prepare_config_and_inputs_for_common()

        snake_case_						       =							(
            self.model_tester.image_size
            if isinstance(self.model_tester.image_size	, collections.abc.Iterable    )
            else (self.model_tester.image_size, self.model_tester.image_size)
        )

        for model_class in self.all_model_classes:
              snake_case_						       =							True
              self.check_hidden_states_output(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              snake_case_						       =							True

              self.check_hidden_states_output(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )

  @slow
  def 		_lowerCamelCase	(self  : Tuple    )   ->       Optional[int]:
        for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              snake_case_						       =							DinatModel.from_pretrained(lowerCAmelCase_    )
              self.assertIsNotNone(lowerCAmelCase_    )


  def 		_lowerCamelCase	(self  : List[str]    )   ->       List[str]:
        snake_case_       ,	snake_case_						       =							self.model_tester.prepare_config_and_inputs_for_common()

        snake_case_						       =							_config_zero_init(lowerCAmelCase_    )
        for model_class in self.all_model_classes:
              snake_case_						       =							model_class(config=lowerCAmelCase_    )
              for name, param in model.named_parameters():
                    if "embeddings" not in name and 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'''	, )



@require_natten
@require_vision
@require_torch
class  lowerCAmelCase__   (       unittest.TestCase       ):



  '''simple docstring'''






  @cached_property
  def 		_lowerCamelCase	(self  : Optional[Any]    )   ->       Any:
        return AutoImageProcessor.from_pretrained("shi-labs/dinat-mini-in1k-224"    ) if is_vision_available() else None


  @slow
  def 		_lowerCamelCase	(self  : int    )   ->       Union[str, Any]:
        snake_case_						       =							DinatForImageClassification.from_pretrained("shi-labs/dinat-mini-in1k-224"    ).to(lowerCAmelCase_    )
        snake_case_						       =							self.default_image_processor

        snake_case_						       =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
        snake_case_						       =							image_processor(images=lowerCAmelCase_	, return_tensors="pt"    ).to(lowerCAmelCase_    )

        # forward pass
        with torch.no_grad():
              snake_case_						       =							model(**lowerCAmelCase_    )

        # verify the logits
        snake_case_						       =							torch.Size((1, 10_00)    )
        self.assertEqual(outputs.logits.shape	, lowerCAmelCase_    )
        snake_case_						       =							torch.tensor([-0.1_545, -0.7_667, 0.4_642]    ).to(lowerCAmelCase_    )
        self.assertTrue(torch.allclose(outputs.logits[0, :3]	, lowerCAmelCase_	, atol=1E-4    )    )



@require_torch
@require_natten
class  lowerCAmelCase__   (       unittest.TestCase					,							UpperCamelCase       ):



  '''simple docstring'''






  _UpperCamelCase						= (DinatBackbone,) if is_torch_available() else ()
  _UpperCamelCase						= DinatConfig

  def 		_lowerCamelCase	(self  : int    )   ->       Tuple:
        snake_case_						       =							DinatModelTester(self    )



'''simple docstring'''





import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp
def   snake_case_       (							_a     :							Optional[int] ,					_a     :							Dict ,					_a     :							Optional[Any] ,					_a     :							Optional[Any] ,					_a     :							int = 4_096			):
      snake_case_       ,	snake_case_       ,	snake_case_						       =							key.shape[-3:]
      snake_case_						       =							value.shape[-1]
      snake_case_						       =							min(_a ,					_a			)
      snake_case_						       =							query / jnp.sqrt(_a			)

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

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

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

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

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

      snake_case_       ,	snake_case_       ,	snake_case_						       =							jax.lax.map(f=_a ,					xs=jnp.arange(0 ,					_a ,					_a			)			)

      snake_case_						       =							jnp.max(_a ,					axis=0 ,					keepdims=_a			)
      snake_case_						       =							jnp.exp(chunk_max - global_max			)

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

      snake_case_						       =							chunk_values.sum(axis=0			)
      snake_case_						       =							jnp.expand_dims(_a ,					-1			).sum(axis=0			)

      return all_values / all_weights
def   snake_case_       (							_a     :							Union[str, Any] ,					_a     :							Optional[int] ,					_a     :							List[Any] ,					_a     :							str=jax.lax.Precision.HIGHEST ,					_a     :							int = 1_024 ,					_a     :							int = 4_096			):
      snake_case_       ,	snake_case_       ,	snake_case_						       =							query.shape[-3:]

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

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

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

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



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






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

  def 		_lowerCamelCase	(self  : Any    )   ->       Optional[Any]:
        snake_case_						       =							self.dim_head * self.heads
        snake_case_						       =							self.dim_head**-0.5

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

        snake_case_						       =							nn.Dense(self.query_dim	, dtype=self.dtype	, name="to_out_0"    )
        snake_case_						       =							nn.Dropout(rate=self.dropout    )

  def 		_lowerCamelCase	(self  : Any	, lowerCAmelCase_  : Optional[Any]    )   ->       Tuple:
        snake_case_       ,	snake_case_       ,	snake_case_						       =							tensor.shape
        snake_case_						       =							self.heads
        snake_case_						       =							tensor.reshape(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, dim // head_size    )
        snake_case_						       =							jnp.transpose(lowerCAmelCase_	, (0, 2, 1, 3)    )
        snake_case_						       =							tensor.reshape(batch_size * head_size	, lowerCAmelCase_	, dim // head_size    )
        return tensor

  def 		_lowerCamelCase	(self  : List[Any]	, lowerCAmelCase_  : Optional[int]    )   ->       List[str]:
        snake_case_       ,	snake_case_       ,	snake_case_						       =							tensor.shape
        snake_case_						       =							self.heads
        snake_case_						       =							tensor.reshape(batch_size // head_size	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
        snake_case_						       =							jnp.transpose(lowerCAmelCase_	, (0, 2, 1, 3)    )
        snake_case_						       =							tensor.reshape(batch_size // head_size	, lowerCAmelCase_	, dim * head_size    )
        return tensor


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

        snake_case_						       =							self.query(lowerCAmelCase_    )
        snake_case_						       =							self.key(lowerCAmelCase_    )
        snake_case_						       =							self.value(lowerCAmelCase_    )

        snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )
        snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )
        snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )

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

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

              snake_case_						       =							jax_memory_efficient_attention(
                  lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, query_chunk_size=lowerCAmelCase_	, key_chunk_size=40_96 * 4    )

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

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

        snake_case_						       =							self.reshape_batch_dim_to_heads(lowerCAmelCase_    )
        snake_case_						       =							self.proj_attn(lowerCAmelCase_    )
        return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 42
  _UpperCamelCase						= 42
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= False
  _UpperCamelCase						= jnp.floataa
  _UpperCamelCase						= False

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


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

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

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

        return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 42
  _UpperCamelCase						= 42
  _UpperCamelCase						= 1
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= False
  _UpperCamelCase						= False
  _UpperCamelCase						= jnp.floataa
  _UpperCamelCase						= False

  def 		_lowerCamelCase	(self  : Optional[Any]    )   ->       Tuple:
        snake_case_						       =							nn.GroupNorm(num_groups=32	, epsilon=1E-5    )

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

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

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


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

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

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

        snake_case_						       =							hidden_states + residual
        return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= jnp.floataa

  def 		_lowerCamelCase	(self  : str    )   ->       List[str]:
        # The second linear layer needs to be called
        # net_2 for now to match the index of the Sequential layer
        snake_case_						       =							FlaxGEGLU(self.dim	, self.dropout	, self.dtype    )
        snake_case_						       =							nn.Dense(self.dim	, dtype=self.dtype    )


  def __call__(self  : int	, lowerCAmelCase_  : Tuple	, lowerCAmelCase_  : Any=True    )   ->       int:
        snake_case_						       =							self.net_a(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
        snake_case_						       =							self.net_a(lowerCAmelCase_    )
        return hidden_states



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= jnp.floataa

  def 		_lowerCamelCase	(self  : str    )   ->       Optional[Any]:
        snake_case_						       =							self.dim * 4
        snake_case_						       =							nn.Dense(inner_dim * 2	, dtype=self.dtype    )
        snake_case_						       =							nn.Dropout(rate=self.dropout    )


  def __call__(self  : Tuple	, lowerCAmelCase_  : Dict	, lowerCAmelCase_  : Optional[Any]=True    )   ->       Union[str, Any]:
        snake_case_						       =							self.proj(lowerCAmelCase_    )
        snake_case_       ,	snake_case_						       =							jnp.split(lowerCAmelCase_	, 2	, axis=2    )
        return self.dropout_layer(hidden_linear * nn.gelu(lowerCAmelCase_    )	, deterministic=lowerCAmelCase_    )

