code: def t5x_attention_lookup(params, i, prefix, layer_name="attention"):
      """Returns the KOQV parameters of (self-)attention. Does not transpose."""
      k = params[f"{prefix}/layers_{i}/{layer_name}/key/kernel"]
      o = params[f"{prefix}/layers_{i}/{layer_name}/out/kernel"]
      q = params[f"{prefix}/layers_{i}/{layer_name}/query/kernel"]
      v = params[f"{prefix}/layers_{i}/{layer_name}/value/kernel"]
      return k, o, q, v
  
file_path: colossalai/examples/language/openmoe/model/convert_openmoe_ckpt.py
imports: import argparse
  import collections
  import torch
  from flax import traverse_util
  from modeling_openmoe import OpenMoeForCausalLM
  from t5x import checkpoints
  from transformers import LlamaConfig
  from transformers.utils import logging
label_desc: Use image node_id 1 for calling a global function with example usage: t5x_attention_lookup(params, i, prefix, layer_name) and returns: k,  o,  q,  v