code: def rename_state_keys(state, keys_regex, replacement):
      """Rename keys from state that match a regex; replacement can use capture groups"""
      regex = re.compile(keys_regex)
      return {
          (k if not regex.findall(k) else regex.sub(replacement, k)): v
          for k, v in state.items()
      }
  
imports: import re
  from typing import List, Optional
  import torch
  from pytext.utils.usage import log_class_usage
  from pytorch.text.fb.nn.modules.transformer import Transformer
  from torch import nn
label_desc: Use image node_id 2 for calling a global function with example usage: rename_state_keys(state, keys_regex, replacement) and returns: unknown