class_docstr: `AugmentedLSTMCell` implements a AugmentedLSTM cell.
  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.
      use_highway (bool): If `True` we append a highway network to the
      outputs of the LSTM.
      use_bias (bool): If `True` we use a bias in our LSTM calculations, otherwise
      we don't.
  
  Attributes:
      input_linearity (nn.Module): Fused weight matrix which
          computes a linear function over the input.
      state_linearity (nn.Module): Fused weight matrix which
          computes a linear function over the states.
class_docstr_tok: ["`","AugmentedLSTMCell","`","implements","a","AugmentedLSTM","cell",".","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.","use_highway","(","bool",")",":","If","`","True","`","we","append","a","highway","network","to","the","outputs","of","the","LSTM",".","use_bias","(","bool",")",":","If","`","True","`","we","use","a","bias","in","our","LSTM","calculations",",","otherwise","we","do","n't",".","Attributes",":","input_linearity","(","nn.Module",")",":","Fused","weight","matrix","which","computes","a","linear","function","over","the","input",".","state_linearity","(","nn.Module",")",":","Fused","weight","matrix","which","computes","a","linear","function","over","the","states","."]
code: def __init__(
      self,
      embed_dim: int,
      lstm_dim: int,
      use_highway: bool,
      use_bias: bool = True,
  ):
      super().__init__()
      self.embed_dim = embed_dim
      self.lstm_dim = lstm_dim
      self.use_highway = use_highway
      self.use_bias = use_bias
  
      if use_highway:
          self._highway_inp_proj_start = 5 * self.lstm_dim
          self._highway_inp_proj_end = 6 * self.lstm_dim
  
          # fused linearity of input to input_gate,
          # forget_gate, memory_init, output_gate, highway_gate,
          # and the actual highway value
          self.input_linearity = nn.Linear(
              self.embed_dim,
              self._highway_inp_proj_end,
              bias=self.use_bias,
          )
          # fused linearity of input to input_gate,
          # forget_gate, memory_init, output_gate, highway_gate
          self.state_linearity = nn.Linear(
              self.lstm_dim, self._highway_inp_proj_start, bias=True
          )
      else:
          # If there's no highway layer then we have a standard
          # LSTM. The 4 comes from fusing input, forget, memory, output
          # gates/inputs.
          self.input_linearity = nn.Linear(
              self.embed_dim, 4 * self.lstm_dim, bias=self.use_bias
          )
          self.state_linearity = nn.Linear(
              self.lstm_dim, 4 * self.lstm_dim, bias=True
          )
      self.reset_parameters()
      log_class_usage(__class__)
  
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 1 to create a new AugmentedLSTMCell object from inherited base classes: nn with example: obj = AugmentedLSTMCell(embed_dim, lstm_dim, use_highway, use_bias)