

import inspect
import unittest
import warnings

from transformers import DeiTConfig
from transformers.models.auto import get_values
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
			import torch
			from torch import nn

			from transformers import (
			    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
			    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
			    MODEL_MAPPING,
			    DeiTForImageClassification,
			    DeiTForImageClassificationWithTeacher,
			    DeiTForMaskedImageModeling,
			    DeiTModel,
			)
			from transformers.models.deit.modeling_deit import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
			from PIL import Image

			from transformers import DeiTImageProcessor






class _UpperCamelCase							:
			'''simple docstring'''







			def __init__(			self							: Any      ,      a							: List[str]      ,      a							: Any=13      ,      a							: str=30      ,      a							: Optional[int]=2      ,      a							: List[str]=3      ,      a							: str=True      ,      a							: Union[str, Any]=True      ,      a							: Optional[Any]=32      ,      a							: List[Any]=5      ,      a							: str=4      ,      a							: int=37      ,      a							: Optional[int]="gelu"      ,      a							: Union[str, Any]=0.1      ,      a							: List[str]=0.1      ,      a							: int=10      ,      a							: Dict=0.02      ,      a							: str=3      ,      a							: Optional[Any]=None      ,      a							: List[Any]=2      ,      )							->       Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       parent
				SCREAMING_SNAKE_CASE   :      List[Any]				     =       batch_size
				SCREAMING_SNAKE_CASE   :      Any				     =       image_size
				SCREAMING_SNAKE_CASE   :      List[str]				     =       patch_size
				SCREAMING_SNAKE_CASE   :      Dict				     =       num_channels
				SCREAMING_SNAKE_CASE   :      Any				     =       is_training
				SCREAMING_SNAKE_CASE   :      int				     =       use_labels
				SCREAMING_SNAKE_CASE   :      int				     =       hidden_size
				SCREAMING_SNAKE_CASE   :      int				     =       num_hidden_layers
				SCREAMING_SNAKE_CASE   :      int				     =       num_attention_heads
				SCREAMING_SNAKE_CASE   :      Tuple				     =       intermediate_size
				SCREAMING_SNAKE_CASE   :      Any				     =       hidden_act
				SCREAMING_SNAKE_CASE   :      Dict				     =       hidden_dropout_prob
				SCREAMING_SNAKE_CASE   :      List[Any]				     =       attention_probs_dropout_prob
				SCREAMING_SNAKE_CASE   :      List[Any]				     =       type_sequence_label_size
				SCREAMING_SNAKE_CASE   :      Any				     =       initializer_range
				SCREAMING_SNAKE_CASE   :      List[Any]				     =       scope
				SCREAMING_SNAKE_CASE   :      Any				     =       encoder_stride

				# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
				SCREAMING_SNAKE_CASE   :      int				     =       (image_size // patch_size) ** 2
				SCREAMING_SNAKE_CASE   :      Tuple				     =       num_patches + 2
			def 							__UpperCamelCase (			self							: Dict	)							->       Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Optional[int]				     =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

				SCREAMING_SNAKE_CASE   :      Tuple				     =       self.get_config()

				return config, pixel_values, labels
			def 							__UpperCamelCase (			self							: Union[str, Any]	)							->       Union[str, Any]:

				"""simple docstring"""

				return DeiTConfig(
				    image_size=self.image_size      ,      patch_size=self.patch_size      ,      num_channels=self.num_channels      ,      hidden_size=self.hidden_size      ,      num_hidden_layers=self.num_hidden_layers      ,      num_attention_heads=self.num_attention_heads      ,      intermediate_size=self.intermediate_size      ,      hidden_act=self.hidden_act      ,      hidden_dropout_prob=self.hidden_dropout_prob      ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,      is_decoder=lowercase__      ,      initializer_range=self.initializer_range      ,      encoder_stride=self.encoder_stride      ,      )
			def 							__UpperCamelCase (			self							: Union[str, Any]      ,      a							: int      ,      a							: Dict      ,      a							: Dict	)							->       str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      List[str]				     =       DeiTModel(config=lowercase__	)
				model.to(lowercase__	)
				model.eval()
				SCREAMING_SNAKE_CASE   :      List[Any]				     =       model(lowercase__	)
				self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)	)
			def 							__UpperCamelCase (			self							: Tuple      ,      a							: List[Any]      ,      a							: Optional[Any]      ,      a							: Tuple	)							->       Union[str, Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Optional[int]				     =       DeiTForMaskedImageModeling(config=lowercase__	)
				model.to(lowercase__	)
				model.eval()
				SCREAMING_SNAKE_CASE   :      Optional[int]				     =       model(lowercase__	)
				self.parent.assertEqual(
				    result.reconstruction.shape      ,      (self.batch_size, self.num_channels, self.image_size, self.image_size)	)

				# test greyscale images
				SCREAMING_SNAKE_CASE   :      str				     =       1
				SCREAMING_SNAKE_CASE   :      List[str]				     =       DeiTForMaskedImageModeling(lowercase__	)
				model.to(lowercase__	)
				model.eval()

				SCREAMING_SNAKE_CASE   :      int				     =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
				SCREAMING_SNAKE_CASE   :      List[Any]				     =       model(lowercase__	)
				self.parent.assertEqual(result.reconstruction.shape      ,      (self.batch_size, 1, self.image_size, self.image_size)	)
			def 							__UpperCamelCase (			self							: Union[str, Any]      ,      a							: List[Any]      ,      a							: Tuple      ,      a							: Union[str, Any]	)							->       List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Optional[int]				     =       self.type_sequence_label_size
				SCREAMING_SNAKE_CASE   :      int				     =       DeiTForImageClassification(lowercase__	)
				model.to(lowercase__	)
				model.eval()
				SCREAMING_SNAKE_CASE   :      Any				     =       model(lowercase__      ,      labels=lowercase__	)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.type_sequence_label_size)	)

				# test greyscale images
				SCREAMING_SNAKE_CASE   :      Any				     =       1
				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       DeiTForImageClassification(lowercase__	)
				model.to(lowercase__	)
				model.eval()

				SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
				SCREAMING_SNAKE_CASE   :      Dict				     =       model(lowercase__      ,      labels=lowercase__	)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.type_sequence_label_size)	)






			def 							__UpperCamelCase (			self							: Dict	)							->       str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       self.prepare_config_and_inputs()
				(
				    SCREAMING_SNAKE_CASE		
				)   :      str				     =       config_and_inputs
				SCREAMING_SNAKE_CASE   :      Tuple				     =       {"pixel_values": pixel_values}
				return config, inputs_dict






