class_docstr: `AugmentedLstm` implements a one-layer single directional
  AugmentedLSTM layer. AugmentedLSTM is an LSTM which optionally
  appends an optional highway network to the output layer. Furthermore the
  dropout controls the level of variational dropout done.
  
  # Parameters
  
  input_size : `int`
      The number of expected features in the input.
  hidden_size : `int`
      Number of features in the hidden state of the LSTM.
      Defaults to 32.
  go_forward : `bool`
      Whether to compute features left to right (forward)
      or right to left (backward).
  recurrent_dropout_probability : `float`
      Variational dropout probability to use. Defaults to 0.0.
  use_highway : `bool`
      If `True` we append a highway network to the outputs of the LSTM.
  use_input_projection_bias : `bool`
      If `True` we use a bias in our LSTM calculations, otherwise we don't.
  
  # Attributes
  
  cell : `AugmentedLSTMCell`
      `AugmentedLSTMCell` that is applied at every timestep.
class_docstr_tok: ["`","AugmentedLstm","`","implements","a","one-layer","single","directional","AugmentedLSTM","layer",".","AugmentedLSTM","is","an","LSTM","which","optionally","appends","an","optional","highway","network","to","the","output","layer",".","Furthermore","the","dropout","controls","the","level","of","variational","dropout","done",".","#","Parameters","input_size",":","`","int","`","The","number","of","expected","features","in","the","input",".","hidden_size",":","`","int","`","Number","of","features","in","the","hidden","state","of","the","LSTM",".","Defaults","to","32.","go_forward",":","`","bool","`","Whether","to","compute","features","left","to","right","(","forward",")","or","right","to","left","(","backward",")",".","recurrent_dropout_probability",":","`","float","`","Variational","dropout","probability","to","use",".","Defaults","to","0.0.","use_highway",":","`","bool","`","If","`","True","`","we","append","a","highway","network","to","the","outputs","of","the","LSTM",".","use_input_projection_bias",":","`","bool","`","If","`","True","`","we","use","a","bias","in","our","LSTM","calculations",",","otherwise","we","do","n't",".","#","Attributes","cell",":","`","AugmentedLSTMCell","`","`","AugmentedLSTMCell","`","that","is","applied","at","every","timestep","."]
docstr: Warning: Would be better to use the BiAugmentedLstm class in a regular model
  
  Given an input batch of sequential data such as word embeddings, produces a single layer unidirectional
  AugmentedLSTM representation of the sequential input and new state tensors.
  
  # Parameters
  
  inputs : `PackedSequence`
      `bsize` sequences of shape `(len, input_dim)` each, in PackedSequence format
  states : `Tuple[torch.Tensor, torch.Tensor]`
      Tuple of tensors containing the initial hidden state and
      the cell state of each element in the batch. Each of these tensors have a dimension of
      (1 x bsize x nhid). Defaults to `None`.
  
  # Returns
  
  `Tuple[PackedSequence, Tuple[torch.Tensor, torch.Tensor]]`
      AugmentedLSTM representation of input and the state of the LSTM `t = seq_len`.
      Shape of representation is (bsize x seq_len x representation_dim).
      Shape of each state is (1 x bsize x nhid).
