docstr: Parameters
  ----------
  x
      - layout = 'NT'
          Shape (batch_size, seq_length, C_in)
      - layout = 'TN'
          Shape (seq_length, batch_size, C_in)
  
  layer_states
      - layout = 'NT'
          Shape (2, batch_size, prev_len, C_in)
      - layout = 'TN'
          Shape (2, prev_len, batch_size, C_in)
docstr_tok: ["Parameters","--","--","--","--","--","x","-","layout","=","'NT'","Shape","(","batch_size",",","seq_length",",","C_in",")","-","layout","=","'TN'","Shape","(","seq_length",",","batch_size",",","C_in",")","layer_states","-","layout","=","'NT'","Shape","(","2",",","batch_size",",","prev_len",",","C_in",")","-","layout","=","'TN'","Shape","(","2",",","prev_len",",","batch_size",",","C_in",")"]
code: def forward(self, x, layer_states):
      """
  
      Parameters
      ----------
      x
          - layout = 'NT'
              Shape (batch_size, seq_length, C_in)
          - layout = 'TN'
              Shape (seq_length, batch_size, C_in)
  
      layer_states
          - layout = 'NT'
              Shape (2, batch_size, prev_len, C_in)
          - layout = 'TN'
              Shape (2, prev_len, batch_size, C_in)
      """
      x = self.ln(x)
      if self._layout == "NT":
          batch_axis, time_axis = 0, 1
          prev_len = npx.shape_array(layer_states)[2]
      else:
          batch_axis, time_axis = 1, 0
          prev_len = npx.shape_array(layer_states)[1]
  
      query, key, value = np.split(self.qkv(x), 3, axis=-1)
      if layer_states is not None:
          prev_key, prev_value = layer_states[0], layer_states[1]
          key = np.concatenate([prev_key, key], axis=time_axis)
          value = np.concatenate([prev_value, value], axis=time_axis)
      new_states = np.stack([key, value], axis=0)
  
      # gen mask
      query_pos = npx.arange_like(query, axis=time_axis)
      if prev_len is not None:
          query_pos = query_pos + prev_len
      key_pos = npx.arange_like(key, axis=time_axis)
      # (query_len, key_len)
      mask = (
          npx.reshape(key_pos, (1, -1))
          <= npx.reshape(query_pos, (-1, 1))
      ).astype(self._dtype)
      # broadcast to (batch_size, query_len, key_len)
      mask = npx.broadcast_like(
          np.expand_dims(mask, axis=0),
          query,
          lhs_axes=0,
          rhs_axes=batch_axis,
      )
  
      query = npx.reshape(query, (-2, -2, self._num_heads, -1))
      key = npx.reshape(key, (-2, -2, self._num_heads, -1))
      value = npx.reshape(value, (-2, -2, self._num_heads, -1))
  
      out, [_, attn_weight] = self.attention_cell(
          query, key, value, mask
      )
      out = self.out_proj(out)
      out = self.hidden_dropout(out)
  
      return out, 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 2 for calling the GPT2SelfAttentionLayer obj's underlying member method code with example usage: obj.forward(x, layer_states) and returns: out,  new_states