"""simple docstring"""





import json
import os
import shutil
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoConfig, BertConfig, GPTaConfig
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / "utils"))

from test_module.custom_configuration import CustomConfig  # noqa E402


_UpperCamelCase   :						Any      =    {
    "return_dict": False,
    "output_hidden_states": True,
    "output_attentions": True,
    "torchscript": True,
    "torch_dtype": "float16",
    "use_bfloat16": True,
    "tf_legacy_loss": True,
    "pruned_heads": {"a": 1},
    "tie_word_embeddings": False,
    "is_decoder": True,
    "cross_attention_hidden_size": 1_28,
    "add_cross_attention": True,
    "tie_encoder_decoder": True,
    "max_length": 50,
    "min_length": 3,
    "do_sample": True,
    "early_stopping": True,
    "num_beams": 3,
    "num_beam_groups": 3,
    "diversity_penalty": 0.5,
    "temperature": 2.0,
    "top_k": 10,
    "top_p": 0.7,
    "typical_p": 0.2,
    "repetition_penalty": 0.8,
    "length_penalty": 0.8,
    "no_repeat_ngram_size": 5,
    "encoder_no_repeat_ngram_size": 5,
    "bad_words_ids": [1, 2, 3],
    "num_return_sequences": 3,
    "chunk_size_feed_forward": 5,
    "output_scores": True,
    "return_dict_in_generate": True,
    "forced_bos_token_id": 2,
    "forced_eos_token_id": 3,
    "remove_invalid_values": True,
    "architectures": ["BertModel"],
    "finetuning_task": "translation",
    "id2label": {0: "label"},
    "label2id": {"label": "0"},
    "tokenizer_class": "BertTokenizerFast",
    "prefix": "prefix",
    "bos_token_id": 6,
    "pad_token_id": 7,
    "eos_token_id": 8,
    "sep_token_id": 9,
    "decoder_start_token_id": 10,
    "exponential_decay_length_penalty": (5, 1.0_1),
    "suppress_tokens": [0, 1],
    "begin_suppress_tokens": 2,
    "task_specific_params": {"translation": "some_params"},
    "problem_type": "regression",
}


@is_staging_test
class      UpperCAmelCase_	(    unittest.TestCase):


      @classmethod
      def 		_UpperCAmelCase       (   cls							)       ->					Dict:
             lowercase__			:		Union[str, Any]       =	TOKEN
             HfFolder.save_token(a							)


      @classmethod
      def 		_UpperCAmelCase       (   cls							)       ->					Union[str, Any]:
             try:
                    delete_repo(token=cls._token		,			repo_id='test-config'							)
             except HTTPError:
                    pass

             try:
                    delete_repo(token=cls._token		,			repo_id='valid_org/test-config-org'							)
             except HTTPError:
                    pass

             try:
                    delete_repo(token=cls._token		,			repo_id='test-dynamic-config'							)
             except HTTPError:
                    pass


      def 		_UpperCAmelCase       (   self							)       ->					List[Any]:
             lowercase__			:		List[Any]       =	BertConfig(
                 vocab_size=9_9		,			hidden_size=3_2		,			num_hidden_layers=5		,			num_attention_heads=4		,			intermediate_size=3_7							)
             config.push_to_hub('test-config'		,			use_auth_token=self._token							)

             lowercase__			:		Any       =	BertConfig.from_pretrained(f"""{USER}/test-config"""							)
             for k, v in config.to_dict().items():
                    if k != "transformers_version":
                           self.assertEqual(a		,			getattr(a		,			a							)							)

        # Reset repo
             delete_repo(token=self._token		,			repo_id='test-config'							)

             # Push to hub via save_pretrained
             with tempfile.TemporaryDirectory() as tmp_dir:
                    config.save_pretrained(a		,			repo_id='test-config'		,			push_to_hub=a		,			use_auth_token=self._token							)

             lowercase__			:		Union[str, Any]       =	BertConfig.from_pretrained(f"""{USER}/test-config"""							)
             for k, v in config.to_dict().items():
                    if k != "transformers_version":
                           self.assertEqual(a		,			getattr(a		,			a							)							)


      def 		_UpperCAmelCase       (   self							)       ->					List[str]:
             lowercase__			:		str       =	BertConfig(
                 vocab_size=9_9		,			hidden_size=3_2		,			num_hidden_layers=5		,			num_attention_heads=4		,			intermediate_size=3_7							)
             config.push_to_hub('valid_org/test-config-org'		,			use_auth_token=self._token							)

             lowercase__			:		Tuple       =	BertConfig.from_pretrained('valid_org/test-config-org'							)
             for k, v in config.to_dict().items():
                    if k != "transformers_version":
                           self.assertEqual(a		,			getattr(a		,			a							)							)

        # Reset repo
             delete_repo(token=self._token		,			repo_id='valid_org/test-config-org'							)

             # Push to hub via save_pretrained
             with tempfile.TemporaryDirectory() as tmp_dir:
                    config.save_pretrained(
                        a		,			repo_id='valid_org/test-config-org'		,			push_to_hub=a		,			use_auth_token=self._token							)

             lowercase__			:		Any       =	BertConfig.from_pretrained('valid_org/test-config-org'							)
             for k, v in config.to_dict().items():
                    if k != "transformers_version":
                           self.assertEqual(a		,			getattr(a		,			a							)							)


      def 		_UpperCAmelCase       (   self							)       ->					int:
             CustomConfig.register_for_auto_class()
             lowercase__			:		Optional[Any]       =	CustomConfig(attribute=4_2							)

             config.push_to_hub('test-dynamic-config'		,			use_auth_token=self._token							)

             # This has added the proper auto_map field to the config
             self.assertDictEqual(config.auto_map		,			{'AutoConfig': 'custom_configuration.CustomConfig'}							)

             lowercase__			:		int       =	AutoConfig.from_pretrained(f"""{USER}/test-dynamic-config"""		,			trust_remote_code=a							)
             # Can't make an isinstance check because the new_config is from the FakeConfig class of a dynamic module
             self.assertEqual(new_config.__class__.__name__		,			'CustomConfig'							)
             self.assertEqual(new_config.attribute		,			4_2							)





