code: def test_amp_transformers():
      batch_size = 16
      seq_length = 32
      units = 8
      num_heads = 8
  
      in_data = mx.np.random.uniform(
          size=(batch_size, seq_length, units), dtype="float32"
      )
      mask = mx.np.random.randint(
          0, 2, (batch_size, seq_length, seq_length), dtype="int32"
      )
  
      net = MultiHeadAttention(units, num_heads)
      net.initialize()
  
      check_amp_fuse(net, [in_data, mask], None)
      check_amp_fuse(
          net, [in_data, mask], None, ["sg_onednn_fully_connected_0"]
      )
  
file_path: mxnet/tests/python/dnnl/subgraphs/test_amp_subgraph.py
imports: import json
  import mxnet
  from mxnet import amp
  from mxnet.gluon import nn
  from mxnet.test_utils import assert_almost_equal
  from subgraph_common import SG_PASS_NAME, QUANTIZE_SG_PASS_NAME
  from test_matmul_subgraph import MultiHeadAttention
  import os
  import sys
  from pathlib import Path
  from amp.common import check_amp_net_stats
label_desc: Use image node_id 5 for calling a global function with example usage: test_amp_transformers() without return types