

'''simple docstring'''




import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


_a				: List[str] 							=						logging.get_logger(__name__)



class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def __init__(     self     )							->				Optional[int]:


    """simple docstring"""
    super().__init__()
    if version.parse(torch.__version__     ) < version.parse('1.12.0'     ):
      raise ImportError(
          F'''You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use '''
          'PytorchGELUTanh. Please upgrade torch.'     )





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return nn.functional.gelu(a_   , approximate='tanh'     )




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi     ) * (input + 0.04_4715 * torch.pow(a_   , 3.0     ))     ))




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def __init__(     self   , a_ = False     )							->				List[Any]:


    """simple docstring"""
    super().__init__()
    if use_gelu_python:
      UpperCAmelCase													=	self._gelu_python
    else:
      UpperCAmelCase													=	nn.functional.gelu
  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0     )     ))





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return self.act(a_     )




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return 0.5 * input * (1.0 + torch.tanh(input * 0.79_7884_5608 * (1.0 + 0.04_4715 * input * input)     ))




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return input * torch.sigmoid(1.702 * input     )




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def __init__(     self   , a_   , a_     )							->				Union[str, Any]:


    """simple docstring"""
    if min > max:
      raise ValueError(F'''min should be < max (got min: {min}, max: {max})'''     )

    super().__init__()
    UpperCAmelCase													=	min
    UpperCAmelCase													=	max





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return torch.clip(gelu(a_     )   , self.min   , self.max     )




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def __init__(     self     )							->				Optional[Any]:


    """simple docstring"""
    super().__init__()
    UpperCAmelCase													=	math.sqrt(2 / math.pi     )





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.04_4715 * torch.pow(a_   , 3     ))     ))




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return nn.functional.silu(a_     )




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def __init__(     self     )							->				Tuple:


    """simple docstring"""
    super().__init__()
    if version.parse(torch.__version__     ) < version.parse('1.9.0'     ):
      UpperCAmelCase													=	self._mish_python
    else:
      UpperCAmelCase													=	nn.functional.mish
  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return input * torch.tanh(nn.functional.softplus(a_     )     )





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return self.act(a_     )




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def 						snake_case_			(     self   , a_     )							->				Tensor:


    """simple docstring"""
    return input




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def 						snake_case_			(     self   , a_   , a_=0.70_7107   , a_=0.28_2095     )							->				Optional[int]:


    """simple docstring"""
    UpperCAmelCase													=	(input - mu).div(sigma * math.sqrt(2.0     )     )
    return 0.5 * (1.0 + torch.erf(a_     ))




class        lowercase_	(		nn.Module ):

  '''simple docstring'''





  def 						snake_case_			(     self   , a_     )							->				Tuple:


    """simple docstring"""
    UpperCAmelCase													=	nn.functional.relu(a_     )
    UpperCAmelCase													=	torch.square(a_     )
    return squared




class        lowercase_	(		snake_case_ ):

  '''simple docstring'''





  def __getitem__(     self   , a_     )							->				Any:


    """simple docstring"""
    UpperCAmelCase													=	super().__getitem__(a_     )
    UpperCAmelCase							,      UpperCAmelCase													=	content if isinstance(a_   , a_     ) else (content, {})
    return cls(**a_     )




_a				: Tuple 							=						{
    'gelu': GELUActivation,
    'gelu_10': (ClippedGELUActivation, {'min': -10, 'max': 10}),
    'gelu_fast': FastGELUActivation,
    'gelu_new': NewGELUActivation,
    'gelu_python': (GELUActivation, {'use_gelu_python': True}),
    'gelu_pytorch_tanh': PytorchGELUTanh,
    'gelu_accurate': AccurateGELUActivation,
    'laplace': LaplaceActivation,
    'linear': LinearActivation,
    'mish': MishActivation,
    'quick_gelu': QuickGELUActivation,
    'relu': nn.ReLU,
    'relu2': ReLUSquaredActivation,
    'relu6': nn.ReLUa,
    'sigmoid': nn.Sigmoid,
    'silu': SiLUActivation,
    'swish': SiLUActivation,
    'tanh': nn.Tanh,
}
_a				: Optional[Any] 							=						ClassInstantier(ACTaCLS)


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   int						):
  if activation_string in ACTaFN:
    return ACTaFN[activation_string]
  else:
    raise KeyError(f'''function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()						)}'''						)


