code: def multihead_attention_formula(
      module: nn.MultiheadAttention | nas_nn.MutableMultiheadAttention,
      query: ShapeTensor,
      key: ShapeTensor,
      *args: Any,
      **kwargs,
  ) -> tuple[MutableShape, MutableShape | None]:
      shape = list(query.real_shape)  # type: ignore
      attn_shape = MutableShape(
          *shape[:-1], _getattr(module, "embed_dim")
      )
      key_shape = ensure_shape(key)
  
      weights_shape = None
      if kwargs.get("need_weights", True):
          batch_first = module.batch_first
          if module.batch_first is not _getattr(module, "batch_first"):
              _logger.warning(
                  "The batch_first attribute of the module is different from the batch_first "
                  "attribute of the formula. The shape inference result may be incorrect. "
                  "Assuming batch_first to be %s.",
                  batch_first,
              )
  
          if len(shape) == 2:
              # unbatched
              N, L = None, shape[0]
              S = key_shape[0]
          elif batch_first:
              N, L = shape[0], shape[1]
              S = key_shape[1]
          else:
              L, N = shape[0], shape[1]
              S = key_shape[0]
  
          if kwargs.get("average_attn_weights", True):
              if N is None:
                  weights_shape = MutableShape(L, S)
              else:
                  weights_shape = MutableShape(N, L, S)
          else:
              num_heads = _getattr(module, "num_heads")
              if N is None:
                  weights_shape = MutableShape(num_heads, L, S)
              else:
                  weights_shape = MutableShape(N, num_heads, L, S)
  
      return attn_shape, weights_shape
  
imports: from __future__ import annotations
  import logging
  import warnings
  from typing import Callable, Type, Tuple, Any, cast
  import torch
  from torch import nn
  import nni.nas.nn.pytorch
  from nni.mutable import MutableExpression
  from .shape import Formula, ShapeTensor, MutableShape, extract_shape_info, switch_case_shape_info, shape_inference
  from ._attrs import _getattr, tuple_2_t
label_desc: Use image node_id 11 for calling a global function with example usage: multihead_attention_formula(module, query, key) and returns: attn_shape,  weights_shape