class_docstr: `AugmentedLSTMUnidirectional` 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 controlls the level of variational dropout done.
  
  Args:
      embed_dim (int): The number of expected features in the input.
      lstm_dim (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: ["`","AugmentedLSTMUnidirectional","`","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","controlls","the","level","of","variational","dropout","done",".","Args",":","embed_dim","(","int",")",":","The","number","of","expected","features","in","the","input",".","lstm_dim","(","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: DO NOT USE THIS LAYER DIRECTLY, INSTEAD USE the AugmentedLSTM class
  
  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.
  
  Args:
      inputs (PackedSequence): Input tensor of shape
          (bsize x seq_len x input_dim).
      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 num_directions * nhid). Defaults to `None`.
  
  Returns:
      Tuple[PackedSequence, Tuple[torch.Tensor, torch.Tensor]]:
          AgumentedLSTM 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",":","DO","NOT","USE","THIS","LAYER","DIRECTLY",",","INSTEAD","USE","the","AugmentedLSTM","class","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",".","Args",":","inputs","(","PackedSequence",")",":","Input","tensor","of","shape","(","bsize","x","seq_len","x","input_dim",")",".","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","num_directions","*","nhid",")",".","Defaults","to","`","None","`",".","Returns",":","Tuple","[","PackedSequence",",","Tuple","[","torch.Tensor",",","torch.Tensor","]","]",":","AgumentedLSTM","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: DO NOT USE THIS LAYER DIRECTLY, INSTEAD USE the AugmentedLSTM class
  
      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.
  
      Args:
          inputs (PackedSequence): Input tensor of shape
              (bsize x seq_len x input_dim).
          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 num_directions * nhid). Defaults to `None`.
  
      Returns:
          Tuple[PackedSequence, Tuple[torch.Tensor, torch.Tensor]]:
              AgumentedLSTM 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).
  
      """
      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 = self.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: import math
  from typing import Optional, Tuple
  import torch
  import torch.nn
  from pytext.config import ConfigBase
  from pytext.utils.usage import log_class_usage
  from torch.nn.utils.rnn import pack_padded_sequence, PackedSequence, pad_packed_sequence
  from .representation_base import RepresentationBase
label_desc: Use image node_id 3 for calling the AugmentedLSTMUnidirectional obj's underlying member method code with example usage: obj.forward(inputs, states) and returns: output_accumulator,  final_state