@require_torch
class _UpperCamelCase							(							_UpperCAmelCase							,	_UpperCAmelCase							,	unittest.TestCase   ):
			'''simple docstring'''







			lowerCamelCase__           =(
			    (
			        DeiTModel,
			        DeiTForImageClassification,
			        DeiTForImageClassificationWithTeacher,
			        DeiTForMaskedImageModeling,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase__           =(
			    {
			        'feature-extraction': DeiTModel,
			        'image-classification': (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
			    }
			    if is_torch_available()
			    else {}
			)

			lowerCamelCase__           =False
			lowerCamelCase__           =False
			lowerCamelCase__           =False
			def 							__UpperCamelCase (			self							: Union[str, Any]	)							->       Union[str, Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      List[str]				     =       DeiTModelTester(self	)
				SCREAMING_SNAKE_CASE   :      List[str]				     =       ConfigTester(self      ,      config_class=lowercase__      ,      has_text_modality=lowercase__      ,      hidden_size=37	)
			def 							__UpperCamelCase (			self							: Union[str, Any]	)							->       Tuple:

				"""simple docstring"""

				self.config_tester.run_common_tests()
			@unittest.skip(reason="DeiT does not use inputs_embeds"	)
			def 							__UpperCamelCase (			self							: Optional[int]	)							->       Tuple:

				"""simple docstring"""

				pass
			def 							__UpperCamelCase (			self							: Dict	)							->       Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      int				     =       self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
					SCREAMING_SNAKE_CASE   :      Optional[int]				     =       model_class(lowercase__	)
					self.assertIsInstance(model.get_input_embeddings()      ,      (nn.Module)	)
					SCREAMING_SNAKE_CASE   :      List[str]				     =       model.get_output_embeddings()
					self.assertTrue(x is None or isinstance(lowercase__      ,      nn.Linear	)	)
			def 							__UpperCamelCase (			self							: Any	)							->       Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       self.model_tester.prepare_config_and_inputs_for_common()

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

					SCREAMING_SNAKE_CASE   :      Any				     =       ["pixel_values"]
					self.assertListEqual(arg_names[:1]      ,      lowercase__	)
			def 							__UpperCamelCase (			self							: List[str]	)							->       Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Optional[int]				     =       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowercase__	)
			def 							__UpperCamelCase (			self							: List[str]	)							->       List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      str				     =       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_image_modeling(*lowercase__	)
			def 							__UpperCamelCase (			self							: Tuple	)							->       str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Tuple				     =       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_image_classification(*lowercase__	)
			def 							__UpperCamelCase (			self							: Union[str, Any]      ,      a							: str      ,      a							: List[str]      ,      a							: Optional[int]=False	)							->       List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Tuple				     =       super()._prepare_for_class(lowercase__      ,      lowercase__      ,      return_labels=lowercase__	)

				if return_labels:
					if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
						del inputs_dict["labels"]

				return inputs_dict
			def 							__UpperCamelCase (			self							: str	)							->       Union[str, Any]:

				"""simple docstring"""

				if not self.model_tester.is_training:
					return

				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       self.model_tester.prepare_config_and_inputs_for_common()
				SCREAMING_SNAKE_CASE   :      List[str]				     =       True

				for model_class in self.all_model_classes:
					# DeiTForImageClassificationWithTeacher supports inference-only
					if (
					    model_class in get_values(lowercase__	)
					    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
					):
						continue
					SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       model_class(lowercase__	)
					model.to(lowercase__	)
					model.train()
					SCREAMING_SNAKE_CASE   :      Optional[int]				     =       self._prepare_for_class(lowercase__      ,      lowercase__      ,      return_labels=lowercase__	)
					SCREAMING_SNAKE_CASE   :      List[Any]				     =       model(**lowercase__	).loss
					loss.backward()
			def 							__UpperCamelCase (			self							: Dict	)							->       List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      Optional[int]				     =       self.model_tester.prepare_config_and_inputs_for_common()
				if not self.model_tester.is_training:
					return

				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       False
				SCREAMING_SNAKE_CASE   :      Dict				     =       True

				for model_class in self.all_model_classes:
					if model_class in get_values(lowercase__	) or not model_class.supports_gradient_checkpointing:
						continue
					# DeiTForImageClassificationWithTeacher supports inference-only
					if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
						continue
					SCREAMING_SNAKE_CASE   :      Any				     =       model_class(lowercase__	)
					model.gradient_checkpointing_enable()
					model.to(lowercase__	)
					model.train()
					SCREAMING_SNAKE_CASE   :      List[Any]				     =       self._prepare_for_class(lowercase__      ,      lowercase__      ,      return_labels=lowercase__	)
					SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       model(**lowercase__	).loss
					loss.backward()
			def 							__UpperCamelCase (			self							: List[str]	)							->       Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      str				     =       self.model_tester.prepare_config_and_inputs_for_common()

				SCREAMING_SNAKE_CASE   :      List[Any]				     =       [
				    {"title": "multi_label_classification", "num_labels": 2, "dtype": torch.float},
				    {"title": "single_label_classification", "num_labels": 1, "dtype": torch.long},
				    {"title": "regression", "num_labels": 1, "dtype": torch.float},
				]

				for model_class in self.all_model_classes:
					if (
					    model_class
					    not in [
					        *get_values(lowercase__	),
					        *get_values(lowercase__	),
					    ]
					    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
					):
						continue

					for problem_type in problem_types:
						with self.subTest(msg=F"Testing {model_class} with {problem_type['title']}"	):
							SCREAMING_SNAKE_CASE   :      int				     =       problem_type["title"]
							SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       problem_type["num_labels"]

							SCREAMING_SNAKE_CASE   :      List[Any]				     =       model_class(lowercase__	)
							model.to(lowercase__	)
							model.train()

							SCREAMING_SNAKE_CASE   :      int				     =       self._prepare_for_class(lowercase__      ,      lowercase__      ,      return_labels=lowercase__	)

							if problem_type["num_labels"] > 1:
								SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       inputs["labels"].unsqueeze(1	).repeat(1      ,      problem_type["num_labels"]	)

							SCREAMING_SNAKE_CASE   :      Optional[int]				     =       inputs["labels"].to(problem_type["dtype"]	)

							# This tests that we do not trigger the warning form PyTorch "Using a target size that is different
							# to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
							# they have the same size." which is a symptom something in wrong for the regression problem.
							# See https://github.com/huggingface/transformers/issues/11780
							with warnings.catch_warnings(record=lowercase__	) as warning_list:
								SCREAMING_SNAKE_CASE   :      Any				     =       model(**lowercase__	).loss
							for w in warning_list:
								if "Using a target size that is different to the input size" in str(w.message	):
									raise ValueError(
									    F"Something is going wrong in the regression problem: intercepted {w.message}"	)

							loss.backward()






			@slow
			def 							__UpperCamelCase (			self							: Dict	)							->       Optional[Any]:

				"""simple docstring"""

				for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					SCREAMING_SNAKE_CASE   :      Tuple				     =       DeiTModel.from_pretrained(lowercase__	)
					self.assertIsNotNone(lowercase__	)

def  lowerCamelCase__ (					):
	SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
	return image






@require_torch
@require_vision
class _UpperCamelCase							(							unittest.TestCase   ):
			'''simple docstring'''







			@cached_property
			def 							__UpperCamelCase (			self							: Union[str, Any]	)							->       int:

				"""simple docstring"""

				return (
				    DeiTImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224"	)
				    if is_vision_available()
				    else None
				)
			@slow
			def 							__UpperCamelCase (			self							: Union[str, Any]	)							->       Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      str				     =       DeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"	).to(
				    lowercase__	)

				SCREAMING_SNAKE_CASE   :      Tuple				     =       self.default_image_processor
				SCREAMING_SNAKE_CASE   :      Dict				     =       prepare_img()
				SCREAMING_SNAKE_CASE   :      Dict				     =       image_processor(images=lowercase__      ,      return_tensors="pt"	).to(lowercase__	)

				# forward pass
				with torch.no_grad():
					SCREAMING_SNAKE_CASE   :      List[str]				     =       model(**lowercase__	)

				# verify the logits
				SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       torch.Size((1, 1000)	)
				self.assertEqual(outputs.logits.shape      ,      lowercase__	)

				SCREAMING_SNAKE_CASE   :      Dict				     =       torch.tensor([-1.0266, 0.1912, -1.2861]	).to(lowercase__	)

				self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,      lowercase__      ,      atol=1e-4	)	)






			@slow
			@require_accelerate
			@require_torch_gpu
			def 							__UpperCamelCase (			self							: Dict	)							->       List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE   :      str				     =       DeiTModel.from_pretrained(
				    "facebook/deit-base-distilled-patch16-224"      ,      torch_dtype=torch.floataa      ,      device_map="auto"	)
				SCREAMING_SNAKE_CASE   :      List[str]				     =       self.default_image_processor

				SCREAMING_SNAKE_CASE   :      str				     =       prepare_img()
				SCREAMING_SNAKE_CASE   :      str				     =       image_processor(images=lowercase__      ,      return_tensors="pt"	)
				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       inputs.pixel_values.to(lowercase__	)

				# forward pass to make sure inference works in fp16
				with torch.no_grad():
					SCREAMING_SNAKE_CASE   :      List[str]				     =       model(lowercase__	)