# For backwards compatibility with: from activations import gelu_python
_a				: int 							=						get_activation('gelu_python')
_a				: Dict 							=						get_activation('gelu_new')
_a				: Tuple 							=						get_activation('gelu')
_a				: List[Any] 							=						get_activation('gelu_fast')
_a				: Dict 							=						get_activation('quick_gelu')
_a				: Optional[Any] 							=						get_activation('silu')
_a				: Optional[Any] 							=						get_activation('mish')



_a				: Optional[int] 							=						get_activation('linear')





'''simple docstring'''




from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
  import torch

_a				: List[Any] 							=						logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
  from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):


  def       lowerCamelCase__		(						cls      :   Dict						):
    return cls


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Union[Tuple, List]       ,   SCREAMING_SNAKE_CASE      :   Optional[Union[List["torch.Generator"], "torch.Generator"]] = None       ,   SCREAMING_SNAKE_CASE      :   Optional["torch.device"] = None       ,   SCREAMING_SNAKE_CASE      :   Optional["torch.dtype"] = None       ,   SCREAMING_SNAKE_CASE      :   Optional["torch.layout"] = None       ,   ):
  UpperCAmelCase													=	device
  UpperCAmelCase													=	shape[0]

  UpperCAmelCase													=	layout or torch.strided
  UpperCAmelCase													=	device or torch.device('cpu'						)

  if generator is not None:
    UpperCAmelCase													=	generator.device.type if not isinstance(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						) else generator[0].device.type
    if gen_device_type != device.type and gen_device_type == "cpu":
      UpperCAmelCase													=	'cpu'
      if device != "mps":
        logger.info(
            f'''The passed generator was created on \'cpu\' even though a tensor on {device} was expected.'''
            f''' Tensors will be created on \'cpu\' and then moved to {device}. Note that one can probably'''
            f''' slighly speed up this function by passing a generator that was created on the {device} device.'''						)
    elif gen_device_type != device.type and gen_device_type == "cuda":
      raise ValueError(f'''Cannot generate a {device} tensor from a generator of type {gen_device_type}.'''						)

  if isinstance(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						):
    UpperCAmelCase													=	(1,) + shape[1:]
    UpperCAmelCase													=	[
        torch.randn(SCREAMING_SNAKE_CASE       ,   generator=generator[i]       ,   device=SCREAMING_SNAKE_CASE       ,   dtype=SCREAMING_SNAKE_CASE       ,   layout=SCREAMING_SNAKE_CASE						)
        for i in range(SCREAMING_SNAKE_CASE						)
    ]
    UpperCAmelCase													=	torch.cat(SCREAMING_SNAKE_CASE       ,   dim=0						).to(SCREAMING_SNAKE_CASE						)
  else:
    UpperCAmelCase													=	torch.randn(SCREAMING_SNAKE_CASE       ,   generator=SCREAMING_SNAKE_CASE       ,   device=SCREAMING_SNAKE_CASE       ,   dtype=SCREAMING_SNAKE_CASE       ,   layout=SCREAMING_SNAKE_CASE						).to(SCREAMING_SNAKE_CASE						)

  return latents



def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Tuple						):
  if is_torch_version('<'       ,   '2.0.0'						) or not hasattr(SCREAMING_SNAKE_CASE       ,   '_dynamo'						):
    return False
  return isinstance(SCREAMING_SNAKE_CASE       ,   torch._dynamo.eval_frame.OptimizedModule						)


