code: def convert_t5x_to_pytorch(
      variables: dict, *, num_layers: int, moe_interval: int
  ):
      """Converts the parameters from T5X-Flax to Transformers-PyTorch."""
      old = traverse_util.flatten_dict(variables["target"])
      old = {"/".join(k): v for k, v in old.items()}
  
      # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
      split_mlp_wi = True
      print("Split MLP:", split_mlp_wi)
  
      new = collections.OrderedDict()
      print(old.keys())
      for key, value in old.items():
          print(f"{key}: {value.shape}")
  
      # Shared embeddings.
      new["model.embed_tokens.weight"] = old["token_embedder/embedding"]
  
      # Decoder.
      for i in range(num_layers):
          # Block i, layer 0 (Self Attention).
          layer_norm = t5x_layer_norm_lookup(
              old, i, "decoder", "pre_self_attention_layer_norm"
          )
          k, o, q, v = t5x_attention_lookup(
              old, i, "decoder", "self_attention"
          )
          new[f"model.layers.{i}.input_layernorm.weight"] = layer_norm
          new[f"model.layers.{i}.self_attn.k_proj.weight"] = k.T
          new[f"model.layers.{i}.self_attn.o_proj.weight"] = o.T
          new[f"model.layers.{i}.self_attn.q_proj.weight"] = q.T
          new[f"model.layers.{i}.self_attn.v_proj.weight"] = v.T
  
          # Block i, layer 2 (MLP).
          layer_norm = t5x_layer_norm_lookup(
              old, i, "decoder", "pre_mlp_layer_norm"
          )
          new[
              f"model.layers.{i}.post_attention_layernorm.weight"
          ] = layer_norm
  
          if (i + 1) % moe_interval == 0:
              # moe
              gate = t5x_gate_lookup(old, i, "decoder", split_mlp_wi)
              new[f"model.layers.{i}.mlp.gate_weight"] = gate.T
              wi, wo = t5x_experts_lookup(
                  old, i, "decoder", split_mlp_wi
              )
              new[f"model.layers.{i}.mlp.experts.wi_gate"] = wi[0]
              new[f"model.layers.{i}.mlp.experts.wi_up"] = wi[1]
              new[f"model.layers.{i}.mlp.experts.wo"] = wo
              # extra
              layer_norm = t5x_layer_norm_lookup(
                  old, i, "decoder", "pre_extra_mlp_layer_norm"
              )
              new[
                  f"model.layers.{i}.pre_extra_mlp_layernorm.weight"
              ] = layer_norm
              wi, wo = t5x_extra_mlp_lookup(
                  old, i, "decoder", split_mlp_wi
              )
              new[f"model.layers.{i}.extra_mlp.gate_proj.weight"] = wi[
                  0
              ].T
              new[f"model.layers.{i}.extra_mlp.up_proj.weight"] = wi[
                  1
              ].T
              new[f"model.layers.{i}.extra_mlp.down_proj.weight"] = wo.T
          else:
              wi, wo = t5x_mlp_lookup(old, i, "decoder", split_mlp_wi)
              new[f"model.layers.{i}.mlp.gate_proj.weight"] = wi[0].T
              new[f"model.layers.{i}.mlp.up_proj.weight"] = wi[1].T
              new[f"model.layers.{i}.mlp.down_proj.weight"] = wo.T
  
      new["model.norm.weight"] = old["decoder/decoder_norm/scale"]
  
      # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
      if "decoder/logits_dense/kernel" in old:
          new["lm_head.weight"] = old["decoder/logits_dense/kernel"].T
  
      return new
  
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 7 for calling a global function with example usage: convert_t5x_to_pytorch(variables) and returns: new