from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin






@dataclass
class _UpperCamelCase							(							__A   ):
			'''simple docstring'''







			lowerCamelCase__           =42






class _UpperCamelCase							(							__A							,	__A   ):
			'''simple docstring'''







			@register_to_config
			def __init__(			self							: str      ,      a							: int = 16      ,      a							: int = 88      ,      a							: Optional[int] = None      ,      a							: Optional[int] = None      ,      a							: int = 1      ,      a							: float = 0.0      ,      a							: int = 32      ,      a							: Optional[int] = None      ,      a							: bool = False      ,      a							: Optional[int] = None      ,      a							: str = "geglu"      ,      a							: bool = True      ,      a							: bool = True      ,      )							->       str:

				"""simple docstring"""

				super().__init__()
				SCREAMING_SNAKE_CASE   :      str				     =       num_attention_heads
				SCREAMING_SNAKE_CASE   :      Tuple				     =       attention_head_dim
				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       num_attention_heads * attention_head_dim

				SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       in_channels

				SCREAMING_SNAKE_CASE   :      List[Any]				     =       torch.nn.GroupNorm(num_groups=a      ,      num_channels=a      ,      eps=1e-6      ,      affine=a	)
				SCREAMING_SNAKE_CASE   :      List[Any]				     =       nn.Linear(a      ,      a	)

				# 3. Define transformers blocks
				SCREAMING_SNAKE_CASE   :      List[str]				     =       nn.ModuleList(
				    [
				        BasicTransformerBlock(
				            a      ,      a      ,      a      ,      dropout=a      ,      cross_attention_dim=a      ,      activation_fn=a      ,      attention_bias=a      ,      double_self_attention=a      ,      norm_elementwise_affine=a      ,      )
				        for d in range(a	)
				    ]	)

				SCREAMING_SNAKE_CASE   :      Optional[int]				     =       nn.Linear(a      ,      a	)






			def 							__UpperCamelCase (			self							: Any      ,      a							: int      ,      a							: Optional[int]=None      ,      a							: List[str]=None      ,      a							: List[str]=None      ,      a							: Union[str, Any]=1      ,      a							: Any=None      ,      a							: bool = True      ,      )							->       Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE     ,SCREAMING_SNAKE_CASE     ,SCREAMING_SNAKE_CASE     ,SCREAMING_SNAKE_CASE   :      Any				     =       hidden_states.shape
				SCREAMING_SNAKE_CASE   :      Any				     =       batch_frames // num_frames

				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       hidden_states

				SCREAMING_SNAKE_CASE   :      Optional[Any]				     =       hidden_states[None, :].reshape(a      ,      a      ,      a      ,      a      ,      a	)
				SCREAMING_SNAKE_CASE   :      Union[str, Any]				     =       hidden_states.permute(0      ,      2      ,      1      ,      3      ,      4	)

				SCREAMING_SNAKE_CASE   :      Tuple				     =       self.norm(a	)
				SCREAMING_SNAKE_CASE   :      List[Any]				     =       hidden_states.permute(0      ,      3      ,      4      ,      2      ,      1	).reshape(batch_size * height * width      ,      a      ,      a	)

				SCREAMING_SNAKE_CASE   :      List[Any]				     =       self.proj_in(a	)

				# 2. Blocks
				for block in self.transformer_blocks:
					SCREAMING_SNAKE_CASE   :      Optional[int]				     =       block(
					    a      ,      encoder_hidden_states=a      ,      timestep=a      ,      cross_attention_kwargs=a      ,      class_labels=a      ,      )

				# 3. Output
				SCREAMING_SNAKE_CASE   :      List[str]				     =       self.proj_out(a	)
				SCREAMING_SNAKE_CASE   :      int				     =       (
				    hidden_states[None, None, :]
				    .reshape(a      ,      a      ,      a      ,      a      ,      a	)
				    .permute(0      ,      3      ,      4      ,      1      ,      2	)
				    .contiguous()
				)
				SCREAMING_SNAKE_CASE   :      Dict				     =       hidden_states.reshape(a      ,      a      ,      a      ,      a	)

				SCREAMING_SNAKE_CASE   :      int				     =       hidden_states + residual

				if not return_dict:
					return (output,)

				return TransformerTemporalModelOutput(sample=a	)