"""simple docstring"""






from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
						from transformers import (
						    TFAutoModel,
						    TFAutoModelForCausalLM,
						    TFAutoModelForMaskedLM,
						    TFAutoModelForPreTraining,
						    TFAutoModelForQuestionAnswering,
						    TFAutoModelForSeqaSeqLM,
						    TFAutoModelForSequenceClassification,
						    TFAutoModelForTableQuestionAnswering,
						    TFAutoModelForTokenClassification,
						    TFAutoModelWithLMHead,
						    TFBertForMaskedLM,
						    TFBertForPreTraining,
						    TFBertForQuestionAnswering,
						    TFBertForSequenceClassification,
						    TFBertModel,
						    TFFunnelBaseModel,
						    TFFunnelModel,
						    TFGPTaLMHeadModel,
						    TFRobertaForMaskedLM,
						    TFTaForConditionalGeneration,
						    TFTapasForQuestionAnswering,
						)
						from transformers.models.auto.modeling_tf_auto import (
						    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
						    TF_MODEL_FOR_MASKED_LM_MAPPING,
						    TF_MODEL_FOR_PRETRAINING_MAPPING,
						    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
						    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
						    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
						    TF_MODEL_MAPPING,
						)
						from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
						from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
						from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
						from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST

class lowerCamelCase		(_SCREAMING_SNAKE_CASE      ):






					'''simple docstring'''


					a   	=		"new-model"


if is_tf_available():

						class lowerCamelCase		(_SCREAMING_SNAKE_CASE      ):






											'''simple docstring'''


											a   	=		NewModelConfig

