






import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, 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_VIDEO_CLASSIFICATION_MAPPING,
	    TimesformerForVideoClassification,
	    TimesformerModel,
	)
	from transformers.models.timesformer.modeling_timesformer import TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from transformers import VideoMAEImageProcessor



class 				SCREAMING_SNAKE_CASE__		:





					"""simple docstring"""






					def __init__(   self	:    Dict				,						__A	:    Optional[int]				,						__A	:    Optional[Any]=1_3				,						__A	:    List[Any]=1_0				,						__A	:    List[str]=3				,						__A	:    List[str]=2				,						__A	:    Tuple=2				,						__A	:    int=True				,						__A	:    Dict=True				,						__A	:    int=3_2				,						__A	:    Dict=5				,						__A	:    Union[str, Any]=4				,						__A	:    List[Any]=3_7				,						__A	:    Any="gelu"				,						__A	:    List[str]=0.1				,						__A	:    Dict=0.1				,						__A	:    Union[str, Any]=1_0				,						__A	:    Union[str, Any]=0.0_2				,						__A	:    int="divided_space_time"				,						__A	:    Tuple=None				,						):
						snake_case__					: Union[str, Any]  =    parent
						snake_case__					: List[Any]  =    batch_size
						snake_case__					: int  =    image_size
						snake_case__					: Any  =    num_channels
						snake_case__					: List[str]  =    patch_size
						snake_case__					: int  =    num_frames
						snake_case__					: Optional[int]  =    is_training
						snake_case__					: int  =    use_labels
						snake_case__					: str  =    hidden_size
						snake_case__					: str  =    num_hidden_layers
						snake_case__					: List[str]  =    num_attention_heads
						snake_case__					: Any  =    intermediate_size
						snake_case__					: List[str]  =    hidden_act
						snake_case__					: Union[str, Any]  =    hidden_dropout_prob
						snake_case__					: Union[str, Any]  =    attention_probs_dropout_prob
						snake_case__					: Any  =    attention_type
						snake_case__					: Tuple  =    initializer_range
						snake_case__					: Union[str, Any]  =    scope
						snake_case__					: Tuple  =    num_labels

						# in TimeSformer, the number of spatial tokens equals num_frames * num_patches per frame + 1 CLS token
						snake_case__					: List[Any]  =    (image_size // patch_size) ** 2
						snake_case__					: Dict  =    (num_frames) * self.num_patches_per_frame + 1



					def 			_lowercase	(   self	:    str			):
						snake_case__					: int  =    floats_tensor(
						    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]			)

						snake_case__					: Optional[Any]  =    None
						if self.use_labels:
							snake_case__					: Dict  =    ids_tensor([self.batch_size]				,						self.num_labels			)

						snake_case__					: Tuple  =    self.get_config()

						return config, pixel_values, labels



					def 			_lowercase	(   self	:    List[str]			):
						snake_case__					: Dict  =    TimesformerConfig(
						    image_size=self.image_size				,						patch_size=self.patch_size				,						num_channels=self.num_channels				,						num_frames=self.num_frames				,						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				,						initializer_range=self.initializer_range				,						attention_type=self.attention_type				,						)
						snake_case__					: Tuple  =    self.num_labels
						return config



					def 			_lowercase	(   self	:    Union[str, Any]				,						__A	:    Any				,						__A	:    Union[str, Any]				,						__A	:    Union[str, Any]			):
						snake_case__					: Any  =    TimesformerModel(config=snake_case__			)
						model.to(snake_case__			)
						model.eval()
						snake_case__					: Union[str, Any]  =    model(snake_case__			)
						self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, self.seq_length, self.hidden_size)			)



					def 			_lowercase	(   self	:    str				,						__A	:    Union[str, Any]				,						__A	:    Optional[Any]				,						__A	:    int			):
						snake_case__					: Union[str, Any]  =    TimesformerForVideoClassification(snake_case__			)
						model.to(snake_case__			)
						model.eval()

						snake_case__					: str  =    model(snake_case__			)

						# verify the logits shape
						snake_case__					: Optional[Any]  =    torch.Size((self.batch_size, self.num_labels)			)
						self.parent.assertEqual(result.logits.shape				,						snake_case__			)






					def 			_lowercase	(   self	:    Tuple			):
						snake_case__					: List[Any]  =    self.prepare_config_and_inputs()
						snake_case__,		snake_case__,		snake_case__					: Optional[int]  =    config_and_inputs
						snake_case__					: int  =    {"pixel_values": pixel_values}
						return config, inputs_dict



