code: def test_launch_quant_transform4d_0213():
      batch_size, seq_len = kt.bs_sl()
      hidden_dim = kt.hidden_dim
      nhead = kt.nhead
      head_dim = int(hidden_dim / nhead)
      trans_count = random.choice([1, 3])
      print(
          "(batch_size, seq_len, hidden_dim, nhead, trans_count): "
          f"({batch_size}, {seq_len}, {hidden_dim}, {nhead}, {trans_count})"
      )
  
      # shared weights
      vals = kt.rand(
          (trans_count, batch_size, nhead, seq_len, head_dim)
      )
      cmax = kt.topk(vals).to(kt.dtype)
  
      # custom weights
      custom_cmask = kt.randuint8(
          (batch_size, seq_len, trans_count, nhead, head_dim)
      )
      custom_res = kt.randint8(
          (batch_size, seq_len, trans_count, nhead, head_dim)
      )
  
      if kt.dtype == torch.float:
          func = cuda_module.torch_launch_quant_transform4d_0213_fp32
      else:
          func = cuda_module.torch_launch_quant_transform4d_0213_fp16
  
      # [trans_count, batch_size, nhead, seq_len, head_dim] ->
      # [batch_size, seq_len, trans_count, nhead, head_dim]
  
      def custom():
          custom_cmask.zero_()
          func(
              custom_res,
              custom_cmask,
              vals,
              cmax,
              batch_size,
              seq_len,
              hidden_dim,
              nhead,
              trans_count,
          )
          return [custom_res, custom_cmask]
  
      def baseline():
          base = vals.permute(1, 3, 0, 2, 4)
          base_q, base_cmask = kt.quantize(base, cmax)
          return [base_q.contiguous(), base_cmask.contiguous()]
  
      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 28 for calling a global function with example usage: test_launch_quant_transform4d_0213() and returns: custom,  baseline, list, list