@require_tf
class lowerCamelCase		(unittest.TestCase      ):






					'''simple docstring'''


					@slow
					def 			lowerCAmelCase_					(						self			:				Tuple		)							->       List[str]:
									SCREAMING_SNAKE_CASE__		  =				"bert-base-cased"
									SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
									self.assertIsNotNone(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)

									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained(_snake_case		)
									self.assertIsNotNone(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Any		)							->       Union[str, Any]:
									SCREAMING_SNAKE_CASE__		  =				"bert-base-cased"
									SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
									self.assertIsNotNone(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)

									SCREAMING_SNAKE_CASE__		  =				TFAutoModelForPreTraining.from_pretrained(_snake_case		)
									self.assertIsNotNone(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Optional[Any]		)							->       Optional[Any]:
									for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForCausalLM.from_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				TFAutoModelForCausalLM.from_pretrained(_snake_case   , output_loading_info=_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       str:
									for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelWithLMHead.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       Optional[int]:
									for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForMaskedLM.from_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				TFAutoModelForMaskedLM.from_pretrained(_snake_case   , output_loading_info=_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				str		)							->       Any:
									for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForSeqaSeqLM.from_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				TFAutoModelForSeqaSeqLM.from_pretrained(_snake_case   , output_loading_info=_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Tuple		)							->       Tuple:
									# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									for model_name in ["bert-base-uncased"]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForSequenceClassification.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       Optional[Any]:
									# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									for model_name in ["bert-base-uncased"]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForQuestionAnswering.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					@require_tensorflow_probability
					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       Dict:
									for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForTableQuestionAnswering.from_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				TFAutoModelForTableQuestionAnswering.from_pretrained(
													    _snake_case   , output_loading_info=_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					def 			lowerCAmelCase_					(						self			:				int		)							->       Optional[int]:
									SCREAMING_SNAKE_CASE__		  =				TFAutoModelWithLMHead.from_pretrained(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)
									self.assertEqual(model.num_parameters()   , 14410		)
									self.assertEqual(model.num_parameters(only_trainable=_snake_case		)   , 14410		)



					def 			lowerCAmelCase_					(						self			:				List[Any]		)							->       Union[str, Any]:
									SCREAMING_SNAKE_CASE__		  =				TFAutoModelWithLMHead.from_pretrained(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)
									self.assertEqual(model.num_parameters()   , 14410		)
									self.assertEqual(model.num_parameters(only_trainable=_snake_case		)   , 14410		)



					def 			lowerCAmelCase_					(						self			:				Any		)							->       List[Any]:
									# For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("sgugger/funnel-random-tiny"		)
									self.assertIsInstance(_snake_case   , _snake_case		)

									SCREAMING_SNAKE_CASE__		  =				copy.deepcopy(model.config		)
									SCREAMING_SNAKE_CASE__		  =				["FunnelBaseModel"]
									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_config(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)

									with tempfile.TemporaryDirectory() as tmp_dir:
													model.save_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					def 			lowerCAmelCase_					(						self			:				Optional[Any]		)							->       Optional[Any]:
									try:
													AutoConfig.register("new-model"   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				[
													    TFAutoModel,
													    TFAutoModelForCausalLM,
													    TFAutoModelForMaskedLM,
													    TFAutoModelForPreTraining,
													    TFAutoModelForQuestionAnswering,
													    TFAutoModelForSequenceClassification,
													    TFAutoModelForTokenClassification,
													]

													for auto_class in auto_classes:
																	with self.subTest(auto_class.__name__		):
																					# Wrong config class will raise an error
																					with self.assertRaises(_snake_case		):
																									auto_class.register(_snake_case   , _snake_case		)
																					auto_class.register(_snake_case   , _snake_case		)
																					# Trying to register something existing in the Transformers library will raise an error
																					with self.assertRaises(_snake_case		):
																									auto_class.register(_snake_case   , _snake_case		)

																					# Now that the config is registered, it can be used as any other config with the auto-API
																					SCREAMING_SNAKE_CASE__		  =				BertModelTester(self		).get_config()
																					SCREAMING_SNAKE_CASE__		  =				NewModelConfig(**tiny_config.to_dict()		)
																					SCREAMING_SNAKE_CASE__		  =				auto_class.from_config(_snake_case		)
																					self.assertIsInstance(_snake_case   , _snake_case		)

																					with tempfile.TemporaryDirectory() as tmp_dir:
																									model.save_pretrained(_snake_case		)
																									SCREAMING_SNAKE_CASE__		  =				auto_class.from_pretrained(_snake_case		)
																									self.assertIsInstance(_snake_case   , _snake_case		)

									finally:
													if "new-model" in CONFIG_MAPPING._extra_content:
																	del CONFIG_MAPPING._extra_content["new-model"]
													for mapping in (
													    TF_MODEL_MAPPING,
													    TF_MODEL_FOR_PRETRAINING_MAPPING,
													    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
													    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
													    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
													    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
													    TF_MODEL_FOR_MASKED_LM_MAPPING,
													):
																	if NewModelConfig in mapping._extra_content:
																					del mapping._extra_content[NewModelConfig]



					def 			lowerCAmelCase_					(						self			:				Optional[Any]		)							->       Optional[int]:
									with self.assertRaisesRegex(
									    _snake_case   , "bert-base is not a local folder and is not a valid model identifier"		):
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("bert-base"		)



					def 			lowerCAmelCase_					(						self			:				int		)							->       List[str]:
									with self.assertRaisesRegex(
									    _snake_case   , r"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"		):
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained(_snake_case   , revision="aaaaaa"		)



					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       int:
									with self.assertRaisesRegex(
									    _snake_case   , "hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"   , ):
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("hf-internal-testing/config-no-model"		)



					def 			lowerCAmelCase_					(						self			:				Dict		)							->       Tuple:
									with self.assertRaisesRegex(_snake_case   , "Use `from_pt=True` to load this model"		):
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("hf-internal-testing/tiny-bert-pt-only"		)



					def 			lowerCAmelCase_					(						self			:				List[Any]		)							->       Optional[Any]:
									# Make sure we have cached the model.
									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"		)
									with RequestCounter() as counter:
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"		)
													self.assertEqual(counter.get_request_count   , 0		)
													self.assertEqual(counter.head_request_count   , 1		)
													self.assertEqual(counter.other_request_count   , 0		)

									# With a sharded checkpoint
									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"		)
									with RequestCounter() as counter:
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"		)
													self.assertEqual(counter.get_request_count   , 0		)
													self.assertEqual(counter.head_request_count   , 1		)
													self.assertEqual(counter.other_request_count   , 0		)





"""simple docstring"""






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			:				Optional[Any]=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   , )							->       Union[str, Any]:
									super().__init__()
									SCREAMING_SNAKE_CASE__		  =				only_cross_attention

									SCREAMING_SNAKE_CASE__		  =				(num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
									SCREAMING_SNAKE_CASE__		  =				(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:
													SCREAMING_SNAKE_CASE__		  =				AdaLayerNorm(_snake_case   , _snake_case		)
									elif self.use_ada_layer_norm_zero:
													SCREAMING_SNAKE_CASE__		  =				AdaLayerNormZero(_snake_case   , _snake_case		)
									else:
													SCREAMING_SNAKE_CASE__		  =				nn.LayerNorm(_snake_case   , elementwise_affine=_snake_case		)
									SCREAMING_SNAKE_CASE__		  =				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.
													SCREAMING_SNAKE_CASE__		  =				(
													    AdaLayerNorm(_snake_case   , _snake_case		)
													    if self.use_ada_layer_norm
													    else nn.LayerNorm(_snake_case   , elementwise_affine=_snake_case		)
													)
													SCREAMING_SNAKE_CASE__		  =				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:
													SCREAMING_SNAKE_CASE__		  =				None
													SCREAMING_SNAKE_CASE__		  =				None

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

									# let chunk size default to None
									SCREAMING_SNAKE_CASE__		  =				None
									SCREAMING_SNAKE_CASE__		  =				0



					def 			lowerCAmelCase_					(						self			:				Tuple   , _snake_case			:				Optional[int]   , _snake_case			:				int		)							->       List[str]:
									# Sets chunk feed-forward
									SCREAMING_SNAKE_CASE__		  =				chunk_size
									SCREAMING_SNAKE_CASE__		  =				dim



					def 			lowerCAmelCase_					(						self			:				Optional[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   , )							->       Union[str, Any]:
									# Notice that normalization is always applied before the real computation in the following blocks.
									# 1. Self-Attention
									if self.use_ada_layer_norm:
													SCREAMING_SNAKE_CASE__		  =				self.norma(_snake_case   , _snake_case		)
									elif self.use_ada_layer_norm_zero:
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				self.norma(
													    _snake_case   , _snake_case   , _snake_case   , hidden_dtype=hidden_states.dtype		)
									else:
													SCREAMING_SNAKE_CASE__		  =				self.norma(_snake_case		)

									SCREAMING_SNAKE_CASE__		  =				cross_attention_kwargs if cross_attention_kwargs is not None else {}

									SCREAMING_SNAKE_CASE__		  =				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:
													SCREAMING_SNAKE_CASE__		  =				gate_msa.unsqueeze(1		) * attn_output
									SCREAMING_SNAKE_CASE__		  =				attn_output + hidden_states

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

													SCREAMING_SNAKE_CASE__		  =				self.attna(
													    _snake_case   , encoder_hidden_states=_snake_case   , attention_mask=_snake_case   , **_snake_case   , )
													SCREAMING_SNAKE_CASE__		  =				attn_output + hidden_states

									# 3. Feed-forward
									SCREAMING_SNAKE_CASE__		  =				self.norma(_snake_case		)

									if self.use_ada_layer_norm_zero:
													SCREAMING_SNAKE_CASE__		  =				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`."""		)

													SCREAMING_SNAKE_CASE__		  =				norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
													SCREAMING_SNAKE_CASE__		  =				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:
													SCREAMING_SNAKE_CASE__		  =				self.ff(_snake_case		)

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

									SCREAMING_SNAKE_CASE__		  =				ff_output + hidden_states

									return hidden_states



class lowerCamelCase		(nn.Module      ):






					'''simple docstring'''


					def __init__(						self			:				List[str]   , _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[Any]:
									super().__init__()
									SCREAMING_SNAKE_CASE__		  =				int(dim * mult		)
									SCREAMING_SNAKE_CASE__		  =				dim_out if dim_out is not None else dim

									if activation_fn == "gelu":
													SCREAMING_SNAKE_CASE__		  =				GELU(_snake_case   , _snake_case		)
									if activation_fn == "gelu-approximate":
													SCREAMING_SNAKE_CASE__		  =				GELU(_snake_case   , _snake_case   , approximate="tanh"		)
									elif activation_fn == "geglu":
													SCREAMING_SNAKE_CASE__		  =				GEGLU(_snake_case   , _snake_case		)
									elif activation_fn == "geglu-approximate":
													SCREAMING_SNAKE_CASE__		  =				ApproximateGELU(_snake_case   , _snake_case		)

									SCREAMING_SNAKE_CASE__		  =				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 			lowerCAmelCase_					(						self			:				Union[str, Any]   , _snake_case			:				List[str]		)							->       str:
									for module in self.net:
													SCREAMING_SNAKE_CASE__		  =				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"		)							->       Dict:
									super().__init__()
									SCREAMING_SNAKE_CASE__		  =				nn.Linear(_snake_case   , _snake_case		)
									SCREAMING_SNAKE_CASE__		  =				approximate



					def 			lowerCAmelCase_					(						self			:				Tuple   , _snake_case			:				Optional[Any]		)							->       Tuple:
									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 			lowerCAmelCase_					(						self			:				Union[str, Any]   , _snake_case			:				List[Any]		)							->       List[Any]:
									SCREAMING_SNAKE_CASE__		  =				self.proj(_snake_case		)
									SCREAMING_SNAKE_CASE__		  =				self.gelu(_snake_case		)
									return hidden_states



class lowerCamelCase		(nn.Module      ):






					'''simple docstring'''


					def __init__(						self			:				Optional[int]   , _snake_case			:				int   , _snake_case			:				int		)							->       Optional[int]:
									super().__init__()
									SCREAMING_SNAKE_CASE__		  =				nn.Linear(_snake_case   , dim_out * 2		)



					def 			lowerCAmelCase_					(						self			:				Optional[Any]   , _snake_case			:				Optional[Any]		)							->       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 			lowerCAmelCase_					(						self			:				Tuple   , _snake_case			:				Tuple		)							->       Any:
									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				self.proj(_snake_case		).chunk(2   , dim=-1		)
									return hidden_states * self.gelu(_snake_case		)



class lowerCamelCase		(nn.Module      ):






					'''simple docstring'''


					def __init__(						self			:				int   , _snake_case			:				int   , _snake_case			:				int		)							->       List[Any]:
									super().__init__()
									SCREAMING_SNAKE_CASE__		  =				nn.Linear(_snake_case   , _snake_case		)



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



class lowerCamelCase		(nn.Module      ):






					'''simple docstring'''


					def __init__(						self			:				Union[str, Any]   , _snake_case			:				Dict   , _snake_case			:				Any		)							->       Optional[int]:
									super().__init__()
									SCREAMING_SNAKE_CASE__		  =				nn.Embedding(_snake_case   , _snake_case		)
									SCREAMING_SNAKE_CASE__		  =				nn.SiLU()
									SCREAMING_SNAKE_CASE__		  =				nn.Linear(_snake_case   , embedding_dim * 2		)
									SCREAMING_SNAKE_CASE__		  =				nn.LayerNorm(_snake_case   , elementwise_affine=_snake_case		)



					def 			lowerCAmelCase_					(						self			:				int   , _snake_case			:				int   , _snake_case			:				Tuple		)							->       int:
									SCREAMING_SNAKE_CASE__		  =				self.linear(self.silu(self.emb(_snake_case		)		)		)
									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				torch.chunk(_snake_case   , 2		)
									SCREAMING_SNAKE_CASE__		  =				self.norm(_snake_case		) * (1 + scale) + shift
									return x



class lowerCamelCase		(nn.Module      ):






					'''simple docstring'''


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

									SCREAMING_SNAKE_CASE__		  =				CombinedTimestepLabelEmbeddings(_snake_case   , _snake_case		)

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



					def 			lowerCAmelCase_					(						self			:				Dict   , _snake_case			:				int   , _snake_case			:				List[Any]   , _snake_case			:				Any   , _snake_case			:				str=None		)							->       str:
									SCREAMING_SNAKE_CASE__		  =				self.linear(self.silu(self.emb(_snake_case   , _snake_case   , hidden_dtype=_snake_case		)		)		)
									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				emb.chunk(6   , dim=1		)
									SCREAMING_SNAKE_CASE__		  =				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			:				Optional[Any]   , _snake_case			:				int   , _snake_case			:				int   , _snake_case			:				int   , _snake_case			:				Optional[str] = None   , _snake_case			:				float = 1e-5		)							->       Optional[Any]:
									super().__init__()
									SCREAMING_SNAKE_CASE__		  =				num_groups
									SCREAMING_SNAKE_CASE__		  =				eps

									if act_fn is None:
													SCREAMING_SNAKE_CASE__		  =				None
									else:
													SCREAMING_SNAKE_CASE__		  =				get_activation(_snake_case		)

									SCREAMING_SNAKE_CASE__		  =				nn.Linear(_snake_case   , out_dim * 2		)



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

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




