

import unittest

from transformers import (
    MODEL_FOR_OBJECT_DETECTION_MAPPING,
    AutoFeatureExtractor,
    AutoModelForObjectDetection,
    ObjectDetectionPipeline,
    is_vision_available,
    pipeline,
)
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_pytesseract,
    require_tf,
    require_timm,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
 from PIL import Image
else:







 class 	__UpperCamelCase    :

   @staticmethod
   def __lowerCamelCase		(		*_A    ,**_A      ):
       '''simple docstring'''

       pass







@is_pipeline_test
@require_vision
@require_timm
@require_torch
class 	__UpperCamelCase    (    unittest.TestCase						):
  _UpperCAmelCase          =		MODEL_FOR_OBJECT_DETECTION_MAPPING

  def __lowerCamelCase		(		self    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    ObjectDetectionPipeline(model=_A    ,image_processor=_A      )
      return object_detector, ["./tests/fixtures/tests_samples/COCO/000000039769.png"]

  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    object_detector('./tests/fixtures/tests_samples/COCO/000000039769.png'    ,threshold=0.0      )

      self.assertGreater(len(_A      )    ,0      )
      for detected_object in outputs:
          self.assertEqual(
              _A    ,{
                  'score': ANY(_A      ),
                  'label': ANY(_A      ),
                  'box': {'xmin': ANY(_A      ), 'ymin': ANY(_A      ), 'xmax': ANY(_A      ), 'ymax': ANY(_A      )},
              }    ,)

      import datasets

      _lowerCAmelCase       :   List[Any]  =    datasets.load_dataset('hf-internal-testing/fixtures_image_utils'    ,'image'    ,split='test'      )

      _lowerCAmelCase       :   Optional[int]  =    [
          Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      ),
          'http://images.cocodataset.org/val2017/000000039769.jpg',
          # RGBA
          dataset[0]['file'],
          # LA
          dataset[1]['file'],
          # L
          dataset[2]['file'],
      ]
      _lowerCAmelCase       :   Dict  =    object_detector(_A    ,threshold=0.0      )

      self.assertEqual(len(_A      )    ,len(_A      )      )
      for outputs in batch_outputs:
          self.assertGreater(len(_A      )    ,0      )
          for detected_object in outputs:
              self.assertEqual(
                  _A    ,{
                      'score': ANY(_A      ),
                      'label': ANY(_A      ),
                      'box': {'xmin': ANY(_A      ), 'ymin': ANY(_A      ), 'xmax': ANY(_A      ), 'ymax': ANY(_A      )},
                  }    ,)

  @require_tf
  @unittest.skip('Object detection not implemented in TF'      )
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass

  @require_torch
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    'hf-internal-testing/tiny-detr-mobilenetsv3'

      _lowerCAmelCase       :   Optional[int]  =    AutoModelForObjectDetection.from_pretrained(_A      )
      _lowerCAmelCase       :   int  =    AutoFeatureExtractor.from_pretrained(_A      )
      _lowerCAmelCase       :   Tuple  =    ObjectDetectionPipeline(model=_A    ,feature_extractor=_A      )

      _lowerCAmelCase       :   str  =    object_detector('http://images.cocodataset.org/val2017/000000039769.jpg'    ,threshold=0.0      )

      self.assertEqual(
          nested_simplify(_A    ,decimals=4      )    ,[
              {'score': 0.3_3_7_6, 'label': 'LABEL_0', 'box': {'xmin': 159, 'ymin': 120, 'xmax': 480, 'ymax': 359}},
              {'score': 0.3_3_7_6, 'label': 'LABEL_0', 'box': {'xmin': 159, 'ymin': 120, 'xmax': 480, 'ymax': 359}},
          ]    ,)

      _lowerCAmelCase       :   Optional[Any]  =    object_detector(
          [
              'http://images.cocodataset.org/val2017/000000039769.jpg',
              'http://images.cocodataset.org/val2017/000000039769.jpg',
          ]    ,threshold=0.0    ,)

      self.assertEqual(
          nested_simplify(_A    ,decimals=4      )    ,[
              [
                  {'score': 0.3_3_7_6, 'label': 'LABEL_0', 'box': {'xmin': 159, 'ymin': 120, 'xmax': 480, 'ymax': 359}},
                  {'score': 0.3_3_7_6, 'label': 'LABEL_0', 'box': {'xmin': 159, 'ymin': 120, 'xmax': 480, 'ymax': 359}},
              ],
              [
                  {'score': 0.3_3_7_6, 'label': 'LABEL_0', 'box': {'xmin': 159, 'ymin': 120, 'xmax': 480, 'ymax': 359}},
                  {'score': 0.3_3_7_6, 'label': 'LABEL_0', 'box': {'xmin': 159, 'ymin': 120, 'xmax': 480, 'ymax': 359}},
              ],
          ]    ,)

  @require_torch
  @slow
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    'facebook/detr-resnet-50'

      _lowerCAmelCase       :   Optional[int]  =    AutoModelForObjectDetection.from_pretrained(_A      )
      _lowerCAmelCase       :   List[Any]  =    AutoFeatureExtractor.from_pretrained(_A      )
      _lowerCAmelCase       :   List[str]  =    ObjectDetectionPipeline(model=_A    ,feature_extractor=_A      )

      _lowerCAmelCase       :   List[Any]  =    object_detector('http://images.cocodataset.org/val2017/000000039769.jpg'      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=4      )    ,[
              {'score': 0.9_9_8_2, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 70, 'xmax': 175, 'ymax': 117}},
              {'score': 0.9_9_6_0, 'label': 'remote', 'box': {'xmin': 333, 'ymin': 72, 'xmax': 368, 'ymax': 187}},
              {'score': 0.9_9_5_5, 'label': 'couch', 'box': {'xmin': 0, 'ymin': 1, 'xmax': 639, 'ymax': 473}},
              {'score': 0.9_9_8_8, 'label': 'cat', 'box': {'xmin': 13, 'ymin': 52, 'xmax': 314, 'ymax': 470}},
              {'score': 0.9_9_8_7, 'label': 'cat', 'box': {'xmin': 345, 'ymin': 23, 'xmax': 640, 'ymax': 368}},
          ]    ,)

      _lowerCAmelCase       :   Optional[Any]  =    object_detector(
          [
              'http://images.cocodataset.org/val2017/000000039769.jpg',
              'http://images.cocodataset.org/val2017/000000039769.jpg',
          ]      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=4      )    ,[
              [
                  {'score': 0.9_9_8_2, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 70, 'xmax': 175, 'ymax': 117}},
                  {'score': 0.9_9_6_0, 'label': 'remote', 'box': {'xmin': 333, 'ymin': 72, 'xmax': 368, 'ymax': 187}},
                  {'score': 0.9_9_5_5, 'label': 'couch', 'box': {'xmin': 0, 'ymin': 1, 'xmax': 639, 'ymax': 473}},
                  {'score': 0.9_9_8_8, 'label': 'cat', 'box': {'xmin': 13, 'ymin': 52, 'xmax': 314, 'ymax': 470}},
                  {'score': 0.9_9_8_7, 'label': 'cat', 'box': {'xmin': 345, 'ymin': 23, 'xmax': 640, 'ymax': 368}},
              ],
              [
                  {'score': 0.9_9_8_2, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 70, 'xmax': 175, 'ymax': 117}},
                  {'score': 0.9_9_6_0, 'label': 'remote', 'box': {'xmin': 333, 'ymin': 72, 'xmax': 368, 'ymax': 187}},
                  {'score': 0.9_9_5_5, 'label': 'couch', 'box': {'xmin': 0, 'ymin': 1, 'xmax': 639, 'ymax': 473}},
                  {'score': 0.9_9_8_8, 'label': 'cat', 'box': {'xmin': 13, 'ymin': 52, 'xmax': 314, 'ymax': 470}},
                  {'score': 0.9_9_8_7, 'label': 'cat', 'box': {'xmin': 345, 'ymin': 23, 'xmax': 640, 'ymax': 368}},
              ],
          ]    ,)

  @require_torch
  @slow
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    'facebook/detr-resnet-50'

      _lowerCAmelCase       :   Dict  =    pipeline('object-detection'    ,model=_A      )

      _lowerCAmelCase       :   Optional[Any]  =    object_detector('http://images.cocodataset.org/val2017/000000039769.jpg'      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=4      )    ,[
              {'score': 0.9_9_8_2, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 70, 'xmax': 175, 'ymax': 117}},
              {'score': 0.9_9_6_0, 'label': 'remote', 'box': {'xmin': 333, 'ymin': 72, 'xmax': 368, 'ymax': 187}},
              {'score': 0.9_9_5_5, 'label': 'couch', 'box': {'xmin': 0, 'ymin': 1, 'xmax': 639, 'ymax': 473}},
              {'score': 0.9_9_8_8, 'label': 'cat', 'box': {'xmin': 13, 'ymin': 52, 'xmax': 314, 'ymax': 470}},
              {'score': 0.9_9_8_7, 'label': 'cat', 'box': {'xmin': 345, 'ymin': 23, 'xmax': 640, 'ymax': 368}},
          ]    ,)

      _lowerCAmelCase       :   Dict  =    object_detector(
          [
              'http://images.cocodataset.org/val2017/000000039769.jpg',
              'http://images.cocodataset.org/val2017/000000039769.jpg',
          ]      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=4      )    ,[
              [
                  {'score': 0.9_9_8_2, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 70, 'xmax': 175, 'ymax': 117}},
                  {'score': 0.9_9_6_0, 'label': 'remote', 'box': {'xmin': 333, 'ymin': 72, 'xmax': 368, 'ymax': 187}},
                  {'score': 0.9_9_5_5, 'label': 'couch', 'box': {'xmin': 0, 'ymin': 1, 'xmax': 639, 'ymax': 473}},
                  {'score': 0.9_9_8_8, 'label': 'cat', 'box': {'xmin': 13, 'ymin': 52, 'xmax': 314, 'ymax': 470}},
                  {'score': 0.9_9_8_7, 'label': 'cat', 'box': {'xmin': 345, 'ymin': 23, 'xmax': 640, 'ymax': 368}},
              ],
              [
                  {'score': 0.9_9_8_2, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 70, 'xmax': 175, 'ymax': 117}},
                  {'score': 0.9_9_6_0, 'label': 'remote', 'box': {'xmin': 333, 'ymin': 72, 'xmax': 368, 'ymax': 187}},
                  {'score': 0.9_9_5_5, 'label': 'couch', 'box': {'xmin': 0, 'ymin': 1, 'xmax': 639, 'ymax': 473}},
                  {'score': 0.9_9_8_8, 'label': 'cat', 'box': {'xmin': 13, 'ymin': 52, 'xmax': 314, 'ymax': 470}},
                  {'score': 0.9_9_8_7, 'label': 'cat', 'box': {'xmin': 345, 'ymin': 23, 'xmax': 640, 'ymax': 368}},
              ],
          ]    ,)

  @require_torch
  @slow
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    0.9_9_8_5
      _lowerCAmelCase       :   str  =    'facebook/detr-resnet-50'

      _lowerCAmelCase       :   Optional[int]  =    pipeline('object-detection'    ,model=_A      )

      _lowerCAmelCase       :   Optional[Any]  =    object_detector('http://images.cocodataset.org/val2017/000000039769.jpg'    ,threshold=_A      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=4      )    ,[
              {'score': 0.9_9_8_8, 'label': 'cat', 'box': {'xmin': 13, 'ymin': 52, 'xmax': 314, 'ymax': 470}},
              {'score': 0.9_9_8_7, 'label': 'cat', 'box': {'xmin': 345, 'ymin': 23, 'xmax': 640, 'ymax': 368}},
          ]    ,)







  @require_torch
  @require_pytesseract
  @slow
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    'Narsil/layoutlmv3-finetuned-funsd'
      _lowerCAmelCase       :   Tuple  =    0.9_9_9_3

      _lowerCAmelCase       :   str  =    pipeline('object-detection'    ,model=_A    ,threshold=_A      )

      _lowerCAmelCase       :   Any  =    object_detector(
          'https://huggingface.co/spaces/impira/docquery/resolve/2359223c1837a7587402bda0f2643382a6eefeab/invoice.png'      )
      self.assertEqual(
          nested_simplify(_A    ,decimals=4      )    ,[
              {'score': 0.9_9_9_3, 'label': 'I-ANSWER', 'box': {'xmin': 294, 'ymin': 254, 'xmax': 343, 'ymax': 264}},
              {'score': 0.9_9_9_3, 'label': 'I-ANSWER', 'box': {'xmin': 294, 'ymin': 254, 'xmax': 343, 'ymax': 264}},
          ]    ,)








