


import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin



class 							_UpperCamelCase (						A,A,A						):

   '''simple docstring'''







   @register_to_config
   def __init__(  self  : Union[str, Any]       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : float       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : bool = False       ,  ):






         '''simple docstring'''



         super().__init__()

         __lowerCamelCase     :      int       						=						nn.Embedding(_lowerCamelCase       ,  _lowerCamelCase							)

         __lowerCamelCase     :      str       						=						nn.Embedding(_lowerCamelCase       ,  _lowerCamelCase							)
         __lowerCamelCase     :      str       						=						False

         __lowerCamelCase     :      int       						=						nn.Dropout(p=_lowerCamelCase							)

         __lowerCamelCase     :      Any       						=						TaConfig(
             vocab_size=_lowerCamelCase       ,  d_model=_lowerCamelCase       ,  num_heads=_lowerCamelCase       ,  d_kv=_lowerCamelCase       ,  d_ff=_lowerCamelCase       ,  dropout_rate=_lowerCamelCase       ,  feed_forward_proj=_lowerCamelCase       ,  is_decoder=_lowerCamelCase       ,  is_encoder_decoder=_lowerCamelCase       ,  )

         __lowerCamelCase     :      Optional[Any]       						=						nn.ModuleList()
         for lyr_num in range(_lowerCamelCase							):
               __lowerCamelCase     :      Union[str, Any]       						=						TaBlock(_lowerCamelCase							)
               self.encoders.append(_lowerCamelCase							)

         __lowerCamelCase     :      str       						=						TaLayerNorm(_lowerCamelCase							)
         __lowerCamelCase     :      List[str]       						=						nn.Dropout(p=_lowerCamelCase							)


   def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						self.token_embedder(_lowerCamelCase							)

         __lowerCamelCase     :      Tuple       						=						encoder_input_tokens.shape[1]
         __lowerCamelCase     :      Dict       						=						torch.arange(_lowerCamelCase       ,  device=encoder_input_tokens.device							)
         x += self.position_encoding(_lowerCamelCase							)

         __lowerCamelCase     :      List[str]       						=						self.dropout_pre(_lowerCamelCase							)

         # inverted the attention mask
         __lowerCamelCase     :      Dict       						=						encoder_input_tokens.size()
         __lowerCamelCase     :      Dict       						=						self.get_extended_attention_mask(_lowerCamelCase       ,  _lowerCamelCase							)

         for lyr in self.encoders:
               __lowerCamelCase     :      List[Any]       						=						lyr(_lowerCamelCase       ,  _lowerCamelCase							)[0]
         __lowerCamelCase     :      str       						=						self.layer_norm(_lowerCamelCase							)

         return self.dropout_post(_lowerCamelCase							), encoder_inputs_mask







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__UpperCamelCase    :						Optional[Any]     =							{
    'configuration_swinv2': ['SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP', 'Swinv2Config'],
}


try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __UpperCamelCase    :						List[Any]     =							[
        'SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST',
        'Swinv2ForImageClassification',
        'Swinv2ForMaskedImageModeling',
        'Swinv2Model',
        'Swinv2PreTrainedModel',
    ]


if TYPE_CHECKING:
    from .configuration_swinva import SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinvaConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_swinva import (
            SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST,
            SwinvaForImageClassification,
            SwinvaForMaskedImageModeling,
            SwinvaModel,
            SwinvaPreTrainedModel,
        )


else:
    import sys

    __UpperCamelCase    :						List[Any]     =							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