@require_torch
class 				SCREAMING_SNAKE_CASE__		(							SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  unittest.TestCase		):





					"""simple docstring"""






					a_      =						(TimesformerModel, TimesformerForVideoClassification) if is_torch_available() else ()
					a_      =						(
					    {"feature-extraction": TimesformerModel, "video-classification": TimesformerForVideoClassification}
					    if is_torch_available()
					    else {}
					)

					a_      =						False
					a_      =						False
					a_      =						False
					a_      =						False



					def 			_lowercase	(   self	:    List[str]			):
						snake_case__					: Tuple  =    TimesformerModelTester(self			)
						snake_case__					: List[str]  =    ConfigTester(
						    self				,						config_class=snake_case__				,						has_text_modality=snake_case__				,						hidden_size=3_7			)



					def 			_lowercase	(   self	:    List[str]				,						__A	:    int				,						__A	:    Union[str, Any]				,						__A	:    List[str]=False			):
						snake_case__					: int  =    copy.deepcopy(snake_case__			)

						if return_labels:
							if model_class in get_values(snake_case__			):
								snake_case__					: Union[str, Any]  =    torch.zeros(
								    self.model_tester.batch_size				,						dtype=torch.long				,						device=snake_case__			)

						return inputs_dict



					def 			_lowercase	(   self	:    Union[str, Any]			):
						self.config_tester.run_common_tests()



					@unittest.skip(reason="TimeSformer does not use inputs_embeds"			)
					def 			_lowercase	(   self	:    Any			):
						pass



					def 			_lowercase	(   self	:    Optional[Any]			):
						snake_case__,		snake_case__					: Union[str, Any]  =    self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
							snake_case__					: Any  =    model_class(snake_case__			)
							self.assertIsInstance(model.get_input_embeddings()				,						(nn.Module)			)
							snake_case__					: List[str]  =    model.get_output_embeddings()
							self.assertTrue(x is None or isinstance(snake_case__				,						nn.Linear			)			)



					def 			_lowercase	(   self	:    Optional[Any]			):
						snake_case__,		snake_case__					: Optional[int]  =    self.model_tester.prepare_config_and_inputs_for_common()

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

							snake_case__					: List[Any]  =    ["pixel_values"]
							self.assertListEqual(arg_names[:1]				,						snake_case__			)



					def 			_lowercase	(   self	:    Any			):
						snake_case__					: Tuple  =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*snake_case__			)



					def 			_lowercase	(   self	:    Any			):
						snake_case__					: List[Any]  =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_video_classification(*snake_case__			)



					@slow
					def 			_lowercase	(   self	:    int			):
						for model_name in TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							snake_case__					: int  =    TimesformerModel.from_pretrained(snake_case__			)
							self.assertIsNotNone(snake_case__			)



					def 			_lowercase	(   self	:    Union[str, Any]			):
						if not self.has_attentions:
							pass

						else:
							snake_case__,		snake_case__					: Dict  =    self.model_tester.prepare_config_and_inputs_for_common()
							snake_case__					: Any  =    True

							for model_class in self.all_model_classes:
								snake_case__					: int  =    self.model_tester.seq_length
								snake_case__					: List[str]  =    self.model_tester.num_frames

								snake_case__					: Dict  =    True
								snake_case__					: Tuple  =    False
								snake_case__					: int  =    True
								snake_case__					: Optional[Any]  =    model_class(snake_case__			)
								model.to(snake_case__			)
								model.eval()
								with torch.no_grad():
									snake_case__					: Union[str, Any]  =    model(**self._prepare_for_class(snake_case__				,						snake_case__			)			)
								snake_case__					: Dict  =    outputs.attentions
								self.assertEqual(len(snake_case__			)				,						self.model_tester.num_hidden_layers			)

								# check that output_attentions also work using config
								del inputs_dict["output_attentions"]
								snake_case__					: str  =    True
								snake_case__					: Optional[Any]  =    model_class(snake_case__			)
								model.to(snake_case__			)
								model.eval()
								with torch.no_grad():
									snake_case__					: List[str]  =    model(**self._prepare_for_class(snake_case__				,						snake_case__			)			)
								snake_case__					: Dict  =    outputs.attentions
								self.assertEqual(len(snake_case__			)				,						self.model_tester.num_hidden_layers			)

								# attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
								self.assertListEqual(
								    list(attentions[0].shape[-3:]			)				,						[self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1]				,						)
								snake_case__					: Any  =    len(snake_case__			)

								# Check attention is always last and order is fine
								snake_case__					: List[str]  =    True
								snake_case__					: Dict  =    True
								snake_case__					: Optional[int]  =    model_class(snake_case__			)
								model.to(snake_case__			)
								model.eval()
								with torch.no_grad():
									snake_case__					: Dict  =    model(**self._prepare_for_class(snake_case__				,						snake_case__			)			)

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

								snake_case__					: Optional[int]  =    outputs.attentions

								self.assertEqual(len(snake_case__			)				,						self.model_tester.num_hidden_layers			)

								# attentions has shape (batch_size x num_frames) x num_heads x (num_patches per frame + 1) x (num_patches per frame + 1)
								self.assertListEqual(
								    list(self_attentions[0].shape[-3:]			)				,						[self.model_tester.num_attention_heads, seq_len // num_frames + 1, seq_len // num_frames + 1]				,						)






					def 			_lowercase	(   self	:    List[Any]			):
						def check_hidden_states_output(__A	:    Any				,						__A	:    Optional[Any]				,						__A	:    Tuple			):
							snake_case__					: List[str]  =    model_class(snake_case__			)
							model.to(snake_case__			)
							model.eval()

							with torch.no_grad():
								snake_case__					: Optional[int]  =    model(**self._prepare_for_class(snake_case__				,						snake_case__			)			)

							snake_case__					: Union[str, Any]  =    outputs.hidden_states
							snake_case__					: Any  =    self.model_tester.num_hidden_layers + 1
							self.assertEqual(len(snake_case__			)				,						snake_case__			)

							snake_case__					: int  =    self.model_tester.seq_length

							self.assertListEqual(
							    list(hidden_states[0].shape[-2:]			)				,						[seq_length, self.model_tester.hidden_size]				,						)

						snake_case__,		snake_case__					: Any  =    self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
							snake_case__					: Optional[Any]  =    True
							check_hidden_states_output(snake_case__				,						snake_case__				,						snake_case__			)

							# check that output_hidden_states also work using config
							del inputs_dict["output_hidden_states"]
							snake_case__					: Tuple  =    True

							check_hidden_states_output(snake_case__				,						snake_case__				,						snake_case__			)







def 	SCREAMING_SNAKE_CASE							(					):
	snake_case__					: int  =    hf_hub_download(
	    repo_id="hf-internal-testing/spaghetti-video"	,    filename="eating_spaghetti.npy"	,    repo_type="dataset"  )
	snake_case__					: Union[str, Any]  =    np.load(snake_case_  )
	return list(snake_case_  )





@require_torch
@require_vision
class 				SCREAMING_SNAKE_CASE__		(							unittest.TestCase		):





					"""simple docstring"""






					@cached_property
					def 			_lowercase	(   self	:    Any			):
						# logits were tested with a different mean and std, so we use the same here
						return (
						    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]				,						image_std=[0.5, 0.5, 0.5]			)
						    if is_vision_available()
						    else None
						)






					@slow
					def 			_lowercase	(   self	:    int			):
						snake_case__					: List[Any]  =    TimesformerForVideoClassification.from_pretrained("facebook/timesformer-base-finetuned-k400"			).to(
						    snake_case__			)

						snake_case__					: Dict  =    self.default_image_processor
						snake_case__					: int  =    prepare_video()
						snake_case__					: Any  =    image_processor(video[:8]				,						return_tensors="pt"			).to(snake_case__			)

						# forward pass
						with torch.no_grad():
							snake_case__					: Union[str, Any]  =    model(**snake_case__			)

						# verify the logits
						snake_case__					: Union[str, Any]  =    torch.Size((1, 4_0_0)			)
						self.assertEqual(outputs.logits.shape				,						snake_case__			)

						snake_case__					: int  =    torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]			).to(snake_case__			)

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




