


import torch
from torch import nn


class 	A_ (	nn.Module   ):




      '''simple docstring'''


      def __init__(    self:	Any						,				a:	Any						,				a:	Optional[int]						,				a:	str						,				a:	Dict						,				a:	Optional[int]=1						,				a:	Union[str, Any]=False    ):
       super().__init__()

       __lowerCamelCase		:							str		 =						n_token
       __lowerCamelCase		:							str		 =						d_embed
       __lowerCamelCase		:							List[str]		 =						d_proj

       __lowerCamelCase		:							Union[str, Any]		 =						cutoffs + [n_token]
       __lowerCamelCase		:							List[str]		 =						[0] + self.cutoffs
       __lowerCamelCase		:							Tuple		 =						div_val

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

       if self.n_clusters > 0:
        __lowerCamelCase		:							Tuple		 =						nn.Parameter(torch.zeros(self.n_clusters						,				self.d_embed    )    )
        __lowerCamelCase		:							Optional[int]		 =						nn.Parameter(torch.zeros(self.n_clusters    )    )

       __lowerCamelCase		:							Optional[Any]		 =						nn.ModuleList()
       __lowerCamelCase		:							List[str]		 =						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    )    ):
         __lowerCamelCase ,					__lowerCamelCase		:							Any		 =						self.cutoff_ends[i], self.cutoff_ends[i + 1]
         __lowerCamelCase		:							Optional[int]		 =						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    )    )

       __lowerCamelCase		:							int		 =						keep_order




      def 					_snake_case     (    self:	List[str]						,				a:	List[str]						,				a:	Optional[Any]						,				a:	int						,				a:	int    ):
       if proj is None:
        __lowerCamelCase		:							Any		 =						nn.functional.linear(a						,				a						,				bias=a    )
       else:
        # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
        __lowerCamelCase		:							int		 =						nn.functional.linear(a						,				proj.t().contiguous()    )
        __lowerCamelCase		:							Optional[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 					_snake_case     (    self:	Tuple						,				a:	str						,				a:	Any=None						,				a:	List[str]=False    ):

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

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

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

         weights.append(a    )
         biases.append(a    )

        __lowerCamelCase ,					__lowerCamelCase ,					__lowerCamelCase		:							Tuple		 =						weights[0], biases[0], self.out_projs[0]

        __lowerCamelCase		:							str		 =						self._compute_logit(a						,				a						,				a						,				a    )
        __lowerCamelCase		:							Any		 =						nn.functional.log_softmax(a						,				dim=1    )

        if labels is None:
         __lowerCamelCase		:							Optional[Any]		 =						hidden.new_empty((head_logit.size(0    ), self.n_token)    )
        else:
         __lowerCamelCase		:							Union[str, Any]		 =						torch.zeros_like(a						,				dtype=hidden.dtype						,				device=hidden.device    )

        __lowerCamelCase		:							Dict		 =						0
        __lowerCamelCase		:							Any		 =						[0] + self.cutoffs
        for i in range(len(a    ) - 1    ):
         __lowerCamelCase ,					__lowerCamelCase		:							Any		 =						cutoff_values[i], cutoff_values[i + 1]

         if labels is not None:
          __lowerCamelCase		:							List[Any]		 =						(labels >= l_idx) & (labels < r_idx)
          __lowerCamelCase		:							Tuple		 =						mask_i.nonzero().squeeze()

          if indices_i.numel() == 0:
           continue

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

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

          __lowerCamelCase		:							Tuple		 =						self._compute_logit(a						,				a						,				a						,				a    )
          __lowerCamelCase		:							Any		 =						nn.functional.log_softmax(a						,				dim=1    )
          __lowerCamelCase		:							List[str]		 =						self.cutoffs[0] + i - 1  # No probability for the head cluster
          if labels is not None:
           __lowerCamelCase		:							int		 =						head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
               1						,				target_i[:, None]    ).squeeze(1    )
          else:
           __lowerCamelCase		:							Tuple		 =						head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
           __lowerCamelCase		:							int		 =						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 					_snake_case     (    self:	str						,				a:	str    ):
       if self.n_clusters == 0:
        __lowerCamelCase		:							Optional[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
        __lowerCamelCase ,					__lowerCamelCase		:							List[Any]		 =						[], []
        for i in range(len(self.cutoffs    )    ):
         if self.div_val == 1:
          __lowerCamelCase ,					__lowerCamelCase		:							Optional[Any]		 =						self.cutoff_ends[i], self.cutoff_ends[i + 1]
          __lowerCamelCase		:							Union[str, Any]		 =						self.out_layers[0].weight[l_idx:r_idx]
          __lowerCamelCase		:							Dict		 =						self.out_layers[0].bias[l_idx:r_idx]
         else:
          __lowerCamelCase		:							List[str]		 =						self.out_layers[i].weight
          __lowerCamelCase		:							str		 =						self.out_layers[i].bias

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

         weights.append(a    )
         biases.append(a    )

        __lowerCamelCase ,					__lowerCamelCase ,					__lowerCamelCase		:							Dict		 =						weights[0], biases[0], self.out_projs[0]
        __lowerCamelCase		:							Any		 =						self._compute_logit(a						,				a						,				a						,				a    )

        __lowerCamelCase		:							Tuple		 =						hidden.new_empty((head_logit.size(0    ), self.n_token)    )
        __lowerCamelCase		:							int		 =						nn.functional.log_softmax(a						,				dim=1    )

        __lowerCamelCase		:							List[Any]		 =						[0] + self.cutoffs
        for i in range(len(a    ) - 1    ):
         __lowerCamelCase ,					__lowerCamelCase		:							int		 =						cutoff_values[i], cutoff_values[i + 1]

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

          __lowerCamelCase		:							Tuple		 =						self._compute_logit(a						,				a						,				a						,				a    )
          __lowerCamelCase		:							Tuple		 =						nn.functional.log_softmax(a						,				dim=1    )

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

        return out





from __future__ import annotations
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							List[Any]		 =						str(SCREAMING_SNAKE_CASE__						)
 return len(SCREAMING_SNAKE_CASE__						) == 9 and set(SCREAMING_SNAKE_CASE__						) == set('123456789'						)
def 				UpperCamelCase__					(  ):
 for base_num in range(9_999      ,	4_999      ,	-1						):
  __lowerCamelCase		:							Tuple		 =						100_002 * base_num
  if is_9_pandigital(SCREAMING_SNAKE_CASE__						):
   return candidate

 for base_num in range(333      ,	99      ,	-1						):
  __lowerCamelCase		:							Union[str, Any]		 =						1_002_003 * base_num
  if is_9_pandigital(SCREAMING_SNAKE_CASE__						):
   return candidate

 return None


if __name__ == "__main__":
       print(F"""{solution() = }""")

