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","."]
docstr: Warning: DO NOT USE THIS LAYER DIRECTLY, INSTEAD USE the AugmentedLSTM class
  
  Args:
      x (torch.Tensor): Input tensor of shape
          (bsize x input_dim).
      states (Tuple[torch.Tensor, torch.Tensor]): Tuple of tensors containing
          the hidden state and the cell state of each element in
          the batch. Each of these tensors have a dimension of
          (bsize x nhid). Defaults to `None`.
  
  Returns:
      Tuple[torch.Tensor, torch.Tensor]:
          Returned states. Shape of each state is (bsize x nhid).
docstr_tok: ["Warning",":","DO","NOT","USE","THIS","LAYER","DIRECTLY",",","INSTEAD","USE","the","AugmentedLSTM","class","Args",":","x","(","torch.Tensor",")",":","Input","tensor","of","shape","(","bsize","x","input_dim",")",".","states","(","Tuple","[","torch.Tensor",",","torch.Tensor","]",")",":","Tuple","of","tensors","containing","the","hidden","state","and","the","cell","state","of","each","element","in","the","batch",".","Each","of","these","tensors","have","a","dimension","of","(","bsize","x","nhid",")",".","Defaults","to","`","None","`",".","Returns",":","Tuple","[","torch.Tensor",",","torch.Tensor","]",":","Returned","states",".","Shape","of","each","state","is","(","bsize","x","nhid",")","."]
code: def forward(
      self,
      x: torch.Tensor,
      states=Tuple[torch.Tensor, torch.Tensor],
      variational_dropout_mask: Optional[torch.Tensor] = None,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Warning: DO NOT USE THIS LAYER DIRECTLY, INSTEAD USE the AugmentedLSTM class
  
      Args:
          x (torch.Tensor): Input tensor of shape
              (bsize x input_dim).
          states (Tuple[torch.Tensor, torch.Tensor]): Tuple of tensors containing
              the hidden state and the cell state of each element in
              the batch. Each of these tensors have a dimension of
              (bsize x nhid). Defaults to `None`.
  
      Returns:
          Tuple[torch.Tensor, torch.Tensor]:
              Returned states. Shape of each state is (bsize x nhid).
  
      """
      hidden_state, memory_state = states
      projected_input = self.input_linearity(x)
      projected_state = self.state_linearity(hidden_state)
  
      input_gate = (
          forget_gate
      ) = memory_init = output_gate = highway_gate = None
      if self.use_highway:
          fused_op = (
              projected_input[:, : 5 * self.lstm_dim] + projected_state
          )
          fused_chunked = torch.chunk(fused_op, 5, 1)
          (
              input_gate,
              forget_gate,
              memory_init,
              output_gate,
              highway_gate,
          ) = fused_chunked
          highway_gate = torch.sigmoid(highway_gate)
      else:
          fused_op = projected_input + projected_state
          (
              input_gate,
              forget_gate,
              memory_init,
              output_gate,
          ) = torch.chunk(fused_op, 4, 1)
      input_gate = torch.sigmoid(input_gate)
      forget_gate = torch.sigmoid(forget_gate)
      memory_init = torch.tanh(memory_init)
      output_gate = torch.sigmoid(output_gate)
      memory = input_gate * memory_init + forget_gate * memory_state
      timestep_output: torch.Tensor = output_gate * torch.tanh(memory)
  
      if self.use_highway:
          highway_input_projection = projected_input[
              :,
              self._highway_inp_proj_start : self._highway_inp_proj_end,
          ]
          timestep_output = (
              highway_gate * timestep_output
              + (1 - highway_gate) * highway_input_projection  # noqa
          )
      if variational_dropout_mask is not None and self.training:
          timestep_output = timestep_output * variational_dropout_mask
      return timestep_output, memory
  
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 AugmentedLSTMCell obj's underlying member method code with example usage: obj.forward(x, states, variational_dropout_mask) and returns: timestep_output,  memory