code: def blockdiag_to_dense_mlp_bert(state_dict):
      from src.ops.blockdiag_multiply import (
          blockdiag_weight_to_dense_weight,
      )
  
      names = {
          name
          for name in state_dict
          if re.match(
              "bert.encoder.layer.(\d+).(mlp.fc(1|2)|(intermediate|output).dense).weight",
              name,
          )
      }
      for name in names:
          state_dict[name] = blockdiag_weight_to_dense_weight(
              state_dict[name]
          )
      return state_dict
  
file_path: flash-attention/training/src/utils/checkpoint.py
imports: import re
  from pathlib import Path
  import torch
  import math
  from einops import rearrange
label_desc: Use image node_id 2 for calling a global function with example usage: blockdiag_to_dense_mlp_bert(state_dict) and returns: state_dict