docstr: Parameters
  ----------
  step_data
      Shape (batch_size,)
  states
      It includes :
          - layout = 'NT'
              - mem_data : (batch_size, src_length, C_mem)
              - mem_valid_length : (batch_size,)
              - position : (batch_size,)
              - dec_states : list
  
          - layout = 'TN'
              - mem_data : (src_length, batch_size, C_mem)
              - mem_valid_length : (batch_size,)
              - position : (batch_size,)
              - dec_states : list
  
  Returns
  -------
  out
      Shape (batch_size, C)
  new_states
      Has the same structure as the states
docstr_tok: ["Parameters","--","--","--","--","--","step_data","Shape","(","batch_size",",",")","states","It","includes",":","-","layout","=","'NT'","-","mem_data",":","(","batch_size",",","src_length",",","C_mem",")","-","mem_valid_length",":","(","batch_size",",",")","-","position",":","(","batch_size",",",")","-","dec_states",":","list","-","layout","=","'TN'","-","mem_data",":","(","src_length",",","batch_size",",","C_mem",")","-","mem_valid_length",":","(","batch_size",",",")","-","position",":","(","batch_size",",",")","-","dec_states",":","list","Returns","--","--","--","-","out","Shape","(","batch_size",",","C",")","new_states","Has","the","same","structure","as","the","states"]
code: def forward(self, step_data, states):
      """
  
      Parameters
      ----------
      step_data
          Shape (batch_size,)
      states
          It includes :
              - layout = 'NT'
                  - mem_data : (batch_size, src_length, C_mem)
                  - mem_valid_length : (batch_size,)
                  - position : (batch_size,)
                  - dec_states : list
  
              - layout = 'TN'
                  - mem_data : (src_length, batch_size, C_mem)
                  - mem_valid_length : (batch_size,)
                  - position : (batch_size,)
                  - dec_states : list
  
      Returns
      -------
      out
          Shape (batch_size, C)
      new_states
          Has the same structure as the states
      """
      mem_data, mem_valid_length, position, dec_states = states
      # 1. Get the embedding
      step_data = self.model.tgt_embed_layer(step_data)
      if self.model.scaled_embed:
          step_data = step_data * _np.sqrt(self.model.dec_units)
      if self.model.pos_embed_type is not None:
          step_data = step_data + self.model.tgt_pos_embed_layer(
              position
          )
      out, new_states = self.model.decoder.incremental_decode(
          step_data, dec_states, mem_data, mem_valid_length
      )
      out = self.model.tgt_final_layer(out)
      return out, (mem_data, mem_valid_length, position + 1, new_states)
  
imports: import numpy
  import mxnet
  from mxnet import np, npx
  from mxnet import use_np
  from mxnet.gluon import nn, HybridBlock
  from typing import Optional, Tuple, List
  from ..utils.registry import Registry
  from ..attention_cell import MultiHeadAttentionCell, gen_self_attn_mask, gen_mem_attn_mask
  from ..layers import PositionalEmbedding, PositionwiseFFN, InitializerType
  from ..utils.config import CfgNode
  from ..sequence_sampler import BaseStepDecoder
label_desc: Use image node_id 5 for calling the TransformerInference obj's underlying member method code with example usage: obj.forward(step_data, states) and returns: out,  mem_data,  mem_valid_length,  unknown,  new_states