code: def test_self_attention(
      batch_size, seq_length, units, num_heads, split
  ):
      net = MultiHeadAttention(
          units, num_heads, no_split_case=not split
      )
      in_data = mx.np.random.uniform(
          size=[batch_size, seq_length, units], dtype="float32"
      )
      if split:
          mask = mx.np.random.uniform(
              low=0,
              high=2,
              size=[batch_size, seq_length, seq_length],
              dtype="int32",
          )
      else:
          # key dimension will be expanded by num_heads value to simulate gpt-2 model
          # mask needs to be expanded as well
          mask = mx.np.random.uniform(
              low=0,
              high=2,
              size=[batch_size, seq_length, seq_length * 2],
              dtype="int32",
          )
  
      net.initialize()
      fused_net = net
      net.hybridize()
      ref_out = net(in_data, mask)
  
      fused_net.optimize_for(in_data, mask, backend="ONEDNN")
      out = fused_net(in_data, mask)
      mx.nd.waitall()
  
      assert_almost_equal(out.asnumpy(), ref_out.asnumpy())
  
      calib_data = mx.gluon.data.DataLoader(
          mx.gluon.data.ArrayDataset(in_data, mask),
          batch_size=batch_size,
      )
      qnet = mx.contrib.quant.quantize_net(
          net,
          quantized_dtype="auto",
          exclude_layers=None,
          exclude_layers_match=None,
          calib_data=calib_data,
          calib_mode="naive",
          num_calib_batches=batch_size,
          ctx=mx.cpu(),
      )
  
      qout = qnet(in_data, mask)
      mx.nd.waitall()
  
      min_range = np.min(ref_out.asnumpy())
      max_range = np.max(ref_out.asnumpy())
      atol = 0.1 * max(abs(min_range), abs(max_range))
      assert_almost_equal_with_err(
          qout.asnumpy(),
          ref_out.asnumpy(),
          rtol=0.1,
          atol=atol,
          etol=0.2,
      )
  
file_path: mxnet/tests/python/dnnl/subgraphs/test_matmul_subgraph.py
imports: import copy
  import mxnet
  import numpy
  import pytest
  from mxnet.contrib import quantization
  from mxnet.gluon import nn
  from mxnet.test_utils import assert_almost_equal, assert_almost_equal_with_err
  from mxnet.util import use_np
  import math
label_desc: Use image node_id 1 for calling a global function with example usage: test_self_attention(batch_size, seq_length, units, num_heads, split) without return types