docstr: Parameters
  ----------
  x
      - layout = 'NT'
          Shape (batch_size, seq_length)
      - layout = 'TN'
          Shape (seq_length, batch_size)
  
  states
      The previous states
  
      - layout = 'NT'
          Shape (num_layers, 2, batch_size, prev_len, C_in)]
      - layout = 'TN'
          Shape (num_layers, 2, prev_len, batch_size, C_in)]
  
  Returns
  -------
  new_x
      Output
  
      - layout = 'NT'
          Shape (batch_size, seq_length, C_out)
      - layout = 'TN'
          Shape (seq_length, batch_size, C_out)
  
  new_states
      The new states
  
      - layout = 'NT'
          Shape (num_layers, 2, batch_size, prev_len + seq_length, C_in)
      - layout = 'TN'
          Shape (num_layers, 2, prev_len + seq_length, batch_size, C_in)
docstr_tok: ["Parameters","--","--","--","--","--","x","-","layout","=","'NT'","Shape","(","batch_size",",","seq_length",")","-","layout","=","'TN'","Shape","(","seq_length",",","batch_size",")","states","The","previous","states","-","layout","=","'NT'","Shape","(","num_layers",",","2",",","batch_size",",","prev_len",",","C_in",")","]","-","layout","=","'TN'","Shape","(","num_layers",",","2",",","prev_len",",","batch_size",",","C_in",")","]","Returns","--","--","--","-","new_x","Output","-","layout","=","'NT'","Shape","(","batch_size",",","seq_length",",","C_out",")","-","layout","=","'TN'","Shape","(","seq_length",",","batch_size",",","C_out",")","new_states","The","new","states","-","layout","=","'NT'","Shape","(","num_layers",",","2",",","batch_size",",","prev_len","+","seq_length",",","C_in",")","-","layout","=","'TN'","Shape","(","num_layers",",","2",",","prev_len","+","seq_length",",","batch_size",",","C_in",")"]
code: def forward(self, x, states):
      """
  
      Parameters
      ----------
      x
          - layout = 'NT'
              Shape (batch_size, seq_length)
          - layout = 'TN'
              Shape (seq_length, batch_size)
  
      states
          The previous states
  
          - layout = 'NT'
              Shape (num_layers, 2, batch_size, prev_len, C_in)]
          - layout = 'TN'
              Shape (num_layers, 2, prev_len, batch_size, C_in)]
  
      Returns
      -------
      new_x
          Output
  
          - layout = 'NT'
              Shape (batch_size, seq_length, C_out)
          - layout = 'TN'
              Shape (seq_length, batch_size, C_out)
  
      new_states
          The new states
  
          - layout = 'NT'
              Shape (num_layers, 2, batch_size, prev_len + seq_length, C_in)
          - layout = 'TN'
              Shape (num_layers, 2, prev_len + seq_length, batch_size, C_in)
  
      """
      prev_len = (
          npx.shape_array(states)[3]
          if self._layout == "NT"
          else npx.shape_array(states)[2]
      )
      x = self.get_initial_embedding(x, prev_len)
  
      if self._layout != self._compute_layout:
          x = np.swapaxes(x, 0, 1)
          states = np.swapaxes(states, 2, 3)
  
      new_states = []
      for layer_idx in range(self._num_layers):
          layer_states = None if states is None else states[layer_idx]
          x, new_layer_states = self._layers[layer_idx](x, layer_states)
          new_states.append(new_layer_states)
      new_states = np.stack(new_states, axis=0)
  
      x = self._final_ln(x)
      if self._layout != self._compute_layout:
          x = np.swapaxes(x, 0, 1)
          new_states = np.swapaxes(new_states, 2, 3)
      return x, new_states
  
imports: import os
  from typing import Tuple
  import mxnet
  from mxnet import use_np, np, npx
  from mxnet.gluon import HybridBlock, nn
  from .base import BACKBONE_REGISTRY
  from ..base import get_model_zoo_home_dir, get_repo_model_zoo_url, get_model_zoo_checksum_dir
  from ..utils.config import CfgNode
  from ..utils.misc import load_checksum_stats, download
  from ..utils.registry import Registry
  from ..initializer import TruncNorm
  from ..attention_cell import MultiHeadAttentionCell
  from ..layers import get_activation, PositionalEmbedding
  from ..data.tokenizers import HuggingFaceByteBPETokenizer
label_desc: Use image node_id 3 for calling the GPT2Model obj's underlying member method code with example usage: obj.forward(x, states) and returns: x,  new_states