

"""simple docstring"""




import torch
from torch import nn







class      __lowercase			(    nn.Module					):

  '''simple docstring'''





  def __init__(  self		:			List[Any]					,A		:			List[Any]					,A		:			Any					,A		:			Optional[int]					,A		:			List[Any]					,A		:			Optional[int]=1					,A		:			List[Any]=False    ):






   '''simple docstring'''
   super().__init__()

   UpperCAmelCase__					:  Optional[Any]					   =					n_token
   UpperCAmelCase__					:  Tuple					   =					d_embed
   UpperCAmelCase__					:  Tuple					   =					d_proj

   UpperCAmelCase__					:  Optional[Any]					   =					cutoffs + [n_token]
   UpperCAmelCase__					:  List[Any]					   =					[0] + self.cutoffs
   UpperCAmelCase__					:  Any					   =					div_val

   UpperCAmelCase__					:  List[Any]					   =					self.cutoffs[0]
   UpperCAmelCase__					:  Tuple					   =					len(self.cutoffs    ) - 1
   UpperCAmelCase__					:  Optional[Any]					   =					self.shortlist_size + self.n_clusters

   if self.n_clusters > 0:
    UpperCAmelCase__					:  List[Any]					   =					nn.Parameter(torch.zeros(self.n_clusters					,self.d_embed    )    )
    UpperCAmelCase__					:  Dict					   =					nn.Parameter(torch.zeros(self.n_clusters    )    )

   UpperCAmelCase__					:  Dict					   =					nn.ModuleList()
   UpperCAmelCase__					:  int					   =					nn.ParameterList()

   if div_val == 1:
    for i in range(len(self.cutoffs    )    ):
     if d_proj != d_embed:
      self.out_projs.append(nn.Parameter(torch.FloatTensor(A					,A    )    )    )
     else:
      self.out_projs.append(A    )

    self.out_layers.append(nn.Linear(A					,A    )    )
   else:
    for i in range(len(self.cutoffs    )    ):
     UpperCAmelCase__					:  Tuple					   =					self.cutoff_ends[i], self.cutoff_ends[i + 1]
     UpperCAmelCase__					:  Optional[Any]					   =					d_embed // (div_val**i)

     self.out_projs.append(nn.Parameter(torch.FloatTensor(A					,A    )    )    )

     self.out_layers.append(nn.Linear(A					,r_idx - l_idx    )    )

   UpperCAmelCase__					:  Dict					   =					keep_order

  def    __lowercase       (  self		:			List[Any]					,A		:			Union[str, Any]					,A		:			Dict					,A		:			Optional[int]					,A		:			int    ):






   '''simple docstring'''
   if proj is None:
    UpperCAmelCase__					:  List[Any]					   =					nn.functional.linear(A					,A					,bias=A    )
   else:
    # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
    UpperCAmelCase__					:  Dict					   =					nn.functional.linear(A					,proj.t().contiguous()    )
    UpperCAmelCase__					:  Any					   =					nn.functional.linear(A					,A					,bias=A    )
    # else:
    #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
    #     if bias is not None:
    #         logit = logit + bias

   return logit

  def    __lowercase       (  self		:			str					,A		:			Union[str, Any]					,A		:			Optional[Any]=None					,A		:			Optional[Any]=False    ):






   '''simple docstring'''

   if labels is not None:
    # Shift so that tokens < n predict n
    UpperCAmelCase__					:  Any					   =					hidden[..., :-1, :].contiguous()
    UpperCAmelCase__					:  Dict					   =					labels[..., 1:].contiguous()
    UpperCAmelCase__					:  str					   =					hidden.view(-1					,hidden.size(-1    )    )
    UpperCAmelCase__					:  Tuple					   =					labels.view(-1    )
    if hidden.size(0    ) != labels.size(0    ):
     raise RuntimeError("""Input and labels should have the same size in the batch dimension."""    )
   else:
    UpperCAmelCase__					:  str					   =					hidden.view(-1					,hidden.size(-1    )    )

   if self.n_clusters == 0:
    UpperCAmelCase__					:  Union[str, Any]					   =					self._compute_logit(A					,self.out_layers[0].weight					,self.out_layers[0].bias					,self.out_projs[0]    )
    if labels is not None:
     UpperCAmelCase__					:  Optional[int]					   =					labels != -100
     UpperCAmelCase__					:  List[Any]					   =					torch.zeros_like(A					,dtype=hidden.dtype					,device=hidden.device    )
     UpperCAmelCase__					:  List[str]					   =					(
         -nn.functional.log_softmax(A					,dim=-1    )[mask].gather(1					,labels[mask].unsqueeze(1    )    ).squeeze(1    )
     )
    else:
     UpperCAmelCase__					:  Tuple					   =					nn.functional.log_softmax(A					,dim=-1    )
   else:
    # construct weights and biases
    UpperCAmelCase__					:  Optional[int]					   =					[], []
    for i in range(len(self.cutoffs    )    ):
     if self.div_val == 1:
      UpperCAmelCase__					:  Tuple					   =					self.cutoff_ends[i], self.cutoff_ends[i + 1]
      UpperCAmelCase__					:  Optional[Any]					   =					self.out_layers[0].weight[l_idx:r_idx]
      UpperCAmelCase__					:  int					   =					self.out_layers[0].bias[l_idx:r_idx]
     else:
      UpperCAmelCase__					:  Optional[int]					   =					self.out_layers[i].weight
      UpperCAmelCase__					:  Union[str, Any]					   =					self.out_layers[i].bias

     if i == 0:
      UpperCAmelCase__					:  List[str]					   =					torch.cat([weight_i, self.cluster_weight]					,dim=0    )
      UpperCAmelCase__					:  Optional[int]					   =					torch.cat([bias_i, self.cluster_bias]					,dim=0    )

     weights.append(A    )
     biases.append(A    )

    UpperCAmelCase__					:  Any					   =					weights[0], biases[0], self.out_projs[0]

    UpperCAmelCase__					:  Union[str, Any]					   =					self._compute_logit(A					,A					,A					,A    )
    UpperCAmelCase__					:  Optional[Any]					   =					nn.functional.log_softmax(A					,dim=1    )

    if labels is None:
     UpperCAmelCase__					:  Dict					   =					hidden.new_empty((head_logit.size(0    ), self.n_token)    )
    else:
     UpperCAmelCase__					:  Optional[Any]					   =					torch.zeros_like(A					,dtype=hidden.dtype					,device=hidden.device    )

    UpperCAmelCase__					:  Tuple					   =					0
    UpperCAmelCase__					:  Any					   =					[0] + self.cutoffs
    for i in range(len(A    ) - 1    ):
     UpperCAmelCase__					:  Optional[Any]					   =					cutoff_values[i], cutoff_values[i + 1]

     if labels is not None:
      UpperCAmelCase__					:  Optional[int]					   =					(labels >= l_idx) & (labels < r_idx)
      UpperCAmelCase__					:  Optional[int]					   =					mask_i.nonzero().squeeze()

      if indices_i.numel() == 0:
       continue

      UpperCAmelCase__					:  Dict					   =					labels.index_select(0					,A    ) - l_idx
      UpperCAmelCase__					:  List[Any]					   =					head_logprob.index_select(0					,A    )
      UpperCAmelCase__					:  List[str]					   =					hidden.index_select(0					,A    )
     else:
      UpperCAmelCase__					:  Any					   =					hidden

     if i == 0:
      if labels is not None:
       UpperCAmelCase__					:  Any					   =					head_logprob_i.gather(1					,target_i[:, None]    ).squeeze(1    )
      else:
       UpperCAmelCase__					:  int					   =					head_logprob[:, : self.cutoffs[0]]
     else:
      UpperCAmelCase__					:  Union[str, Any]					   =					weights[i], biases[i], self.out_projs[i]

      UpperCAmelCase__					:  int					   =					self._compute_logit(A					,A					,A					,A    )
      UpperCAmelCase__					:  List[Any]					   =					nn.functional.log_softmax(A					,dim=1    )
      UpperCAmelCase__					:  Tuple					   =					self.cutoffs[0] + i - 1  # No probability for the head cluster
      if labels is not None:
       UpperCAmelCase__					:  Tuple					   =					head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
           1					,target_i[:, None]    ).squeeze(1    )
      else:
       UpperCAmelCase__					:  Tuple					   =					head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
       UpperCAmelCase__					:  Any					   =					logprob_i

     if labels is not None:
      if (hasattr(self					,"""keep_order"""    ) and self.keep_order) or keep_order:
       out.index_copy_(0					,A					,-logprob_i    )
      else:
       out[offset : offset + logprob_i.size(0    )].copy_(-logprob_i    )
      offset += logprob_i.size(0    )

   return out




  def    __lowercase       (  self		:			Union[str, Any]					,A		:			List[str]    ):






   '''simple docstring'''
   if self.n_clusters == 0:
    UpperCAmelCase__					:  int					   =					self._compute_logit(A					,self.out_layers[0].weight					,self.out_layers[0].bias					,self.out_projs[0]    )
    return nn.functional.log_softmax(A					,dim=-1    )
   else:
    # construct weights and biases
    UpperCAmelCase__					:  Optional[Any]					   =					[], []
    for i in range(len(self.cutoffs    )    ):
     if self.div_val == 1:
      UpperCAmelCase__					:  Tuple					   =					self.cutoff_ends[i], self.cutoff_ends[i + 1]
      UpperCAmelCase__					:  Any					   =					self.out_layers[0].weight[l_idx:r_idx]
      UpperCAmelCase__					:  Optional[Any]					   =					self.out_layers[0].bias[l_idx:r_idx]
     else:
      UpperCAmelCase__					:  Optional[int]					   =					self.out_layers[i].weight
      UpperCAmelCase__					:  int					   =					self.out_layers[i].bias

     if i == 0:
      UpperCAmelCase__					:  List[str]					   =					torch.cat([weight_i, self.cluster_weight]					,dim=0    )
      UpperCAmelCase__					:  List[Any]					   =					torch.cat([bias_i, self.cluster_bias]					,dim=0    )

     weights.append(A    )
     biases.append(A    )

    UpperCAmelCase__					:  Tuple					   =					weights[0], biases[0], self.out_projs[0]
    UpperCAmelCase__					:  int					   =					self._compute_logit(A					,A					,A					,A    )

    UpperCAmelCase__					:  Union[str, Any]					   =					hidden.new_empty((head_logit.size(0    ), self.n_token)    )
    UpperCAmelCase__					:  Dict					   =					nn.functional.log_softmax(A					,dim=1    )

    UpperCAmelCase__					:  Union[str, Any]					   =					[0] + self.cutoffs
    for i in range(len(A    ) - 1    ):
     UpperCAmelCase__					:  Union[str, Any]					   =					cutoff_values[i], cutoff_values[i + 1]

     if i == 0:
      UpperCAmelCase__					:  Optional[int]					   =					head_logprob[:, : self.cutoffs[0]]
     else:
      UpperCAmelCase__					:  int					   =					weights[i], biases[i], self.out_projs[i]

      UpperCAmelCase__					:  Optional[Any]					   =					self._compute_logit(A					,A					,A					,A    )
      UpperCAmelCase__					:  List[Any]					   =					nn.functional.log_softmax(A					,dim=1    )

      UpperCAmelCase__					:  Any					   =					head_logprob[:, -i] + tail_logprob_i
      UpperCAmelCase__					:  Dict					   =					logprob_i

    return out








"""simple docstring"""




def   lowerCAmelCase						(					__UpperCamelCase = 1000      ):
 '''simple docstring'''





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


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





