code: def _count_mhattn(
      module: nn.MultiheadAttention | nas_nn.MultiheadAttention,
      input: tuple[MutableShape, MutableShape, MutableShape],
      output: MutableShape,
      count_bias: bool = True,
  ) -> FlopsResult:
      # https://github.com/sovrasov/flops-counter.pytorch/blob/0d9518e2ca4cae9b2a0f43dd8c3297c30e33a565/ptflops/pytorch_ops.py#L159
  
      q, k, v = input[:3]  # there could be more than 3 arguments.
  
      qdim, kdim, vdim = q[-1], k[-1], v[-1]
      if not module.batch_first or len(q) == 2:
          # No batch dimension or batch dimension is the second dimension
          L, S = q[0], v[0]
      else:
          L, S = q[1], v[1]
  
      num_heads = _getattr(module, "num_heads")
      embed_dim = _getattr(module, "embed_dim")
      try:
          bias = _getattr(
              module, "bias", expected_type=bool
          )  # Mutable version
      except AttributeError:
          bias = _getattr(
              module, "in_proj_bias", expected_type=bool
          )  # nn version
      head_dim = embed_dim // num_heads
  
      params = (qdim + kdim + vdim) * (  # always count bias
          embed_dim + bias
      ) + embed_dim * (  # in_proj
          embed_dim + bias
      )  # out_proj
  
      flops = (
          L * qdim
          + L * embed_dim * (qdim + (bias and count_bias))  # Q scaling
          + S * embed_dim * (kdim + (bias and count_bias))  # in_proj Q
          + S * embed_dim * (vdim + (bias and count_bias))  # in_proj K
          + num_heads  # in_proj V
          * (  # attention heads: heads * (L * S * dim * 2 + 1)
              L * S * head_dim
              + L * S  # QK^T
              + L * S * head_dim  # softmax  # AV
          )
          + L
          * embed_dim
          * (embed_dim + (bias and count_bias))  # out_proj
      )
  
      return FlopsResult(flops, params)
  
imports: from __future__ import annotations
  import logging
  from dataclasses import dataclass, fields
  from functools import partial, reduce
  from typing import Any, NamedTuple, Callable, cast
  from torch import nn
  import nni.nas.nn.pytorch
  from nni.mutable import Sample, Mutable, MutableExpression
  from nni.nas.nn.pytorch import ModelSpace
  from nni.nas.profiler import Profiler, ExpressionProfiler
  from .utils import MutableShape, ShapeTensor, submodule_input_output_shapes, standardize_arguments, argument_in_spec, concat_name
  from .utils._attrs import tuple_n_t, _getattr
label_desc: Use image node_id 7 for calling a global function with example usage: _count_mhattn(module, input, output, count_bias) and returns: FlopsResult