code: def test_mhattn():
      mhattn = nn.MultiheadAttention(6, 3)
      query = torch.randn(8, 1, 6)
      key = torch.randn(7, 1, 6)
      value = torch.randn(7, 1, 6)
      profiler = FlopsParamsProfiler(mhattn, (query, key, value))
      result = profiler.profile({})
      assert result.flops == 2148
      assert result.params == sum(
          p.numel() for p in mhattn.parameters()
      )
  
      mhattn = nn.MultiheadAttention(
          6, 3, batch_first=True, kdim=5, vdim=4, bias=True
      )
      query = torch.randn(1, 8, 6)
      key = torch.randn(1, 7, 5)
      value = torch.randn(1, 7, 4)
      profiler = FlopsParamsProfiler(mhattn, (query, key, value))
      result = profiler.profile({})
      assert result.flops == 2022
      assert (
          result.params == 147
      )  # torch created extra 3 parameters. Cannot check.
  
      mhattn = nn.MultiheadAttention(6, 3, kdim=5, vdim=4, bias=False)
      query = torch.randn(8, 1, 6)
      key = torch.randn(7, 1, 5)
      value = torch.randn(7, 1, 4)
      profiler = FlopsParamsProfiler(mhattn, (query, key, value))
      result = profiler.profile({})
      assert result.flops == 1842
      assert result.params == sum(
          p.numel() for p in mhattn.parameters()
      )
  
file_path: nni/test/ut/nas/profiler/test_flops.py
imports: import pytest
  import torch
  from torch import nn
  import nni
  import nni.nas.nn.pytorch
  from nni.mutable import Categorical
  from nni.nas.profiler.pytorch.flops import FlopsParamsProfiler, FlopsProfiler, NumParamsProfiler, FlopsResult, register_flops_formula
  from ut.nas.nn.models import MODELS
label_desc: Use image node_id 4 for calling a global function with example usage: test_mhattn() without return types