class      UpperCAmelCase_	(    unittest.TestCase):


      def 		_UpperCAmelCase       (   self							)       ->					Dict:
             lowercase__			:		List[Any]       =	GPTaConfig()

             # attempt to modify each of int/float/bool/str config records and verify they were updated
             lowercase__			:		Optional[int]       =	c.n_embd + 1  # int
             lowercase__			:		Optional[int]       =	c.resid_pdrop + 1.0  # float
             lowercase__			:		List[Any]       =	not c.scale_attn_weights  # bool
             lowercase__			:		Union[str, Any]       =	c.summary_type + 'foo'  # str
             c.update_from_string(
                 f"""n_embd={n_embd},resid_pdrop={resid_pdrop},scale_attn_weights={scale_attn_weights},summary_type={summary_type}"""							)
             self.assertEqual(a		,			c.n_embd		,			'mismatch for key: n_embd'							)
             self.assertEqual(a		,			c.resid_pdrop		,			'mismatch for key: resid_pdrop'							)
             self.assertEqual(a		,			c.scale_attn_weights		,			'mismatch for key: scale_attn_weights'							)
             self.assertEqual(a		,			c.summary_type		,			'mismatch for key: summary_type'							)


      def 		_UpperCAmelCase       (   self							)       ->					Union[str, Any]:
             lowercase__			:		List[Any]       =	PretrainedConfig()
             lowercase__			:		Dict       =	[key for key in base_config.__dict__ if key not in config_common_kwargs]
             # If this part of the test fails, you have arguments to addin config_common_kwargs above.
             self.assertListEqual(
                 a		,			['is_encoder_decoder', '_name_or_path', '_commit_hash', 'transformers_version']							)
             lowercase__			:		Dict       =	[key for key, value in config_common_kwargs.items() if value == getattr(a		,			a							)]
             if len(a							) > 0:
                    raise ValueError(
                        'The following keys are set with the default values in'
                        ' `test_configuration_common.config_common_kwargs` pick another value for them:'
                        f""" {", ".join(a							)}."""							)


      def 		_UpperCAmelCase       (   self							)       ->					Optional[int]:
             with self.assertRaises(a							):
                    # config is in subfolder, the following should not work without specifying the subfolder
                    lowercase__			:		Union[str, Any]       =	BertConfig.from_pretrained('hf-internal-testing/tiny-random-bert-subfolder'							)

             lowercase__			:		str       =	BertConfig.from_pretrained('hf-internal-testing/tiny-random-bert-subfolder'		,			subfolder='bert'							)

             self.assertIsNotNone(a							)


      def 		_UpperCAmelCase       (   self							)       ->					Union[str, Any]:
             # A mock response for an HTTP head request to emulate server down
             lowercase__			:		List[Any]       =	mock.Mock()
             lowercase__			:		Optional[int]       =	5_0_0
             lowercase__			:		Any       =	{}
             lowercase__			:		Union[str, Any]       =	HTTPError
             lowercase__			:		Any       =	{}

             # Download this model to make sure it's in the cache.
             lowercase__			:		Optional[Any]       =	BertConfig.from_pretrained('hf-internal-testing/tiny-random-bert'							)

             # Under the mock environment we get a 500 error when trying to reach the model.
             with mock.patch('requests.Session.request'		,			return_value=a							) as mock_head:
                    lowercase__			:		Any       =	BertConfig.from_pretrained('hf-internal-testing/tiny-random-bert'							)
                    # This check we did call the fake head request
                    mock_head.assert_called()


      def 		_UpperCAmelCase       (   self							)       ->					Union[str, Any]:
             # This test is for deprecated behavior and can be removed in v5
             lowercase__			:		Union[str, Any]       =	BertConfig.from_pretrained(
                 'https://huggingface.co/hf-internal-testing/tiny-random-bert/resolve/main/config.json'							)


      def 		_UpperCAmelCase       (   self							)       ->					Dict:
             lowercase__			:		str       =	AutoConfig.from_pretrained('bert-base-cased'							)
             lowercase__			:		List[str]       =	['config.4.0.0.json']

             with tempfile.TemporaryDirectory() as tmp_dir:
                    configuration.save_pretrained(a							)
                    lowercase__			:		str       =	2
                    json.dump(configuration.to_dict()		,			open(os.path.join(a		,			'config.4.0.0.json'							)		,			'w'							)							)

                    # This should pick the new configuration file as the version of Transformers is > 4.0.0
                    lowercase__			:		Tuple       =	AutoConfig.from_pretrained(a							)
                    self.assertEqual(new_configuration.hidden_size		,			2							)

                    # Will need to be adjusted if we reach v42 and this test is still here.
                    # Should pick the old configuration file as the version of Transformers is < 4.42.0
                    lowercase__			:		Any       =	['config.42.0.0.json']
                    lowercase__			:		Optional[int]       =	7_6_8
                    configuration.save_pretrained(a							)
                    shutil.move(os.path.join(a		,			'config.4.0.0.json'							)		,			os.path.join(a		,			'config.42.0.0.json'							)							)
                    lowercase__			:		Dict       =	AutoConfig.from_pretrained(a							)
                    self.assertEqual(new_configuration.hidden_size		,			7_6_8							)


      def 		_UpperCAmelCase       (   self							)       ->					List[Any]:
             # This repo has two configuration files, one for v4.0.0 and above with a different hidden size.
             lowercase__			:		Dict       =	'hf-internal-testing/test-two-configs'

             import transformers as new_transformers

             lowercase__			:		str       =	'v4.0.0'
             lowercase__    ,      lowercase__			:		str       =	new_transformers.models.auto.AutoConfig.from_pretrained(
                 a		,			return_unused_kwargs=a							)
             self.assertEqual(new_configuration.hidden_size		,			2							)
             # This checks `_configuration_file` ia not kept in the kwargs by mistake.
             self.assertDictEqual(a		,			{}							)

             # Testing an older version by monkey-patching the version in the module it's used.
             import transformers as old_transformers

             lowercase__			:		Dict       =	'v3.0.0'
             lowercase__			:		List[str]       =	old_transformers.models.auto.AutoConfig.from_pretrained(a							)
             self.assertEqual(old_configuration.hidden_size		,			7_6_8							)






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 SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__=0.0,		lowerCamelCase__ = None,		lowerCamelCase__ = "geglu",		lowerCamelCase__ = None,		lowerCamelCase__ = False,		lowerCamelCase__ = False,		lowerCamelCase__ = False,		lowerCamelCase__ = False,		lowerCamelCase__ = True,		lowerCamelCase__ = "layer_norm",		lowerCamelCase__ = False,		):
        super().__init__()
        A :       Dict          =   only_cross_attention

        A :       Tuple          =   (num_embeds_ada_norm is not None) and norm_type == """ada_norm_zero"""
        A :       Union[str, Any]          =   (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:
            A :       Dict          =   AdaLayerNorm(lowerCamelCase__,		lowerCamelCase__		)
        elif self.use_ada_layer_norm_zero:
            A :       List[str]          =   AdaLayerNormZero(lowerCamelCase__,		lowerCamelCase__		)
        else:
            A :       Tuple          =   nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__		)
        A :       Any          =   Attention(
            query_dim=lowerCamelCase__,		heads=lowerCamelCase__,		dim_head=lowerCamelCase__,		dropout=lowerCamelCase__,		bias=lowerCamelCase__,		cross_attention_dim=cross_attention_dim if only_cross_attention else None,		upcast_attention=lowerCamelCase__,		)

        # 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.
            A :       int          =   (
                AdaLayerNorm(lowerCamelCase__,		lowerCamelCase__		)
                if self.use_ada_layer_norm
                else nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__		)
            )
            A :       Dict          =   Attention(
                query_dim=lowerCamelCase__,		cross_attention_dim=cross_attention_dim if not double_self_attention else None,		heads=lowerCamelCase__,		dim_head=lowerCamelCase__,		dropout=lowerCamelCase__,		bias=lowerCamelCase__,		upcast_attention=lowerCamelCase__,		)  # is self-attn if encoder_hidden_states is none
        else:
            A :       Dict          =   None
            A :       Dict          =   None

        # 3. Feed-forward
        A :       Optional[Any]          =   nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__		)
        A :       int          =   FeedForward(lowerCamelCase__,		dropout=lowerCamelCase__,		activation_fn=lowerCamelCase__,		final_dropout=lowerCamelCase__		)

        # let chunk size default to None
        A :       Optional[Any]          =   None
        A :       int          =   0






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__		):
        # Sets chunk feed-forward
        A :       List[str]          =   chunk_size
        A :       int          =   dim






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__ = None,		lowerCamelCase__ = None,		lowerCamelCase__ = None,		lowerCamelCase__ = None,		lowerCamelCase__ = None,		lowerCamelCase__ = None,		):
        # Notice that normalization is always applied before the real computation in the following blocks.
        # 1. Self-Attention
        if self.use_ada_layer_norm:
            A :       Optional[int]          =   self.norma(lowerCamelCase__,		lowerCamelCase__		)
        elif self.use_ada_layer_norm_zero:
            A     ,							A     ,							A     ,							A     ,							A :       Tuple          =   self.norma(
                lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		hidden_dtype=hidden_states.dtype		)
        else:
            A :       Tuple          =   self.norma(lowerCamelCase__		)

        A :       Dict          =   cross_attention_kwargs if cross_attention_kwargs is not None else {}

        A :       str          =   self.attna(
            lowerCamelCase__,		encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,		attention_mask=lowerCamelCase__,		**lowerCamelCase__,		)
        if self.use_ada_layer_norm_zero:
            A :       List[Any]          =   gate_msa.unsqueeze(1		) * attn_output
        A :       str          =   attn_output + hidden_states

        # 2. Cross-Attention
        if self.attna is not None:
            A :       Optional[int]          =   (
                self.norma(lowerCamelCase__,		lowerCamelCase__		) if self.use_ada_layer_norm else self.norma(lowerCamelCase__		)
            )

            A :       Optional[Any]          =   self.attna(
                lowerCamelCase__,		encoder_hidden_states=lowerCamelCase__,		attention_mask=lowerCamelCase__,		**lowerCamelCase__,		)
            A :       Dict          =   attn_output + hidden_states

        # 3. Feed-forward
        A :       str          =   self.norma(lowerCamelCase__		)

        if self.use_ada_layer_norm_zero:
            A :       Tuple          =   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`.'''		)

            A :       Any          =   norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
            A :       Optional[Any]          =   torch.cat(
                [self.ff(lowerCamelCase__		) for hid_slice in norm_hidden_states.chunk(lowerCamelCase__,		dim=self._chunk_dim		)],		dim=self._chunk_dim,		)
        else:
            A :       Dict          =   self.ff(lowerCamelCase__		)

        if self.use_ada_layer_norm_zero:
            A :       Optional[Any]          =   gate_mlp.unsqueeze(1		) * ff_output

        A :       Optional[int]          =   ff_output + hidden_states

        return hidden_states






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__ = None,		lowerCamelCase__ = 4,		lowerCamelCase__ = 0.0,		lowerCamelCase__ = "geglu",		lowerCamelCase__ = False,		):
        super().__init__()
        A :       str          =   int(dim * mult		)
        A :       Optional[Any]          =   dim_out if dim_out is not None else dim

        if activation_fn == "gelu":
            A :       Dict          =   GELU(lowerCamelCase__,		lowerCamelCase__		)
        if activation_fn == "gelu-approximate":
            A :       Optional[int]          =   GELU(lowerCamelCase__,		lowerCamelCase__,		approximate="""tanh"""		)
        elif activation_fn == "geglu":
            A :       Any          =   GEGLU(lowerCamelCase__,		lowerCamelCase__		)
        elif activation_fn == "geglu-approximate":
            A :       Optional[Any]          =   ApproximateGELU(lowerCamelCase__,		lowerCamelCase__		)

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






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        for module in self.net:
            A :       int          =   module(lowerCamelCase__		)
        return hidden_states






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__ = "none"		):
        super().__init__()
        A :       Optional[int]          =   nn.Linear(lowerCamelCase__,		lowerCamelCase__		)
        A :       int          =   approximate






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        if gate.device.type != "mps":
            return F.gelu(lowerCamelCase__,		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,		lowerCamelCase__		):
        A :       List[Any]          =   self.proj(lowerCamelCase__		)
        A :       Union[str, Any]          =   self.gelu(lowerCamelCase__		)
        return hidden_states






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__		):
        super().__init__()
        A :       Union[str, Any]          =   nn.Linear(lowerCamelCase__,		dim_out * 2		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        if gate.device.type != "mps":
            return F.gelu(lowerCamelCase__		)
        # mps: gelu is not implemented for float16
        return F.gelu(gate.to(dtype=torch.floataa		)		).to(dtype=gate.dtype		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        A     ,							A :       Any          =   self.proj(lowerCamelCase__		).chunk(2,		dim=-1		)
        return hidden_states * self.gelu(lowerCamelCase__		)






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__		):
        super().__init__()
        A :       Union[str, Any]          =   nn.Linear(lowerCamelCase__,		lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        A :       List[Any]          =   self.proj(lowerCamelCase__		)
        return x * torch.sigmoid(1.702 * x		)






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__		):
        super().__init__()
        A :       Dict          =   nn.Embedding(lowerCamelCase__,		lowerCamelCase__		)
        A :       Tuple          =   nn.SiLU()
        A :       Tuple          =   nn.Linear(lowerCamelCase__,		embedding_dim * 2		)
        A :       List[str]          =   nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__		):
        A :       int          =   self.linear(self.silu(self.emb(lowerCamelCase__		)		)		)
        A     ,							A :       Optional[int]          =   torch.chunk(lowerCamelCase__,		2		)
        A :       Tuple          =   self.norm(lowerCamelCase__		) * (1 + scale) + shift
        return x






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__		):
        super().__init__()

        A :       Union[str, Any]          =   CombinedTimestepLabelEmbeddings(lowerCamelCase__,		lowerCamelCase__		)

        A :       List[Any]          =   nn.SiLU()
        A :       str          =   nn.Linear(lowerCamelCase__,		6 * embedding_dim,		bias=lowerCamelCase__		)
        A :       List[Any]          =   nn.LayerNorm(lowerCamelCase__,		elementwise_affine=lowerCamelCase__,		eps=1e-6		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__=None		):
        A :       Tuple          =   self.linear(self.silu(self.emb(lowerCamelCase__,		lowerCamelCase__,		hidden_dtype=lowerCamelCase__		)		)		)
        A     ,							A     ,							A     ,							A     ,							A     ,							A :       Optional[Any]          =   emb.chunk(6,		dim=1		)
        A :       List[str]          =   self.norm(lowerCamelCase__		) * (1 + scale_msa[:, None]) + shift_msa[:, None]
        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__ = None,		lowerCamelCase__ = 1e-5		):
        super().__init__()
        A :       int          =   num_groups
        A :       Dict          =   eps

        if act_fn is None:
            A :       Union[str, Any]          =   None
        else:
            A :       Optional[int]          =   get_activation(lowerCamelCase__		)

        A :       Dict          =   nn.Linear(lowerCamelCase__,		out_dim * 2		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__		):
        if self.act:
            A :       Dict          =   self.act(lowerCamelCase__		)
        A :       Optional[int]          =   self.linear(lowerCamelCase__		)
        A :       int          =   emb[:, :, None, None]
        A     ,							A :       Tuple          =   emb.chunk(2,		dim=1		)

        A :       Any          =   F.group_norm(lowerCamelCase__,		self.num_groups,		eps=self.eps		)
        A :       List[str]          =   x * (1 + scale) + shift
        return x


