docstr: Incrementally generate the output given the decoder input.
  
  Parameters
  ----------
  data
      Shape (batch_size, C_in)
  states
      The previous states, contains
      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,)
  mem_attn_mask
      The attention mask between data and the memory
      Has shape (batch_size, 1, mem_length)
  
  Returns
  -------
  out
      Shape (batch_size, C_out)
  updated_states
      - new_key
          Shape (batch_size, prev_seq_length + 1, num_heads, C_key)
      - new_value
          Shape (batch_size, prev_seq_length + 1, 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",",","contains","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",",",")","mem_attn_mask","The","attention","mask","between","data","and","the","memory","Has","shape","(","batch_size",",","1",",","mem_length",")","Returns","--","--","--","-","out","Shape","(","batch_size",",","C_out",")","updated_states","-","new_key","Shape","(","batch_size",",","prev_seq_length","+","1",",","num_heads",",","C_key",")","-","new_value","Shape","(","batch_size",",","prev_seq_length","+","1",",","num_heads",",","C_value",")"]
code: def incremental_decode(
      self, data, states, mem, mem_valid_length, mem_attn_mask=None
  ):
      """Incrementally generate the output given the decoder input.
  
      Parameters
      ----------
      data
          Shape (batch_size, C_in)
      states
          The previous states, contains
          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,)
      mem_attn_mask
          The attention mask between data and the memory
          Has shape (batch_size, 1, mem_length)
  
      Returns
      -------
      out
          Shape (batch_size, C_out)
      updated_states
          - new_key
              Shape (batch_size, prev_seq_length + 1, num_heads, C_key)
          - new_value
              Shape (batch_size, prev_seq_length + 1, num_heads, C_value)
      """
      batch_size = data.shape[0]
      if self.layout == "NT":
          time_axis = 1
      else:
          time_axis = 0
      data = data.unsqueeze(time_axis)
      residual = data
      if self._pre_norm:
          data = self.ln_in(data)
      # Shape (B, prev_L, #Head, C_K), (B, prev_L, #Head, C_V)
      #  or (prev_L, B, #Head, C_K), (prev_L, B, #Head, C_V)
      prev_key, prev_value = states
      if mem_attn_mask is None:
          mem_attn_mask = gen_mem_attn_mask(
              mem, mem_valid_length, data, None, layout=self.layout
          )
      # 1. Get the causal self-attention value, we need to attend to both the current data
      # and the previous stored key/values
      # Shape (B, 1, 3 * num_heads * C_key)
      #  or (1, B, 3 * num_heads * C_key)
      step_qkv = self.attn_in_qkv(data)
      step_query, step_key, step_value = th.split(
          step_qkv, self._units, dim=-1
      )
      step_query = th.reshape(
          step_query,
          shape=(
              step_query.shape[0],
              step_query.shape[1],
              self._num_heads,
              -1,
          ),
      )
      step_key = th.reshape(
          step_key,
          shape=(
              step_key.shape[0],
              step_key.shape[1],
              self._num_heads,
              -1,
          ),
      )
      step_value = th.reshape(
          step_value,
          shape=(
              step_value.shape[0],
              step_value.shape[1],
              self._num_heads,
              -1,
          ),
      )
      new_key = th.cat([prev_key, step_key], dim=time_axis)
      new_value = th.cat([prev_value, step_value], dim=time_axis)
      out, [_, attn_weight] = self.self_attention(
          step_query, new_key, new_value, None
      )
      out = self.proj_in(out)
      out = self.dropout_layer(out)
      out = out + residual
      if not self._pre_norm:
          out = self.ln_in(out)
      # 2. Attend to the contextual memory
      data = out
      residual = data
      if self._pre_norm:
          data = self.ln_inter(data)
      out, _ = self.inter_attention(
          th.reshape(
              self.attn_inter_q(data),
              shape=(data.shape[0], data.shape[1], self._num_heads, -1),
          ),
          th.reshape(
              self.attn_inter_k(mem),
              shape=(mem.shape[0], mem.shape[1], self._num_heads, -1),
          ),
          th.reshape(
              self.attn_inter_v(mem),
              shape=(mem.shape[0], mem.shape[1], self._num_heads, -1),
          ),
          mem_attn_mask,
      )
      out = self.proj_inter(out)
      out = self.dropout_layer(out)
      out = out + residual
      if not self._pre_norm:
          out = self.ln_inter(out)
      # 3. Encode the output via an FFN layer
      out = self.ffn(out)
      out = th.reshape(out, shape=(batch_size, -1))
      return out, (new_key, new_value)
  
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 8 for calling the TransformerDecoderLayer obj's underlying member method code with example usage: obj.incremental_decode(data, states, mem, mem_valid_length, mem_attn_mask) and returns: out,  new_key,  new_value