code: 
  '''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    )
  
  
style_context: 
  '''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_    )
  
  