code: import math
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.autograd import Variable
  
  from embed_regularize import embedded_dropout
  from locked_dropout import LockedDropout
  from weight_drop import WeightDrop
  
  class RNNModel_v2(nn.Module):
      """Container module with an encoder, a recurrent module, and a decoder."""
  
      def __init__(self, rnn_type, ntoken, ninp, nhid, nhidlast, nlayers, 
                   dropout=0.5, dropouth=0.5, dropouti=0.5, dropoute=0.1, wdrop=0, 
                   tie_weights=False, ldropout=0.5, n_experts=10):
          super(RNNModel_v2, self).__init__()
          self.lockdrop = LockedDropout()
          self.encoder = nn.Embedding(ntoken*n_experts, ninp)
  
          self.rnns = [torch.nn.LSTM(ninp if l == 0 else nhid, nhid if l != nlayers - 1 else nhidlast, 1, dropout=0) for l in range(nlayers)]
          if wdrop:
              self.rnns = [WeightDrop(rnn, ['weight_hh_l0'], dropout=wdrop) for rnn in self.rnns]
          self.rnns = torch.nn.ModuleList(self.rnns)
  
          self.prior = nn.Linear(ninp, 1, bias=False)
          #self.latent = nn.Sequential(nn.Linear(nhidlast, ninp), nn.Tanh())
          self.decoder = nn.Linear(ninp, ntoken*n_experts)
          #self.proj_expert = nn.Linear(ntoken*n_experts, ntoken)
          #print(ninp, ntoken)
          # Optionally tie weights as in:
          # "Using the Output Embedding to Improve Language Models" (Press & Wolf 2016)
          # https://arxiv.org/abs/1608.05859
          # and
          # "Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling" (Inan et al. 2016)
          # https://arxiv.org/abs/1611.01462
          if tie_weights:
              #if nhid != ninp:
              #    raise ValueError('When using the tied flag, nhid must be equal to emsize')
              self.decoder.weight = self.encoder.weight
  
          self.init_weights()
  
          self.rnn_type = rnn_type
          self.ninp = ninp
          self.nhid = nhid
          self.nhidlast = nhidlast
          self.nlayers = nlayers
          self.dropout = dropout
          self.dropouti = dropouti
          self.dropouth = dropouth
          self.dropoute = dropoute
          self.ldropout = ldropout
          self.dropoutl = ldropout
          self.n_experts = n_experts
          self.ntoken = ntoken
  
          size = 0
          for p in self.parameters():
              size += p.nelement()
          print('param size: {}'.format(size))
  
      def init_weights(self):
          initrange = 0.1
          self.encoder.weight.data.uniform_(-initrange, initrange)
          self.decoder.bias.data.fill_(0)
          self.decoder.weight.data.uniform_(-initrange, initrange)
  
      def forward(self, input, hidden, return_h=False, return_prob=False):
          length, batch_size = input.size()
          #print(input)
          input = input * self.n_experts
          # (length, batch_size * n_experts)
          input_experts = torch.cat([input+i for i in range(self.n_experts)], dim=1)
          # (length, batch_size * n_experts, emb_dim)
          #print(input_experts)
          emb = embedded_dropout(self.encoder, input_experts, dropout=self.dropoute if self.training else 0)
          # (length, batch_size * n_experts, emb_dim)
          emb = self.lockdrop(emb, self.dropouti)
          # (length, batch_size, n_experts, emb_dim)
          emb = emb.view(length, self.n_experts, batch_size, -1).permute(0, 2, 1, 3).contiguous()
          # (length, batch_size, n_experts, 1)
          prior_logit = self.prior(emb).squeeze(3)
          # (length, batch_size, n_experts)
          prior = nn.functional.softmax(prior_logit.view(-1, self.n_experts)).view(-1, batch_size, self.n_experts)
          emb = (emb * prior.unsqueeze(3)).sum(2)
  
          raw_output = emb
          new_hidden = []
          #raw_output, hidden = self.rnn(emb, hidden)
          raw_outputs = []
          outputs = []
          for l, rnn in enumerate(self.rnns):
              current_input = raw_output
              raw_output, new_h = rnn(raw_output, hidden[l])
              new_hidden.append(new_h)
              raw_outputs.append(raw_output)
              if l != self.nlayers - 1:
                  #self.hdrop(raw_output)
                  raw_output = self.lockdrop(raw_output, self.dropouth)
                  outputs.append(raw_output)
          hidden = new_hidden
  
          output = self.lockdrop(raw_output, self.dropout) # (length, batch_size, dim_last)
          outputs.append(output)
  
          #print(output)
          logit = self.decoder(output.view(-1, self.ninp)) # (length*batch_size, ntok*n_expert)
  
          prob = nn.functional.softmax(logit).view(-1, batch_size, self.ntoken * self.n_experts)
          #prob = (prob * prior.unsqueeze(3).permute(0, 2, 1, 3)).sum(1) # (length, batch_size, ntoken*n_experts)
          prob = prob.view(-1, self.n_experts, self.ntoken).sum(1)
          if return_prob:
              model_output = prob
          else:
              log_prob = torch.log(prob.add_(1e-8))
              model_output = log_prob
  
          model_output = model_output.view(-1, batch_size, self.ntoken)
  
          if return_h:
              return model_output, hidden, raw_outputs, outputs
          return model_output, hidden
  
      def init_hidden(self, bsz):
          weight = next(self.parameters()).data
          return [(Variable(weight.new(1, bsz, self.nhid if l != self.nlayers - 1 else self.nhidlast).zero_()),
                   Variable(weight.new(1, bsz, self.nhid if l != self.nlayers - 1 else self.nhidlast).zero_()))
                  for l in range(self.nlayers)]
  
  if __name__ == '__main__':
      model = RNNModel('LSTM', 10, 12, 12, 12, 2)
      input = Variable(torch.LongTensor(13, 9).random_(0, 10))
      hidden = model.init_hidden(9)
      model(input, hidden)
  
      # input = Variable(torch.LongTensor(13, 9).random_(0, 10))
      # hidden = model.init_hidden(9)
      # print(model.sample(input, hidden, 5, 6, 1, 2, sample_latent=True).size())
  
  
apis: torch.nn.LSTM torch.nn.Linear torch.nn.functional.softmax torch.nn.Embedding torch.nn.ModuleList torch.LongTensor