






import operator as op


SCREAMING_SNAKE_CASE__		:							str 				=				'scaler.pt'
SCREAMING_SNAKE_CASE__		:							Any 				=				'pytorch_model'
SCREAMING_SNAKE_CASE__		:							Any 				=				'random_states'
SCREAMING_SNAKE_CASE__		:							List[str] 				=				'optimizer'
SCREAMING_SNAKE_CASE__		:							List[str] 				=				'scheduler'
SCREAMING_SNAKE_CASE__		:							Tuple 				=				'pytorch_model.bin'
SCREAMING_SNAKE_CASE__		:							Optional[Any] 				=				'pytorch_model.bin.index.json'
SCREAMING_SNAKE_CASE__		:							List[str] 				=				'model.safetensors'
SCREAMING_SNAKE_CASE__		:							Dict 				=				'model.safetensors.index.json'
SCREAMING_SNAKE_CASE__		:							Optional[Any] 				=				'1.10.2'
SCREAMING_SNAKE_CASE__		:							Optional[Any] 				=				'py38'
SCREAMING_SNAKE_CASE__		:							Tuple 				=				'4.17.0'
SCREAMING_SNAKE_CASE__		:							Any 				=				['ml.p3.16xlarge', 'ml.p3dn.24xlarge', 'ml.p4dn.24xlarge']
SCREAMING_SNAKE_CASE__		:							int 				=				['FULL_SHARD', 'SHARD_GRAD_OP', 'NO_SHARD', 'HYBRID_SHARD', 'HYBRID_SHARD_ZERO2']
SCREAMING_SNAKE_CASE__		:							List[str] 				=				['TRANSFORMER_BASED_WRAP', 'SIZE_BASED_WRAP', 'NO_WRAP']
SCREAMING_SNAKE_CASE__		:							int 				=				['BACKWARD_PRE', 'BACKWARD_POST', 'NO_PREFETCH']
SCREAMING_SNAKE_CASE__		:							Optional[int] 				=				['FULL_STATE_DICT', 'LOCAL_STATE_DICT', 'SHARDED_STATE_DICT']
SCREAMING_SNAKE_CASE__		:							Optional[int] 				=				'2.0.1'
SCREAMING_SNAKE_CASE__		:							Tuple 				=				['pdsh', 'standard', 'openmpi', 'mvapich']
SCREAMING_SNAKE_CASE__		:							Union[str, Any] 				=				['default', 'reduce-overhead', 'max-autotune']

SCREAMING_SNAKE_CASE__		:							List[Any] 				=				{'>': op.gt, '>=': op.ge, '==': op.eq, '!=': op.ne, '<=': op.le, '<': op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
SCREAMING_SNAKE_CASE__		:							Tuple 				=				[
    'nnodes',
    'nproc_per_node',
    'rdzv_backend',
    'rdzv_endpoint',
    'rdzv_id',
    'rdzv_conf',
    'standalone',
    'max_restarts',
    'monitor_interval',
    'start_method',
    'role',
    'module',
    'm',
    'no_python',
    'run_path',
    'log_dir',
    'r',
    'redirects',
    't',
    'tee',
    'node_rank',
    'master_addr',
    'master_port',
]

SCREAMING_SNAKE_CASE__		:							List[str] 				=				['DEEPSPEED', 'MULTI_GPU', 'FSDP', 'MEGATRON_LM']
SCREAMING_SNAKE_CASE__		:							List[Any] 				=				['DEEPSPEED', 'MULTI_XPU', 'FSDP']








"""simple docstring"""




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     lowercase(_lowercase    ,						_lowercase    ,						_lowercase       ):



     @register_to_config
     def __init__(							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE = False     ,				)		->     Union[str, Any]:







         """simple docstring"""



         super().__init__()

         a__		=			nn.Embedding(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )

         a__		=			nn.Embedding(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )
         a__		=			False

         a__		=			nn.Dropout(p=__SCREAMING_SNAKE_CASE   )

         a__		=			TaConfig(
             vocab_size=__SCREAMING_SNAKE_CASE     ,				d_model=__SCREAMING_SNAKE_CASE     ,				num_heads=__SCREAMING_SNAKE_CASE     ,				d_kv=__SCREAMING_SNAKE_CASE     ,				d_ff=__SCREAMING_SNAKE_CASE     ,				dropout_rate=__SCREAMING_SNAKE_CASE     ,				feed_forward_proj=__SCREAMING_SNAKE_CASE     ,				is_decoder=__SCREAMING_SNAKE_CASE     ,				is_encoder_decoder=__SCREAMING_SNAKE_CASE     ,				)

         a__		=			nn.ModuleList()
         for lyr_num in range(__SCREAMING_SNAKE_CASE   ):
             a__		=			TaBlock(__SCREAMING_SNAKE_CASE   )
             self.encoders.append(__SCREAMING_SNAKE_CASE   )

         a__		=			TaLayerNorm(__SCREAMING_SNAKE_CASE   )
         a__		=			nn.Dropout(p=__SCREAMING_SNAKE_CASE   )




     def 			lowercase__  (							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )		->     List[Any]:







         """simple docstring"""



         a__		=			self.token_embedder(__SCREAMING_SNAKE_CASE   )

         a__		=			encoder_input_tokens.shape[1]
         a__		=			torch.arange(__SCREAMING_SNAKE_CASE     ,				device=encoder_input_tokens.device   )
         x += self.position_encoding(__SCREAMING_SNAKE_CASE   )

         a__		=			self.dropout_pre(__SCREAMING_SNAKE_CASE   )

         # inverted the attention mask
         a__		=			encoder_input_tokens.size()
         a__		=			self.get_extended_attention_mask(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )

         for lyr in self.encoders:
             a__		=			lyr(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )[0]
         a__		=			self.layer_norm(__SCREAMING_SNAKE_CASE   )

         return self.dropout_post(__SCREAMING_SNAKE_CASE   ), encoder_inputs_mask


