code: def test_launch_transform_0213():
      batch_size, seq_len = kt.bs_sl()
      hidden_dim = kt.hidden_dim
      nhead = kt.nhead
      head_dim = int(hidden_dim / nhead)
      print(
          "(batch_size, seq_len, hidden_dim, nhead): "
          f"({batch_size}, {seq_len}, {hidden_dim}, {nhead})"
      )
  
      vals = kt.rand((batch_size, seq_len, hidden_dim))
      custom_res = kt.rand((batch_size, nhead, seq_len, head_dim))
  
      if kt.dtype == torch.float:
          func = cuda_module.torch_launch_transform_0213_fp32
      else:
          func = cuda_module.torch_launch_transform_0213_fp16
  
      # [sz0, sz1, sz2, sz3] -> [sz0, sz2, sz1, sz3]
      def custom():
          func(vals, custom_res, batch_size, seq_len, nhead, head_dim)
          return kt.norm_res_list(custom_res)
  
      def baseline():
          base = vals.reshape(
              (batch_size, seq_len, nhead, head_dim)
          ).transpose(1, 2)
          return kt.norm_res_list(base)
  
      return custom, baseline
  
imports: import random
  import torch
  from torch._C import ScriptModule, dtype
  from torch.nn import functional
  from lightseq.training.ops.pytorch.builder import KernelBuilder, AdamBuilder
  from tests.util import TestDecorator
  from torch_crf import CRF
label_desc: Use image node_id 3 for calling a global function with example usage: test_launch_transform_0213() and returns: custom,  baseline, kt, kt