


'''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     __lowerCAmelCase		(lowercase_    ,	lowercase_    ,	lowercase_	):





      '''simple docstring'''






      @register_to_config
      def __init__(self       :       List[str]  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       float  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       bool = False  ,		):

             '''simple docstring'''







             super().__init__()

             lowercase__      =					nn.Embedding(UpperCamelCase  ,		UpperCamelCase					)

             lowercase__      =					nn.Embedding(UpperCamelCase  ,		UpperCamelCase					)
             lowercase__      =					False

             lowercase__      =					nn.Dropout(p=UpperCamelCase					)

             lowercase__      =					TaConfig(
                 vocab_size=UpperCamelCase  ,		d_model=UpperCamelCase  ,		num_heads=UpperCamelCase  ,		d_kv=UpperCamelCase  ,		d_ff=UpperCamelCase  ,		dropout_rate=UpperCamelCase  ,		feed_forward_proj=UpperCamelCase  ,		is_decoder=UpperCamelCase  ,		is_encoder_decoder=UpperCamelCase  ,		)

             lowercase__      =					nn.ModuleList()
             for lyr_num in range(UpperCamelCase					):
                    lowercase__      =					TaBlock(UpperCamelCase					)
                    self.encoders.append(UpperCamelCase					)

             lowercase__      =					TaLayerNorm(UpperCamelCase					)
             lowercase__      =					nn.Dropout(p=UpperCamelCase					)






      def  UpperCamelCase__						(self       :       Dict  ,		UpperCamelCase       :       List[Any]  ,		UpperCamelCase       :       Dict					):

             '''simple docstring'''







             lowercase__      =					self.token_embedder(UpperCamelCase					)

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

             lowercase__      =					self.dropout_pre(UpperCamelCase					)

             # inverted the attention mask
             lowercase__      =					encoder_input_tokens.size()
             lowercase__      =					self.get_extended_attention_mask(UpperCamelCase  ,		UpperCamelCase					)

             for lyr in self.encoders:
                    lowercase__      =					lyr(UpperCamelCase  ,		UpperCamelCase					)[0]
             lowercase__      =					self.layer_norm(UpperCamelCase					)

             return self.dropout_post(UpperCamelCase					), encoder_inputs_mask








'''simple docstring'''

def _SCREAMING_SNAKE_CASE (A = 1_000			)			->			int:






       """simple docstring"""







       return sum(2 * a * ((a - 1) // 2) for a in range(3    ,					n + 1			)			)


if __name__ == "__main__":
 print(solution())




