class_docstr: This class is designed to support the `relative position <https://arxiv.org/pdf/1803.02155v2.pdf>`__ in attention.
  
  Different from the absolute position embedding,
  the relative position embedding encodes relative distance between input tokens and learn the pairwise relations of them.
  It is commonly calculated via a look-up table with learnable parameters,
  interacting with queries and keys in self-attention modules.
  
  This class is different from PyTorch's built-in ``nn.MultiheadAttention`` in:
  
  1. It supports relative position embedding.
  2. It only supports self attention.
  3. It uses fixed dimension for each head, rather than fixed total dimension.
docstr: Count the FLOPs of :class:`RelativePositionSelfAttention`.
  
  RPE module is ignored in this computation.
docstr_tok: ["Count","the","FLOPs","of",":","class",":","`","RelativePositionSelfAttention","`",".","RPE","module","is","ignored","in","this","computation","."]
code: def _count_flops(
      self, x: tuple[MutableShape], y: tuple[MutableShape]
  ) -> FlopsResult:
      """Count the FLOPs of :class:`RelativePositionSelfAttention`.
  
      RPE module is ignored in this computation.
      """
      _, N, __ = x[0]
  
      # Dimension working inside.
      interm_dim = self.head_dim * self.num_heads
  
      params = (
          3 * self.embed_dim * (interm_dim + self.qkv_bias)
          +  # in_proj
          # skip RPE
          interm_dim
          * (self.embed_dim + 1)  # out_proj, bias always true
      )
  
      flops = (
          N * interm_dim * self.embed_dim * 3
          + N * N * interm_dim  # in_proj
          + N * interm_dim * N  # QK^T
          + N * N * interm_dim  # RPE (k)
          + N * interm_dim * N  # AV
          + N * interm_dim * self.embed_dim  # RPE (v)  # out_proj
      )
  
      return FlopsResult(flops, params)
  
imports: from __future__ import annotations
  from typing import Tuple, cast, Any, Dict
  import torch
  import torch.nn
  import torch.nn.functional
  import nni
  from nni.mutable import Categorical, MutableExpression, frozen_context, ensure_frozen, Mutable
  from nni.nas.nn.pytorch import ParametrizedModule, ModelSpace, MutableModule, Repeat, MutableConv2d, MutableLinear, MutableLayerNorm
  from nni.nas.space import current_model
  from nni.nas.oneshot.pytorch.supermodule.operation import MixedOperation
  from nni.nas.profiler.pytorch.flops import FlopsResult
  from nni.nas.profiler.pytorch.utils import MutableShape, ShapeTensor, profiler_leaf_module
  from .utils.pretrained import load_pretrained_weight
  from .utils.nn import DropPath
label_desc: Use image node_id 5 for calling the RelativePositionSelfAttention obj's underlying member method code with example usage: obj._count_flops(x, y) and returns: FlopsResult