from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings





@maybe_allow_in_graph
class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Dict     ,  snake_case__     :     int     ,  snake_case__     :     int     ,  snake_case__     :     int     ,  snake_case__     :     Tuple=0.0     ,  snake_case__     :     Optional[int] = None     ,  snake_case__     :     str = "geglu"     ,  snake_case__     :     Optional[int] = None     ,  snake_case__     :     bool = False     ,  snake_case__     :     bool = False     ,  snake_case__     :     bool = False     ,  snake_case__     :     bool = False     ,  snake_case__     :     bool = True     ,  snake_case__     :     str = "layer_norm"     ,  snake_case__     :     bool = False     ,  )      ->	Dict:
				super().__init__()
				_lowerCamelCase				   =  only_cross_attention

				_lowerCamelCase				   =  (num_embeds_ada_norm is not None) and norm_type == 'ada_norm_zero'
				_lowerCamelCase				   =  (num_embeds_ada_norm is not None) and norm_type == 'ada_norm'

				if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
						raise ValueError(
						    f"""`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"""
						    f""" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."""		)

				# Define 3 blocks. Each block has its own normalization layer.
				# 1. Self-Attn
				if self.use_ada_layer_norm:
						_lowerCamelCase				   =  AdaLayerNorm(snake_case__     ,  snake_case__		)
				elif self.use_ada_layer_norm_zero:
						_lowerCamelCase				   =  AdaLayerNormZero(snake_case__     ,  snake_case__		)
				else:
						_lowerCamelCase				   =  nn.LayerNorm(snake_case__     ,  elementwise_affine=snake_case__		)
				_lowerCamelCase				   =  Attention(
				    query_dim=snake_case__     ,  heads=snake_case__     ,  dim_head=snake_case__     ,  dropout=snake_case__     ,  bias=snake_case__     ,  cross_attention_dim=cross_attention_dim if only_cross_attention else None     ,  upcast_attention=snake_case__     ,  )

				# 2. Cross-Attn
				if cross_attention_dim is not None or double_self_attention:
						# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
						# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
						# the second cross attention block.
						_lowerCamelCase				   =  (
						    AdaLayerNorm(snake_case__     ,  snake_case__		)
						    if self.use_ada_layer_norm
						    else nn.LayerNorm(snake_case__     ,  elementwise_affine=snake_case__		)
						)
						_lowerCamelCase				   =  Attention(
						    query_dim=snake_case__     ,  cross_attention_dim=cross_attention_dim if not double_self_attention else None     ,  heads=snake_case__     ,  dim_head=snake_case__     ,  dropout=snake_case__     ,  bias=snake_case__     ,  upcast_attention=snake_case__     ,  )  # is self-attn if encoder_hidden_states is none
				else:
						_lowerCamelCase				   =  None
						_lowerCamelCase				   =  None

				# 3. Feed-forward
				_lowerCamelCase				   =  nn.LayerNorm(snake_case__     ,  elementwise_affine=snake_case__		)
				_lowerCamelCase				   =  FeedForward(snake_case__     ,  dropout=snake_case__     ,  activation_fn=snake_case__     ,  final_dropout=snake_case__		)

				# let chunk size default to None
				_lowerCamelCase				   =  None
				_lowerCamelCase				   =  0
		def   _snake_case       (       self     :     str     ,  snake_case__     :     Optional[int]     ,  snake_case__     :     int		)      ->	Optional[Any]:
				# Sets chunk feed-forward
				_lowerCamelCase				   =  chunk_size
				_lowerCamelCase				   =  dim



		def   _snake_case       (       self     :     Union[str, Any]     ,  snake_case__     :     torch.FloatTensor     ,  snake_case__     :     Optional[torch.FloatTensor] = None     ,  snake_case__     :     Optional[torch.FloatTensor] = None     ,  snake_case__     :     Optional[torch.FloatTensor] = None     ,  snake_case__     :     Optional[torch.LongTensor] = None     ,  snake_case__     :     Dict[str, Any] = None     ,  snake_case__     :     Optional[torch.LongTensor] = None     ,  )      ->	List[Any]:
				# Notice that normalization is always applied before the real computation in the following blocks.
				# 1. Self-Attention
				if self.use_ada_layer_norm:
						_lowerCamelCase				   =  self.norma(snake_case__     ,  snake_case__		)
				elif self.use_ada_layer_norm_zero:
						_lowerCamelCase   ,							_lowerCamelCase   ,							_lowerCamelCase   ,							_lowerCamelCase   ,							_lowerCamelCase				   =  self.norma(
						    snake_case__     ,  snake_case__     ,  snake_case__     ,  hidden_dtype=hidden_states.dtype		)
				else:
						_lowerCamelCase				   =  self.norma(snake_case__		)

				_lowerCamelCase				   =  cross_attention_kwargs if cross_attention_kwargs is not None else {}

				_lowerCamelCase				   =  self.attna(
				    snake_case__     ,  encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None     ,  attention_mask=snake_case__     ,  **snake_case__     ,  )
				if self.use_ada_layer_norm_zero:
						_lowerCamelCase				   =  gate_msa.unsqueeze(1		) * attn_output
				_lowerCamelCase				   =  attn_output + hidden_states

				# 2. Cross-Attention
				if self.attna is not None:
						_lowerCamelCase				   =  (
						    self.norma(snake_case__     ,  snake_case__		) if self.use_ada_layer_norm else self.norma(snake_case__		)
						)

						_lowerCamelCase				   =  self.attna(
						    snake_case__     ,  encoder_hidden_states=snake_case__     ,  attention_mask=snake_case__     ,  **snake_case__     ,  )
						_lowerCamelCase				   =  attn_output + hidden_states

				# 3. Feed-forward
				_lowerCamelCase				   =  self.norma(snake_case__		)

				if self.use_ada_layer_norm_zero:
						_lowerCamelCase				   =  norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

				if self._chunk_size is not None:
						# "feed_forward_chunk_size" can be used to save memory
						if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
								raise ValueError(
								    f"""`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."""		)

						_lowerCamelCase				   =  norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
						_lowerCamelCase				   =  torch.cat(
						    [self.ff(snake_case__		) for hid_slice in norm_hidden_states.chunk(snake_case__     ,  dim=self._chunk_dim		)]     ,  dim=self._chunk_dim     ,  )
				else:
						_lowerCamelCase				   =  self.ff(snake_case__		)

				if self.use_ada_layer_norm_zero:
						_lowerCamelCase				   =  gate_mlp.unsqueeze(1		) * ff_output

				_lowerCamelCase				   =  ff_output + hidden_states

				return hidden_states





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Union[str, Any]     ,  snake_case__     :     int     ,  snake_case__     :     Optional[int] = None     ,  snake_case__     :     int = 4     ,  snake_case__     :     float = 0.0     ,  snake_case__     :     str = "geglu"     ,  snake_case__     :     bool = False     ,  )      ->	Optional[int]:
				super().__init__()
				_lowerCamelCase				   =  int(dim * mult		)
				_lowerCamelCase				   =  dim_out if dim_out is not None else dim

				if activation_fn == "gelu":
						_lowerCamelCase				   =  GELU(snake_case__     ,  snake_case__		)
				if activation_fn == "gelu-approximate":
						_lowerCamelCase				   =  GELU(snake_case__     ,  snake_case__     ,  approximate='tanh'		)
				elif activation_fn == "geglu":
						_lowerCamelCase				   =  GEGLU(snake_case__     ,  snake_case__		)
				elif activation_fn == "geglu-approximate":
						_lowerCamelCase				   =  ApproximateGELU(snake_case__     ,  snake_case__		)

				_lowerCamelCase				   =  nn.ModuleList([]		)
				# project in
				self.net.append(snake_case__		)
				# project dropout
				self.net.append(nn.Dropout(snake_case__		)		)
				# project out
				self.net.append(nn.Linear(snake_case__     ,  snake_case__		)		)
				# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
				if final_dropout:
						self.net.append(nn.Dropout(snake_case__		)		)



		def   _snake_case       (       self     :     Tuple     ,  snake_case__     :     List[Any]		)      ->	Optional[int]:
				for module in self.net:
						_lowerCamelCase				   =  module(snake_case__		)
				return hidden_states





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     str     ,  snake_case__     :     int     ,  snake_case__     :     int     ,  snake_case__     :     str = "none"		)      ->	Any:
				super().__init__()
				_lowerCamelCase				   =  nn.Linear(snake_case__     ,  snake_case__		)
				_lowerCamelCase				   =  approximate
		def   _snake_case       (       self     :     List[Any]     ,  snake_case__     :     List[Any]		)      ->	int:
				if gate.device.type != "mps":
						return F.gelu(snake_case__     ,  approximate=self.approximate		)
				# mps: gelu is not implemented for float16
				return F.gelu(gate.to(dtype=torch.floataa		)     ,  approximate=self.approximate		).to(dtype=gate.dtype		)



		def   _snake_case       (       self     :     Optional[int]     ,  snake_case__     :     Union[str, Any]		)      ->	str:
				_lowerCamelCase				   =  self.proj(snake_case__		)
				_lowerCamelCase				   =  self.gelu(snake_case__		)
				return hidden_states





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Dict     ,  snake_case__     :     int     ,  snake_case__     :     int		)      ->	List[Any]:
				super().__init__()
				_lowerCamelCase				   =  nn.Linear(snake_case__     ,  dim_out * 2		)
		def   _snake_case       (       self     :     Optional[Any]     ,  snake_case__     :     Optional[int]		)      ->	Optional[Any]:
				if gate.device.type != "mps":
						return F.gelu(snake_case__		)
				# mps: gelu is not implemented for float16
				return F.gelu(gate.to(dtype=torch.floataa		)		).to(dtype=gate.dtype		)



		def   _snake_case       (       self     :     List[str]     ,  snake_case__     :     Optional[Any]		)      ->	Optional[int]:
				_lowerCamelCase   ,							_lowerCamelCase				   =  self.proj(snake_case__		).chunk(2     ,  dim=-1		)
				return hidden_states * self.gelu(snake_case__		)





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Dict     ,  snake_case__     :     int     ,  snake_case__     :     int		)      ->	Tuple:
				super().__init__()
				_lowerCamelCase				   =  nn.Linear(snake_case__     ,  snake_case__		)



		def   _snake_case       (       self     :     Optional[Any]     ,  snake_case__     :     List[Any]		)      ->	Optional[int]:
				_lowerCamelCase				   =  self.proj(snake_case__		)
				return x * torch.sigmoid(1.702 * x		)





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Optional[Any]     ,  snake_case__     :     Tuple     ,  snake_case__     :     Any		)      ->	Dict:
				super().__init__()
				_lowerCamelCase				   =  nn.Embedding(snake_case__     ,  snake_case__		)
				_lowerCamelCase				   =  nn.SiLU()
				_lowerCamelCase				   =  nn.Linear(snake_case__     ,  embedding_dim * 2		)
				_lowerCamelCase				   =  nn.LayerNorm(snake_case__     ,  elementwise_affine=snake_case__		)



		def   _snake_case       (       self     :     Tuple     ,  snake_case__     :     str     ,  snake_case__     :     Dict		)      ->	Optional[int]:
				_lowerCamelCase				   =  self.linear(self.silu(self.emb(snake_case__		)		)		)
				_lowerCamelCase   ,							_lowerCamelCase				   =  torch.chunk(snake_case__     ,  2		)
				_lowerCamelCase				   =  self.norm(snake_case__		) * (1 + scale) + shift
				return x





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Any     ,  snake_case__     :     Optional[Any]     ,  snake_case__     :     Union[str, Any]		)      ->	Dict:
				super().__init__()

				_lowerCamelCase				   =  CombinedTimestepLabelEmbeddings(snake_case__     ,  snake_case__		)

				_lowerCamelCase				   =  nn.SiLU()
				_lowerCamelCase				   =  nn.Linear(snake_case__     ,  6 * embedding_dim     ,  bias=snake_case__		)
				_lowerCamelCase				   =  nn.LayerNorm(snake_case__     ,  elementwise_affine=snake_case__     ,  eps=1e-6		)



		def   _snake_case       (       self     :     Dict     ,  snake_case__     :     Tuple     ,  snake_case__     :     str     ,  snake_case__     :     str     ,  snake_case__     :     List[Any]=None		)      ->	int:
				_lowerCamelCase				   =  self.linear(self.silu(self.emb(snake_case__     ,  snake_case__     ,  hidden_dtype=snake_case__		)		)		)
				_lowerCamelCase   ,							_lowerCamelCase   ,							_lowerCamelCase   ,							_lowerCamelCase   ,							_lowerCamelCase   ,							_lowerCamelCase				   =  emb.chunk(6     ,  dim=1		)
				_lowerCamelCase				   =  self.norm(snake_case__		) * (1 + scale_msa[:, None]) + shift_msa[:, None]
				return x, gate_msa, shift_mlp, scale_mlp, gate_mlp





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Tuple     ,  snake_case__     :     int     ,  snake_case__     :     int     ,  snake_case__     :     int     ,  snake_case__     :     Optional[str] = None     ,  snake_case__     :     float = 1e-5		)      ->	int:
				super().__init__()
				_lowerCamelCase				   =  num_groups
				_lowerCamelCase				   =  eps

				if act_fn is None:
						_lowerCamelCase				   =  None
				else:
						_lowerCamelCase				   =  get_activation(snake_case__		)

				_lowerCamelCase				   =  nn.Linear(snake_case__     ,  out_dim * 2		)



		def   _snake_case       (       self     :     Union[str, Any]     ,  snake_case__     :     int     ,  snake_case__     :     Optional[Any]		)      ->	Union[str, Any]:
				if self.act:
						_lowerCamelCase				   =  self.act(snake_case__		)
				_lowerCamelCase				   =  self.linear(snake_case__		)
				_lowerCamelCase				   =  emb[:, :, None, None]
				_lowerCamelCase   ,							_lowerCamelCase				   =  emb.chunk(2     ,  dim=1		)

				_lowerCamelCase				   =  F.group_norm(snake_case__     ,  self.num_groups     ,  eps=self.eps		)
				_lowerCamelCase				   =  x * (1 + scale) + shift
				return x