text: """
  adapted from fairseq.models.lstm to work with pytext and separate
  attended states from the initialization.
  The initial states now have to be passed explicitly.
  """
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from typing import Optional, Dict, Any
  
  from fairseq import utils
  from fairseq.modules import AdaptiveSoftmax
  from fairseq.models import FairseqIncrementalDecoder
  from fairseq.models.lstm import LSTMCell, Linear, Embedding, AttentionLayer
  from fairseq.data.dictionary import Dictionary
  
  
  class HREDDecoderImpl(FairseqIncrementalDecoder):
      def __init__(
          self, dictionary: Dictionary, embed_dim: int = 512, hidden_size: int = 512, out_embed_dim: int = 512,
          num_layers: int = 1, dropout_in: float = 0.1, dropout_out: float = 0.1, attention: bool = True,
          encoder_embed_dim: int = 512, encoder_output_units: int = 512, pretrained_embed: Optional[nn.Embedding] = None,
          share_input_output_embed: bool = False, adaptive_softmax_cutoff: Optional[int] = None,
      ):
          super().__init__(dictionary)
          self.dropout_in = dropout_in
          self.dropout_out = dropout_out
          self.hidden_size = hidden_size
          self.share_input_output_embed = share_input_output_embed
          self.need_attn = True
  
          self.adaptive_softmax = None
          num_embeddings = len(dictionary)
          padding_idx = dictionary.pad()
          if pretrained_embed is None:
              self.embed_tokens = Embedding(num_embeddings, embed_dim, padding_idx)
          else:
              self.embed_tokens = pretrained_embed
  
          self.encoder_output_units = encoder_output_units
  
          self.layers = nn.ModuleList([
              LSTMCell(
                  input_size=hidden_size + embed_dim if layer == 0 else hidden_size,
                  hidden_size=hidden_size,
              )
              for layer in range(num_layers)
          ])
          self.attention = AttentionLayer(hidden_size, encoder_output_units, hidden_size) if attention else None
          if hidden_size != out_embed_dim:
              self.additional_fc = Linear(hidden_size, out_embed_dim)
          if adaptive_softmax_cutoff is not None:
              # setting adaptive_softmax dropout to dropout_out for now but can be redefined
              self.adaptive_softmax = AdaptiveSoftmax(num_embeddings, embed_dim, adaptive_softmax_cutoff,
                                                      dropout=dropout_out)
          elif not self.share_input_output_embed:
              self.fc_out = Linear(out_embed_dim, num_embeddings, dropout=dropout_out)
  
      def forward(self,
                  prev_output_tokens: torch.Tensor,
                  encoder_out_dict: Dict[str, Any],
                  incremental_state: Dict[str, Any]):
          encoder_out = encoder_out_dict['encoder_out']
          encoder_padding_mask = encoder_out_dict['encoder_padding_mask']
  
          if incremental_state is not None:
              prev_output_tokens = prev_output_tokens[:, -1:]
          bsz, seqlen = prev_output_tokens.size()
  
          # get outputs from encoder
          encoder_outs, _, _ = encoder_out[:3]
          srclen = encoder_outs.size(1)
  
          # embed tokens
          x = self.embed_tokens(prev_output_tokens)
          x = F.dropout(x, p=self.dropout_in, training=self.training)
  
          # B x T x C -> T x B x C
          x = x.transpose(0, 1)
  
          # initialize previous states (or get from cache during incremental generation)
          cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state')
          if cached_state is not None:
            prev_hiddens, prev_cells, input_feed = cached_state
          else:
            initial_state = encoder_out_dict['decoder_initial_state']
            assert initial_state is not None
            encoder_hiddens, encoder_cells = initial_state
            num_layers = len(self.layers)
            prev_hiddens = [encoder_hiddens[i] for i in range(num_layers)]
            prev_cells = [encoder_cells[i] for i in range(num_layers)]
            input_feed = x.data.new(bsz, self.hidden_size).zero_()
  
          attn_scores = x.data.new(srclen, seqlen, bsz).zero_()
          outs = []
          for j in range(seqlen):
              # input feeding: concatenate context vector from previous time step
              input = torch.cat((x[j, :, :], input_feed), dim=1)
  
              for i, rnn in enumerate(self.layers):
                  # recurrent cell
                  hidden, cell = rnn(input, (prev_hiddens[i], prev_cells[i]))
  
                  # hidden state becomes the input to the next layer
                  input = F.dropout(hidden, p=self.dropout_out, training=self.training)
  
                  # save state for next time step
                  prev_hiddens[i] = hidden
                  prev_cells[i] = cell
  
              # apply attention using the last layer's hidden state
              if self.attention is not None:
                  out, attn_scores[:, j, :] = self.attention(hidden, encoder_outs.transpose(0, 1),
                                                             encoder_padding_mask.transpose(0, 1))
              else:
                  out = hidden
              out = F.dropout(out, p=self.dropout_out, training=self.training)
  
              # input feeding
              input_feed = out
  
              # save final output
              outs.append(out)
  
          # cache previous states (no-op except during incremental generation)
          utils.set_incremental_state(
              self, incremental_state, 'cached_state', (prev_hiddens, prev_cells, input_feed))
  
          # collect outputs across time steps
          x = torch.cat(outs, dim=0).view(seqlen, bsz, self.hidden_size)
  
          # T x B x C -> B x T x C
          x = x.transpose(1, 0)
  
          # srclen x tgtlen x bsz -> bsz x tgtlen x srclen
          if not self.training and self.need_attn:
              attn_scores = attn_scores.transpose(0, 2)
          else:
              attn_scores = None
  
          # project back to size of vocabulary
          if self.adaptive_softmax is None:
              if hasattr(self, 'additional_fc'):
                  x = self.additional_fc(x)
                  x = F.dropout(x, p=self.dropout_out, training=self.training)
              if self.share_input_output_embed:
                  x = F.linear(x, self.embed_tokens.weight)
              else:
                  x = self.fc_out(x)
          return x, attn_scores
  
      def reorder_incremental_state(self, incremental_state, new_order):
          super().reorder_incremental_state(incremental_state, new_order)
          cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state')
          if cached_state is None:
              return
  
          def reorder_state(state):
              if isinstance(state, list):
                  return [reorder_state(state_i) for state_i in state]
              return state.index_select(0, new_order)
  
          new_state = tuple(map(reorder_state, cached_state))
          utils.set_incremental_state(self, incremental_state, 'cached_state', new_state)
  
      def max_positions(self):
          """Maximum output length supported by the decoder."""
          return int(1e5)  # an arbitrary large number
  
      def make_generation_fast_(self, need_attn=False, **kwargs):
          self.need_attn = need_attn
  