
import torch
from torch import nn



class 		UpperCamelCase__  (       nn.Module    ):
   def __init__(self :					int					,				snake_case_ :					str					,				snake_case_ :					Dict					,				snake_case_ :					Dict					,				snake_case_ :					List[str]					,				snake_case_ :					Optional[Any]=1					,				snake_case_ :					Any=False   ):
    super().__init__()

    __a :  Optional[Any]						=  n_token
    __a :  Dict						=  d_embed
    __a :  int						=  d_proj

    __a :  List[str]						=  cutoffs + [n_token]
    __a :  Any						=  [0] + self.cutoffs
    __a :  Optional[int]						=  div_val

    __a :  str						=  self.cutoffs[0]
    __a :  List[Any]						=  len(self.cutoffs   ) - 1
    __a :  str						=  self.shortlist_size + self.n_clusters

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

    __a :  Any						=  nn.ModuleList()
    __a :  Any						=  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(lowerCAmelCase_					,				lowerCAmelCase_   )   )   )
      else:
       self.out_projs.append(lowerCAmelCase_   )

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

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

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

    __a :  Dict						=  keep_order
   def       lowerCAmelCase			(self :					Tuple					,				snake_case_ :					Optional[Any]					,				snake_case_ :					str					,				snake_case_ :					Optional[Any]					,				snake_case_ :					List[Any]   ):
    if proj is None:
     __a :  Union[str, Any]						=  nn.functional.linear(lowerCAmelCase_					,				lowerCAmelCase_					,				bias=lowerCAmelCase_   )
    else:
     # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
     __a :  Tuple						=  nn.functional.linear(lowerCAmelCase_					,				proj.t().contiguous()   )
     __a :  List[Any]						=  nn.functional.linear(lowerCAmelCase_					,				lowerCAmelCase_					,				bias=lowerCAmelCase_   )
     # else:
     #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
     #     if bias is not None:
     #         logit = logit + bias

    return logit
   def       lowerCAmelCase			(self :					Optional[Any]					,				snake_case_ :					Optional[Any]					,				snake_case_ :					int=None					,				snake_case_ :					Any=False   ):

    if labels is not None:
     # Shift so that tokens < n predict n
     __a :  int						=  hidden[..., :-1, :].contiguous()
     __a :  Optional[int]						=  labels[..., 1:].contiguous()
     __a :  str						=  hidden.view(-1					,				hidden.size(-1   )   )
     __a :  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:
     __a :  List[Any]						=  hidden.view(-1					,				hidden.size(-1   )   )

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

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

      weights.append(lowerCAmelCase_   )
      biases.append(lowerCAmelCase_   )

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

     __a :  str						=  self._compute_logit(lowerCAmelCase_					,				lowerCAmelCase_					,				lowerCAmelCase_					,				lowerCAmelCase_   )
     __a :  Optional[Any]						=  nn.functional.log_softmax(lowerCAmelCase_					,				dim=1   )

     if labels is None:
      __a :  List[Any]						=  hidden.new_empty((head_logit.size(0   ), self.n_token)   )
     else:
      __a :  List[Any]						=  torch.zeros_like(lowerCAmelCase_					,				dtype=hidden.dtype					,				device=hidden.device   )

     __a :  Any						=  0
     __a :  Optional[int]						=  [0] + self.cutoffs
     for i in range(len(lowerCAmelCase_   ) - 1   ):
      __a				,       __a :  Any						=  cutoff_values[i], cutoff_values[i + 1]

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

       if indices_i.numel() == 0:
        continue

       __a :  Union[str, Any]						=  labels.index_select(0					,				lowerCAmelCase_   ) - l_idx
       __a :  Union[str, Any]						=  head_logprob.index_select(0					,				lowerCAmelCase_   )
       __a :  Optional[int]						=  hidden.index_select(0					,				lowerCAmelCase_   )
      else:
       __a :  List[str]						=  hidden

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

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

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

    return out







   def       lowerCAmelCase			(self :					Optional[int]					,				snake_case_ :					List[str]   ):
    if self.n_clusters == 0:
     __a :  List[Any]						=  self._compute_logit(lowerCAmelCase_					,				self.out_layers[0].weight					,				self.out_layers[0].bias					,				self.out_projs[0]   )
     return nn.functional.log_softmax(lowerCAmelCase_					,				dim=-1   )
    else:
     # construct weights and biases
     __a				,       __a :  Optional[int]						=  [], []
     for i in range(len(self.cutoffs   )   ):
      if self.div_val == 1:
       __a				,       __a :  Optional[int]						=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
       __a :  Optional[int]						=  self.out_layers[0].weight[l_idx:r_idx]
       __a :  Optional[int]						=  self.out_layers[0].bias[l_idx:r_idx]
      else:
       __a :  Dict						=  self.out_layers[i].weight
       __a :  Optional[int]						=  self.out_layers[i].bias

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

      weights.append(lowerCAmelCase_   )
      biases.append(lowerCAmelCase_   )

     __a				,       __a				,       __a :  str						=  weights[0], biases[0], self.out_projs[0]
     __a :  Dict						=  self._compute_logit(lowerCAmelCase_					,				lowerCAmelCase_					,				lowerCAmelCase_					,				lowerCAmelCase_   )

     __a :  str						=  hidden.new_empty((head_logit.size(0   ), self.n_token)   )
     __a :  Any						=  nn.functional.log_softmax(lowerCAmelCase_					,				dim=1   )

     __a :  Tuple						=  [0] + self.cutoffs
     for i in range(len(lowerCAmelCase_   ) - 1   ):
      __a				,       __a :  List[str]						=  cutoff_values[i], cutoff_values[i + 1]

      if i == 0:
       __a :  Dict						=  head_logprob[:, : self.cutoffs[0]]
      else:
       __a				,       __a				,       __a :  List[str]						=  weights[i], biases[i], self.out_projs[i]

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

       __a :  Optional[Any]						=  head_logprob[:, -i] + tail_logprob_i
       __a :  Optional[int]						=  logprob_i

     return out







def   a_    ( lowerCAmelCase_      :			int	):

   if p < 2:
      raise ValueError('p should not be less than 2!'	)
   elif p == 2:
      return True

   __lowerCAmelCase								=	4
   __lowerCAmelCase								=	(1 << p) - 1
   for _ in range(p - 2	):
      __lowerCAmelCase								=	((s * s) - 2) % m
   return s == 0


if __name__ == "__main__":
     print(lucas_lehmer_test(7))
     print(lucas_lehmer_test(11))

