code: def map_state(self, fn):
      if self.state["src"] is not None:
          self.state["src"] = fn(self.state["src"], 0)
      for layer in self.transformer_layers:
          if hasattr(layer, "context_attn"):
              if layer.context_attn.layer_cache[1]["keys"].numel() != 0:
                  x = fn(layer.context_attn.layer_cache[1]["keys"], 0)
                  y = fn(layer.context_attn.layer_cache[1]["values"], 0)
                  layer.context_attn.layer_cache = True, {
                      "keys": x,
                      "values": y,
                  }
          if isinstance(layer.self_attn, AverageAttention):
              if layer.self_attn.layer_cache[1]["prev_g"].numel() != 0:
                  x = fn(layer.self_attn.layer_cache[1]["prev_g"], 0)
                  layer.self_attn.layer_cache = True, {"prev_g": x}
          else:
              if layer.self_attn.layer_cache[1]["keys"].numel() != 0:
                  x = fn(layer.self_attn.layer_cache[1]["keys"], 0)
                  y = fn(layer.self_attn.layer_cache[1]["values"], 0)
                  if (
                      layer.self_attn.layer_cache[1].get(
                          "key_pad_mask", None
                      )
                      is not None
                  ):
                      z = fn(
                          layer.self_attn.layer_cache[1][
                              "key_pad_mask"
                          ],
                          0,
                      )
                  else:
                      z = None
                  layer.self_attn.layer_cache = True, {
                      "keys": x,
                      "values": y,
                      "key_pad_mask": z,
                  }
  
imports: import torch
  import torch.nn
  from onmt.decoders.decoder import DecoderBase
  from onmt.modules import MultiHeadedAttention, AverageAttention
  from onmt.modules.position_ffn import PositionwiseFeedForward
  from onmt.modules.position_ffn import ActivationFunction
  from onmt.modules.moe import MoE
  from onmt.utils.misc import sequence_mask
  from onmt.modules.rmsnorm import RMSNorm
label_desc: Use image node_id 4 for calling the TransformerDecoderBase obj's underlying member method code with example usage: obj.map_state(fn) without return types