docstr: Incrementally generate the output given the decoder input.
  
  Parameters
  ----------
  data
      Shape (batch_size, C_in)
  states
      The previous states, contain a list of
  
      1. layout = 'NT'
          - prev_multi_key
              Shape (batch_size, prev_seq_length, num_heads, C_key)
          - prev_multi_value
              Shape (batch_size, prev_seq_length, num_heads, C_value)
  
      2. layout = 'TN'
          - prev_multi_key
              Shape (prev_seq_length, batch_size, num_heads, C_key)
          - prev_multi_value
              Shape (prev_seq_length, batch_size, num_heads, C_value)
  mem
      The memory
  
      1. layout = 'NT'
          Shape (batch_size, mem_length, C_mem)
      2. layout = 'TN'
          Shape (mem_length, batch_size, C_mem)
  
  mem_valid_length
      Valid length of the memory
      Shape (batch_size,)
  
  Returns
  -------
  out
      Shape (batch_size, C_out)
  new_states
      The updated states, contain a list of
  
      1. layout = 'NT'
          - new_key
              Shape (batch_size, prev_seq_length + 1, num_heads, C_key)
          - new_value
              Shape (prev_seq_length + 1, batch_size, num_heads, C_value)
  
      2. layout = 'TN'
          - new_key
              Shape (prev_seq_length + 1, batch_size, num_heads, C_key)
          - new_value
              Shape (prev_seq_length + 1, batch_size, num_heads, C_value)
docstr_tok: ["Incrementally","generate","the","output","given","the","decoder","input",".","Parameters","--","--","--","--","--","data","Shape","(","batch_size",",","C_in",")","states","The","previous","states",",","contain","a","list","of","1.","layout","=","'NT'","-","prev_multi_key","Shape","(","batch_size",",","prev_seq_length",",","num_heads",",","C_key",")","-","prev_multi_value","Shape","(","batch_size",",","prev_seq_length",",","num_heads",",","C_value",")","2.","layout","=","'TN'","-","prev_multi_key","Shape","(","prev_seq_length",",","batch_size",",","num_heads",",","C_key",")","-","prev_multi_value","Shape","(","prev_seq_length",",","batch_size",",","num_heads",",","C_value",")","mem","The","memory","1.","layout","=","'NT'","Shape","(","batch_size",",","mem_length",",","C_mem",")","2.","layout","=","'TN'","Shape","(","mem_length",",","batch_size",",","C_mem",")","mem_valid_length","Valid","length","of","the","memory","Shape","(","batch_size",",",")","Returns","--","--","--","-","out","Shape","(","batch_size",",","C_out",")","new_states","The","updated","states",",","contain","a","list","of","1.","layout","=","'NT'","-","new_key","Shape","(","batch_size",",","prev_seq_length","+","1",",","num_heads",",","C_key",")","-","new_value","Shape","(","prev_seq_length","+","1",",","batch_size",",","num_heads",",","C_value",")","2.","layout","=","'TN'","-","new_key","Shape","(","prev_seq_length","+","1",",","batch_size",",","num_heads",",","C_key",")","-","new_value","Shape","(","prev_seq_length","+","1",",","batch_size",",","num_heads",",","C_value",")"]
code: def incremental_decode(self, data, states, mem, mem_valid_length):
      """Incrementally generate the output given the decoder input.
  
      Parameters
      ----------
      data
          Shape (batch_size, C_in)
      states
          The previous states, contain a list of
  
          1. layout = 'NT'
              - prev_multi_key
                  Shape (batch_size, prev_seq_length, num_heads, C_key)
              - prev_multi_value
                  Shape (batch_size, prev_seq_length, num_heads, C_value)
  
          2. layout = 'TN'
              - prev_multi_key
                  Shape (prev_seq_length, batch_size, num_heads, C_key)
              - prev_multi_value
                  Shape (prev_seq_length, batch_size, num_heads, C_value)
      mem
          The memory
  
          1. layout = 'NT'
              Shape (batch_size, mem_length, C_mem)
          2. layout = 'TN'
              Shape (mem_length, batch_size, C_mem)
  
      mem_valid_length
          Valid length of the memory
          Shape (batch_size,)
  
      Returns
      -------
      out
          Shape (batch_size, C_out)
      new_states
          The updated states, contain a list of
  
          1. layout = 'NT'
              - new_key
                  Shape (batch_size, prev_seq_length + 1, num_heads, C_key)
              - new_value
                  Shape (prev_seq_length + 1, batch_size, num_heads, C_value)
  
          2. layout = 'TN'
              - new_key
                  Shape (prev_seq_length + 1, batch_size, num_heads, C_key)
              - new_value
                  Shape (prev_seq_length + 1, batch_size, num_heads, C_value)
  
      """
      # 1. Embed the data
      out = self.dropout_layer(data)
      if self._data_norm:
          out = self.ln_data(out)
      time_axis = 0 if self.layout == "TN" else 1
      # Generate the mem_attn_mask
      time_steps = npx.arange_like(mem, axis=time_axis)  # (mem_length,)
      mem_attn_mask = np.reshape(time_steps, (1, 1, -1)) < np.reshape(
          mem_valid_length, (-1, 1, 1)
      )
      # TODO(sxjscience) Try with boolean masking
      mem_attn_mask = mem_attn_mask.astype(self._dtype)
      new_states = []
      for i in range(self.num_layers):
          if self.recurrent:
              layer = self.layers[0]
          else:
              layer = self.layers[i]
          out, new_state = layer.incremental_decode(
              out, states[i], mem, mem_valid_length, mem_attn_mask
          )
          new_states.append(new_state)
      if self._pre_norm:
          out = self.ln_final(out)
      return out, 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 6 for calling the TransformerDecoder obj's underlying member method code with example usage: obj.incremental_decode(data, states, mem, mem_valid_length) and returns: out,  new_states