code: def forward(
      self,
      rep: torch.Tensor,
      seq_lengths: torch.Tensor,
      states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
  ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
      if states is not None:
          # Transpose states so hidden and context both have shape
          # (num layers, batch size, lstm dim)
          states = (
              states[0].transpose(0, 1).contiguous(),
              states[1].transpose(0, 1).contiguous(),
          )
      else:
          # state has shape (num layers, batch size, lstm dim)
          state = torch.zeros(
              self.config.num_layers,
              rep.size(0),
              self.config.lstm_dim,
              device=torch.cuda.current_device()
              if cuda.CUDA_ENABLED
              else None,
          )
  
          states = (state, state)
  
      # hidden_by_layer is a list of hidden states for each layer of the
      # network, and similarly for context_by_layer
      hidden_by_layer, context_by_layer = states
  
      # Collect the last hidden and context for each layer
      last_hidden_by_layer = []
      last_context_by_layer = []
      rep = rep.transpose(0, 1).contiguous()
  
      for lstm, hidden, context in zip(
          self.lstms, hidden_by_layer, context_by_layer
      ):
          state = (hidden, context)
  
          # We purposefully throw away new_state until we reach the top layer
          # since we only care about passing on the final hidden state
          rep, (last_hidden, last_context) = lstm(
              rep, state, seq_lengths
          )
          last_hidden_by_layer.append(last_hidden)
          last_context_by_layer.append(last_context)
  
      # Make rep have shape (batch size, num layers, hidden size)
      rep = rep.transpose(0, 1).contiguous()
  
      # Make last_hidden and last_context have shape
      # (batch size, num layers, hidden size)
      last_hidden = torch.stack(last_hidden_by_layer).transpose(0, 1)
      last_context = torch.stack(last_context_by_layer).transpose(0, 1)
      return rep, (last_hidden, last_context)
  
file_path: pytext/pytext/models/representations/ordered_neuron_lstm.py
imports: from typing import List, Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from pytext.config import ConfigBase
  from pytext.models.module import Module
  from pytext.utils import cuda
  from pytext.utils.usage import log_class_usage
  from .representation_base import RepresentationBase
label_desc: Use image node_id 2 for calling the OrderedNeuronLSTM obj's underlying member method code with example usage: obj.forward(rep, seq_lengths, states) and returns: rep,  last_hidden,  last_context