"""simple docstring"""






import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin







class 	__UpperCamelCase    (    a__							,							a__						):

  @register_to_config
  def __init__(		self    ,_A = 128    ,_A = 256    ,_A = 2_0_0_0.0    ,_A = 768    ,_A = 12    ,_A = 12    ,_A = 64    ,_A = 2048    ,_A = 0.1    ,):
      '''simple docstring'''

      super().__init__()

      _lowerCAmelCase       :   int  =    nn.Sequential(
          nn.Linear(_A    ,d_model * 4    ,bias=_A      )    ,nn.SiLU()    ,nn.Linear(d_model * 4    ,d_model * 4    ,bias=_A      )    ,nn.SiLU()    ,)

      _lowerCAmelCase       :   Any  =    nn.Embedding(_A    ,_A      )
      _lowerCAmelCase       :   Tuple  =    False

      _lowerCAmelCase       :   Union[str, Any]  =    nn.Linear(_A    ,_A    ,bias=_A      )

      _lowerCAmelCase       :   int  =    nn.Dropout(p=_A      )

      _lowerCAmelCase       :   int  =    nn.ModuleList()
      for lyr_num in range(_A      ):
          # FiLM conditional T5 decoder
          _lowerCAmelCase       :   Any  =    DecoderLayer(d_model=_A    ,d_kv=_A    ,num_heads=_A    ,d_ff=_A    ,dropout_rate=_A      )
          self.decoders.append(_A      )

      _lowerCAmelCase       :   Optional[Any]  =    TaLayerNorm(_A      )

      _lowerCAmelCase       :   List[str]  =    nn.Dropout(p=_A      )
      _lowerCAmelCase       :   Optional[Any]  =    nn.Linear(_A    ,_A    ,bias=_A      )

  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    torch.mul(query_input.unsqueeze(-1      )    ,key_input.unsqueeze(-2      )      )
      return mask.unsqueeze(-3      )







  def __lowerCamelCase		(		self    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase,					_lowerCAmelCase       :   Dict  =    decoder_input_tokens.shape
      assert decoder_noise_time.shape == (batch,)

      # decoder_noise_time is in [0, 1), so rescale to expected timing range.
      _lowerCAmelCase       :   Any  =    get_timestep_embedding(
          decoder_noise_time * self.config.max_decoder_noise_time    ,embedding_dim=self.config.d_model    ,max_period=self.config.max_decoder_noise_time    ,).to(dtype=self.dtype      )

      _lowerCAmelCase       :   Union[str, Any]  =    self.conditioning_emb(_A      ).unsqueeze(1      )

      assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

      _lowerCAmelCase       :   str  =    decoder_input_tokens.shape[1]

      # If we want to use relative positions for audio context, we can just offset
      # this sequence by the length of encodings_and_masks.
      _lowerCAmelCase       :   Union[str, Any]  =    torch.broadcast_to(
          torch.arange(_A    ,device=decoder_input_tokens.device      )    ,(batch, seq_length)    ,)

      _lowerCAmelCase       :   Any  =    self.position_encoding(_A      )

      _lowerCAmelCase       :   str  =    self.continuous_inputs_projection(_A      )
      inputs += position_encodings
      _lowerCAmelCase       :   int  =    self.dropout(_A      )

      # decoder: No padding present.
      _lowerCAmelCase       :   Union[str, Any]  =    torch.ones(
          decoder_input_tokens.shape[:2]    ,device=decoder_input_tokens.device    ,dtype=inputs.dtype      )

      # Translate encoding masks to encoder-decoder masks.
      _lowerCAmelCase       :   Optional[Any]  =    [(x, self.encoder_decoder_mask(_A    ,_A      )) for x, y in encodings_and_masks]

      # cross attend style: concat encodings
      _lowerCAmelCase       :   Dict  =    torch.cat([x[0] for x in encodings_and_encdec_masks]    ,dim=1      )
      _lowerCAmelCase       :   Tuple  =    torch.cat([x[1] for x in encodings_and_encdec_masks]    ,dim=-1      )

      for lyr in self.decoders:
          _lowerCAmelCase       :   Tuple  =    lyr(
              _A    ,conditioning_emb=_A    ,encoder_hidden_states=_A    ,encoder_attention_mask=_A    ,)[0]

      _lowerCAmelCase       :   Any  =    self.decoder_norm(_A      )
      _lowerCAmelCase       :   List[Any]  =    self.post_dropout(_A      )

      _lowerCAmelCase       :   int  =    self.spec_out(_A      )
      return spec_out







class 	__UpperCamelCase    (    nn.Module						):

  def __init__(		self    ,_A    ,_A    ,_A    ,_A    ,_A    ,_A=1E-6      ):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   Optional[Any]  =    nn.ModuleList()

      # cond self attention: layer 0
      self.layer.append(
          TaLayerSelfAttentionCond(d_model=_A    ,d_kv=_A    ,num_heads=_A    ,dropout_rate=_A      )      )

      # cross attention: layer 1
      self.layer.append(
          TaLayerCrossAttention(
              d_model=_A    ,d_kv=_A    ,num_heads=_A    ,dropout_rate=_A    ,layer_norm_epsilon=_A    ,)      )

      # Film Cond MLP + dropout: last layer
      self.layer.append(
          TaLayerFFCond(d_model=_A    ,d_ff=_A    ,dropout_rate=_A    ,layer_norm_epsilon=_A      )      )







  def __lowerCamelCase		(		self    ,_A    ,_A=None    ,_A=None    ,_A=None    ,_A=None    ,_A=None    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    self.layer[0](
          _A    ,conditioning_emb=_A    ,attention_mask=_A    ,)

      if encoder_hidden_states is not None:
          _lowerCAmelCase       :   Any  =    torch.where(encoder_attention_mask > 0    ,0    ,-1E10      ).to(
              encoder_hidden_states.dtype      )

          _lowerCAmelCase       :   str  =    self.layer[1](
              _A    ,key_value_states=_A    ,attention_mask=_A    ,)

      # Apply Film Conditional Feed Forward layer
      _lowerCAmelCase       :   Optional[Any]  =    self.layer[-1](_A    ,_A      )

      return (hidden_states,)







class 	__UpperCamelCase    (    nn.Module						):

  def __init__(		self    ,_A    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   Union[str, Any]  =    TaLayerNorm(_A      )
      _lowerCAmelCase       :   Any  =    TaFiLMLayer(in_features=d_model * 4    ,out_features=_A      )
      _lowerCAmelCase       :   Dict  =    Attention(query_dim=_A    ,heads=_A    ,dim_head=_A    ,out_bias=_A    ,scale_qk=_A      )
      _lowerCAmelCase       :   Tuple  =    nn.Dropout(_A      )







  def __lowerCamelCase		(		self    ,_A    ,_A=None    ,_A=None    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    self.layer_norm(_A      )

      if conditioning_emb is not None:
          _lowerCAmelCase       :   Union[str, Any]  =    self.FiLMLayer(_A    ,_A      )

      # Self-attention block
      _lowerCAmelCase       :   Union[str, Any]  =    self.attention(_A      )

      _lowerCAmelCase       :   Optional[Any]  =    hidden_states + self.dropout(_A      )

      return hidden_states







class 	__UpperCamelCase    (    nn.Module						):

  def __init__(		self    ,_A    ,_A    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   List[str]  =    Attention(query_dim=_A    ,heads=_A    ,dim_head=_A    ,out_bias=_A    ,scale_qk=_A      )
      _lowerCAmelCase       :   Optional[int]  =    TaLayerNorm(_A    ,eps=_A      )
      _lowerCAmelCase       :   Tuple  =    nn.Dropout(_A      )







  def __lowerCamelCase		(		self    ,_A    ,_A=None    ,_A=None    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   Union[str, Any]  =    self.layer_norm(_A      )
      _lowerCAmelCase       :   str  =    self.attention(
          _A    ,encoder_hidden_states=_A    ,attention_mask=attention_mask.squeeze(1      )    ,)
      _lowerCAmelCase       :   Any  =    hidden_states + self.dropout(_A      )
      return layer_output







class 	__UpperCamelCase    (    nn.Module						):

  def __init__(		self    ,_A    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   Optional[int]  =    TaDenseGatedActDense(d_model=_A    ,d_ff=_A    ,dropout_rate=_A      )
      _lowerCAmelCase       :   Tuple  =    TaFiLMLayer(in_features=d_model * 4    ,out_features=_A      )
      _lowerCAmelCase       :   Any  =    TaLayerNorm(_A    ,eps=_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    nn.Dropout(_A      )







  def __lowerCamelCase		(		self    ,_A    ,_A=None      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    self.layer_norm(_A      )
      if conditioning_emb is not None:
          _lowerCAmelCase       :   Union[str, Any]  =    self.film(_A    ,_A      )

      _lowerCAmelCase       :   str  =    self.DenseReluDense(_A      )
      _lowerCAmelCase       :   Tuple  =    hidden_states + self.dropout(_A      )
      return hidden_states







class 	__UpperCamelCase    (    nn.Module						):

  def __init__(		self    ,_A    ,_A    ,_A      ):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   Union[str, Any]  =    nn.Linear(_A    ,_A    ,bias=_A      )
      _lowerCAmelCase       :   Any  =    nn.Linear(_A    ,_A    ,bias=_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    nn.Linear(_A    ,_A    ,bias=_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    nn.Dropout(_A      )
      _lowerCAmelCase       :   int  =    NewGELUActivation()







  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    self.act(self.wi_a(_A      )      )
      _lowerCAmelCase       :   Optional[int]  =    self.wi_a(_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    hidden_gelu * hidden_linear
      _lowerCAmelCase       :   Dict  =    self.dropout(_A      )

      _lowerCAmelCase       :   Dict  =    self.wo(_A      )
      return hidden_states







class 	__UpperCamelCase    (    nn.Module						):

  def __init__(		self    ,_A    ,_A=1E-6      ):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   Union[str, Any]  =    nn.Parameter(torch.ones(_A      )      )
      _lowerCAmelCase       :   Optional[int]  =    eps







  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    hidden_states.to(torch.floataa      ).pow(2      ).mean(-1    ,keepdim=_A      )
      _lowerCAmelCase       :   List[Any]  =    hidden_states * torch.rsqrt(variance + self.variance_epsilon      )

      # convert into half-precision if necessary
      if self.weight.dtype in [torch.floataa, torch.bfloataa]:
          _lowerCAmelCase       :   Optional[int]  =    hidden_states.to(self.weight.dtype      )

      return self.weight * hidden_states







class 	__UpperCamelCase    (    nn.Module						):

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi      ) * (input + 0.0_4_4_7_1_5 * torch.pow(_A    ,3.0      ))      ))







class 	__UpperCamelCase    (    nn.Module						):

  def __init__(		self    ,_A    ,_A      ):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   List[str]  =    nn.Linear(_A    ,out_features * 2    ,bias=_A      )







  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    self.scale_bias(_A      )
      _lowerCAmelCase,					_lowerCAmelCase       :   List[Any]  =    torch.chunk(_A    ,2    ,-1      )
      _lowerCAmelCase       :   List[Any]  =    x * (1 + scale) + shift
      return x





