docstr: Initialize the states required for incremental decoding
  
  Parameters
  ----------
  batch_size
      The batch size
  device
      The device
  dtype
      The data type of the states
  
  Returns
  -------
  states
      A list of states, each includes:
          - init_key
              - layout = 'NT'
                  Shape (batch_size, 0, N, C_key)
              - layout = 'TN'
                  Shape (0, batch_size, N, C_key)
          - init_value :
              - layout = 'NT'
                  Shape (batch_size, 0, N, C_value)
              - layout = 'TN'
                  Shape (0, batch_size, N, C_value)
docstr_tok: ["Initialize","the","states","required","for","incremental","decoding","Parameters","--","--","--","--","--","batch_size","The","batch","size","device","The","device","dtype","The","data","type","of","the","states","Returns","--","--","--","-","states","A","list","of","states",",","each","includes",":","-","init_key","-","layout","=","'NT'","Shape","(","batch_size",",","0",",","N",",","C_key",")","-","layout","=","'TN'","Shape","(","0",",","batch_size",",","N",",","C_key",")","-","init_value",":","-","layout","=","'NT'","Shape","(","batch_size",",","0",",","N",",","C_value",")","-","layout","=","'TN'","Shape","(","0",",","batch_size",",","N",",","C_value",")"]
code: def init_states(self, batch_size, device=None, dtype="float32"):
      """Initialize the states required for incremental decoding
  
      Parameters
      ----------
      batch_size
          The batch size
      device
          The device
      dtype
          The data type of the states
  
      Returns
      -------
      states
          A list of states, each includes:
              - init_key
                  - layout = 'NT'
                      Shape (batch_size, 0, N, C_key)
                  - layout = 'TN'
                      Shape (0, batch_size, N, C_key)
              - init_value :
                  - layout = 'NT'
                      Shape (batch_size, 0, N, C_value)
                  - layout = 'TN'
                      Shape (0, batch_size, N, C_value)
      """
      states = []
      for i in range(self.num_layers):
          if self.recurrent:
              layer = self.layers[0]
          else:
              layer = self.layers[i]
          states.append(
              layer.init_states(
                  batch_size=batch_size, device=device, dtype=dtype
              )
          )
      return 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 6 for calling the TransformerDecoder obj's underlying member method code with example usage: obj.init_states(batch_size, device, dtype) and returns: states