code: def load_t5x_weights_in_t5(model, config, t5x_checkpoint_path):
      """Replaces the params in model witht the T5X converted params."""
      variables = checkpoints.load_t5x_checkpoint(t5x_checkpoint_path)
      converted = convert_t5x_to_pytorch(
          variables,
          num_layers=config.num_hidden_layers,
          moe_interval=config.moe_layer_interval,
      )
      state_dict = make_state_dict(converted)
      model.load_state_dict(state_dict, strict=True)
  
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 9 for calling a global function with example usage: load_t5x_weights_in_t5(model, config, t5x_checkpoint_path) without return types