docstr_tok: ["Warning",":","Would","be","better","to","use","the","BiAugmentedLstm","class","in","a","regular","model","Given","an","input","batch","of","sequential","data","such","as","word","embeddings",",","produces","a","single","layer","unidirectional","AugmentedLSTM","representation","of","the","sequential","input","and","new","state","tensors",".","#","Parameters","inputs",":","`","PackedSequence","`","`","bsize","`","sequences","of","shape","`","(","len",",","input_dim",")","`","each",",","in","PackedSequence","format","states",":","`","Tuple","[","torch.Tensor",",","torch.Tensor","]","`","Tuple","of","tensors","containing","the","initial","hidden","state","and","the","cell","state","of","each","element","in","the","batch",".","Each","of","these","tensors","have","a","dimension","of","(","1","x","bsize","x","nhid",")",".","Defaults","to","`","None","`",".","#","Returns","`","Tuple","[","PackedSequence",",","Tuple","[","torch.Tensor",",","torch.Tensor","]","]","`","AugmentedLSTM","representation","of","input","and","the","state","of","the","LSTM","`","t","=","seq_len","`",".","Shape","of","representation","is","(","bsize","x","seq_len","x","representation_dim",")",".","Shape","of","each","state","is","(","1","x","bsize","x","nhid",")","."]
code: def forward(
      self,
      inputs: PackedSequence,
      states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
  ) -> Tuple[PackedSequence, Tuple[torch.Tensor, torch.Tensor]]:
      """
      Warning: Would be better to use the BiAugmentedLstm class in a regular model
  
      Given an input batch of sequential data such as word embeddings, produces a single layer unidirectional
      AugmentedLSTM representation of the sequential input and new state tensors.
  
      # Parameters
  
      inputs : `PackedSequence`
          `bsize` sequences of shape `(len, input_dim)` each, in PackedSequence format
      states : `Tuple[torch.Tensor, torch.Tensor]`
          Tuple of tensors containing the initial hidden state and
          the cell state of each element in the batch. Each of these tensors have a dimension of
          (1 x bsize x nhid). Defaults to `None`.
  
      # Returns
  
      `Tuple[PackedSequence, Tuple[torch.Tensor, torch.Tensor]]`
          AugmentedLSTM representation of input and the state of the LSTM `t = seq_len`.
          Shape of representation is (bsize x seq_len x representation_dim).
          Shape of each state is (1 x bsize x nhid).
  
      """
      if not isinstance(inputs, PackedSequence):
          raise ConfigurationError(
              "inputs must be PackedSequence but got %s"
              % (type(inputs))
          )
  
      sequence_tensor, batch_lengths = pad_packed_sequence(
          inputs, batch_first=True
      )
      batch_size = sequence_tensor.size()[0]
      total_timesteps = sequence_tensor.size()[1]
      output_accumulator = sequence_tensor.new_zeros(
          batch_size, total_timesteps, self.lstm_dim
      )
      if states is None:
          full_batch_previous_memory = sequence_tensor.new_zeros(
              batch_size, self.lstm_dim
          )
          full_batch_previous_state = sequence_tensor.data.new_zeros(
              batch_size, self.lstm_dim
          )
      else:
          full_batch_previous_state = states[0].squeeze(0)
          full_batch_previous_memory = states[1].squeeze(0)
      current_length_index = batch_size - 1 if self.go_forward else 0
      if self.recurrent_dropout_probability > 0.0:
          dropout_mask = get_dropout_mask(
              self.recurrent_dropout_probability,
              full_batch_previous_memory,
          )
      else:
          dropout_mask = None
  
      for timestep in range(total_timesteps):
          index = (
              timestep
              if self.go_forward
              else total_timesteps - timestep - 1
          )
  
          if self.go_forward:
              while batch_lengths[current_length_index] <= index:
                  current_length_index -= 1
          # If we're going backwards, we are _picking up_ more indices.
          else:
              # First conditional: Are we already at the maximum
              # number of elements in the batch?
              # Second conditional: Does the next shortest
              # sequence beyond the current batch
              # index require computation use this timestep?
              while (
                  current_length_index < (len(batch_lengths) - 1)
                  and batch_lengths[current_length_index + 1] > index
              ):
                  current_length_index += 1
  
          previous_memory = full_batch_previous_memory[
              0 : current_length_index + 1
          ].clone()
          previous_state = full_batch_previous_state[
              0 : current_length_index + 1
          ].clone()
          timestep_input = sequence_tensor[
              0 : current_length_index + 1, index
          ]
          timestep_output, memory = self.cell(
              timestep_input,
              (previous_state, previous_memory),
              dropout_mask[0 : current_length_index + 1]
              if dropout_mask is not None
              else None,
          )
          full_batch_previous_memory = (
              full_batch_previous_memory.data.clone()
          )
          full_batch_previous_state = (
              full_batch_previous_state.data.clone()
          )
          full_batch_previous_memory[
              0 : current_length_index + 1
          ] = memory
          full_batch_previous_state[
              0 : current_length_index + 1
          ] = timestep_output
          output_accumulator[
              0 : current_length_index + 1, index, :
          ] = timestep_output
  
      output_accumulator = pack_padded_sequence(
          output_accumulator, batch_lengths, batch_first=True
      )
  
      # Mimic the pytorch API by returning state in the following shape:
      # (num_layers * num_directions, batch_size, lstm_dim). As this
      # LSTM cannot be stacked, the first dimension here is just 1.
      final_state = (
          full_batch_previous_state.unsqueeze(0),
          full_batch_previous_memory.unsqueeze(0),
      )
      return output_accumulator, final_state
  
imports: from typing import Optional, Tuple
  import torch
  from allennlp.common.checks import ConfigurationError
  from torch.nn.utils.rnn import PackedSequence, pack_padded_sequence, pad_packed_sequence
  from allennlp.nn.initializers import block_orthogonal
  from allennlp.nn.util import get_dropout_mask
label_desc: Use image node_id 2 for calling the AugmentedLstm obj's underlying member method code with example usage: obj.forward(inputs, states) and returns: output_accumulator,  final_state