



"""simple docstring"""






import torch
from torch import nn



class     UpperCAmelCase    (	nn.Module  ):



   def __init__(       self    :	List[str]							,       __lowerCamelCase    :	Dict							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	Optional[Any]							,       __lowerCamelCase    :	Optional[Any]							,       __lowerCamelCase    :	Tuple=1							,       __lowerCamelCase    :	Optional[int]=False				):



     """simple docstring"""
     super().__init__()

     _snake_case   	=      n_token
     _snake_case   	=      d_embed
     _snake_case   	=      d_proj

     _snake_case   	=      cutoffs + [n_token]
     _snake_case   	=      [0] + self.cutoffs
     _snake_case   	=      div_val

     _snake_case   	=      self.cutoffs[0]
     _snake_case   	=      len(self.cutoffs				) - 1
     _snake_case   	=      self.shortlist_size + self.n_clusters

     if self.n_clusters > 0:
       _snake_case   	=      nn.Parameter(torch.zeros(self.n_clusters							,       self.d_embed				)				)
       _snake_case   	=      nn.Parameter(torch.zeros(self.n_clusters				)				)

     _snake_case   	=      nn.ModuleList()
     _snake_case   	=      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				)				):
         _snake_case	,     _snake_case   	=      self.cutoff_ends[i], self.cutoff_ends[i + 1]
         _snake_case   	=      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				)				)

     _snake_case   	=      keep_order



   def    __UpperCAmelCase    (       self    :	Union[str, Any]							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	Union[str, Any]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	Optional[int]				):



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



     """simple docstring"""

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

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

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

         weights.append(__lowerCamelCase				)
         biases.append(__lowerCamelCase				)

       _snake_case	,     _snake_case	,     _snake_case   	=      weights[0], biases[0], self.out_projs[0]

       _snake_case   	=      self._compute_logit(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase				)
       _snake_case   	=      nn.functional.log_softmax(__lowerCamelCase							,       dim=1				)

       if labels is None:
         _snake_case   	=      hidden.new_empty((head_logit.size(0				), self.n_token)				)
       else:
         _snake_case   	=      torch.zeros_like(__lowerCamelCase							,       dtype=hidden.dtype							,       device=hidden.device				)

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

         if labels is not None:
           _snake_case   	=      (labels >= l_idx) & (labels < r_idx)
           _snake_case   	=      mask_i.nonzero().squeeze()

           if indices_i.numel() == 0:
             continue

           _snake_case   	=      labels.index_select(0							,       __lowerCamelCase				) - l_idx
           _snake_case   	=      head_logprob.index_select(0							,       __lowerCamelCase				)
           _snake_case   	=      hidden.index_select(0							,       __lowerCamelCase				)
         else:
           _snake_case   	=      hidden

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

           _snake_case   	=      self._compute_logit(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase				)
           _snake_case   	=      nn.functional.log_softmax(__lowerCamelCase							,       dim=1				)
           _snake_case   	=      self.cutoffs[0] + i - 1  # No probability for the head cluster
           if labels is not None:
             _snake_case   	=      head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                 1							,       target_i[:, None]				).squeeze(1				)
           else:
             _snake_case   	=      head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
             _snake_case   	=      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    __UpperCAmelCase    (       self    :	str							,       __lowerCamelCase    :	Tuple				):



     """simple docstring"""
     if self.n_clusters == 0:
       _snake_case   	=      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
       _snake_case	,     _snake_case   	=      [], []
       for i in range(len(self.cutoffs				)				):
         if self.div_val == 1:
           _snake_case	,     _snake_case   	=      self.cutoff_ends[i], self.cutoff_ends[i + 1]
           _snake_case   	=      self.out_layers[0].weight[l_idx:r_idx]
           _snake_case   	=      self.out_layers[0].bias[l_idx:r_idx]
         else:
           _snake_case   	=      self.out_layers[i].weight
           _snake_case   	=      self.out_layers[i].bias

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

         weights.append(__lowerCamelCase				)
         biases.append(__lowerCamelCase				)

       _snake_case	,     _snake_case	,     _snake_case   	=      weights[0], biases[0], self.out_projs[0]
       _snake_case   	=      self._compute_logit(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase				)

       _snake_case   	=      hidden.new_empty((head_logit.size(0				), self.n_token)				)
       _snake_case   	=      nn.functional.log_softmax(__lowerCamelCase							,       dim=1				)

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

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

           _snake_case   	=      self._compute_logit(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase				)
           _snake_case   	=      nn.functional.log_softmax(__lowerCamelCase							,       dim=1				)

           _snake_case   	=      head_logprob[:, -i] + tail_logprob_i
           _snake_case   	=      logprob_i

       return out










"""simple docstring"""






from itertools import product




def 			snake_case		(		lowerCAmelCase_  ,						lowerCAmelCase_	)						->				list[int]:
  _snake_case   	=      sides_number
  _snake_case   	=      max_face_number * dice_number
  _snake_case   	=      [0] * (max_total + 1)

  _snake_case   	=      1
  _snake_case   	=      range(lowerCAmelCase_  ,						max_face_number + 1	)
  for dice_numbers in product(lowerCAmelCase_  ,						repeat=lowerCAmelCase_	):
    _snake_case   	=      sum(lowerCAmelCase_	)
    totals_frequencies[total] += 1

  return totals_frequencies




def 			snake_case		(		)						->				float:
  _snake_case   	=      total_frequency_distribution(
      sides_number=4  ,						dice_number=9	)
  _snake_case   	=      total_frequency_distribution(
      sides_number=6  ,						dice_number=6	)

  _snake_case   	=      0
  _snake_case   	=      9
  _snake_case   	=      4 * 9
  _snake_case   	=      6
  for peter_total in range(lowerCAmelCase_  ,						max_peter_total + 1	):
    peter_wins_count += peter_totals_frequencies[peter_total] * sum(
        colin_totals_frequencies[min_colin_total:peter_total]	)

  _snake_case   	=      (4**9) * (6**6)
  _snake_case   	=      peter_wins_count / total_games_number

  _snake_case   	=      round(lowerCAmelCase_  ,						ndigits=7	)

  return rounded_peter_win_probability


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





