class_docstr: Flatten/unflatten the incremental state for use with TorchScripting.
  
  Typically, the incremental state will be stored as a Dict[int, Dict[str, Dict[str,
  torch.Tensor]]], where the 3 dictionary levels map decoder layer, attention type,
  and previous key/value/mask, respectively. However, TorchScript expects dicts to be
  of type Dict[str, torch.Tensor], and thus all input incremental states when
  TorchScripting will have to be of that type. We thus unflatten the input incremental
  state, already of type Dict[str, torch.Tensor], to pass it into whatever method
  needs it, and we flatten it again after the updated incremental state is passed back
  out.
  
  This is a base class that provides methods for flattening/unflattening: subclasses
  will call these methods as the incremental state is passed into and out of their own
  methods.
docstr: Flatten the input incremental state.
  
  For instance, structured_incr_state[0]['self_attn']['prev_key'] will be stored
  in flat_incr_state['layer_0__self_attn__prev_key'].
docstr_tok: ["Flatten","the","input","incremental","state",".","For","instance",",","structured_incr_state","[","0","]","[","'self_attn","'","]","[","'prev_key","'","]","will","be","stored","in","flat_incr_state","[","'layer_0__self_attn__prev_key","'","]","."]
code: def _flatten_incr_state(
      self,
      structured_incr_state: Dict[
          int, Dict[str, Dict[str, torch.Tensor]]
      ],
  ) -> Dict[str, torch.Tensor]:
      """
      Flatten the input incremental state.
  
      For instance, structured_incr_state[0]['self_attn']['prev_key'] will be stored
      in flat_incr_state['layer_0__self_attn__prev_key'].
      """
      flat_incr_state = {}
      for layer_idx, dict1 in structured_incr_state.items():
          for attn_type, dict2 in dict1.items():
              for state_type, state in dict2.items():
                  key = f"{layer_idx:d}__{attn_type}__{state_type}"
                  flat_incr_state[key] = state
      return flat_incr_state
  
imports: from collections import defaultdict
  from typing import Dict, List, Optional, Tuple
  import torch.jit
  from parlai.agents.bart.bart import BartAgent
  from parlai.core.dict import DictionaryAgent
  from parlai.core.torch_agent import TorchAgent
  from parlai.utils.bpe import Gpt2BpeHelper, SubwordBPEHelper
  from parlai.torchscript.tokenizer import ScriptableDictionaryAgent
  from torch import nn
label_desc: Use image node_id 3 for calling the BaseIncrStateFlattener obj's underlying member method code with example usage: obj._flatten_incr_state(structured_incr_state) and returns: flat_incr_state