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
      mem_length = mem.shape[time_axis]
      # Generate the mem_attn_mask
      time_steps = th.arange(
          mem_length, device=data.device
      )  # (mem_length,)
      mem_attn_mask = time_steps.reshape(
          (1, 1, -1)
      ) < mem_valid_length.reshape((-1, 1, 1))
  
      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: from typing import Union, Optional, List
  from torch import nn
  import torch
  from ..layers import PositionwiseFFN
  from ..attention_cell import MultiHeadAttentionCell, gen_mem_attn_mask, gen_self_attn_mask
  from ..utils import to_torch_dtype
label_desc: Use image node_id 7 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