class_docstr: Vanilla self-attention
class_docstr_tok: ["Vanilla","self-attention"]
code: def get_attention_map(self, x, return_map=False):
      B, N, C = x.shape
      qkv = (
          self.qkv(x)
          .reshape(B, N, 3, self.num_heads, C // self.num_heads)
          .permute(2, 0, 3, 1, 4)
      )
      q, k = qkv[0], qkv[1]
      attn_map = (q @ k.transpose(-2, -1)) * self.scale
      attn_map = attn_map.softmax(dim=-1).mean(0)
  
      img_size = int(N**0.5)
      ind = torch.arange(img_size).view(1, -1) - torch.arange(
          img_size
      ).view(-1, 1)
      indx = ind.repeat(img_size, img_size)
      indy = ind.repeat_interleave(img_size, dim=0).repeat_interleave(
          img_size, dim=1
      )
      indd = indx**2 + indy**2
      distances = indd**0.5
      distances = distances.to("cuda")
  
      dist = torch.einsum("nm,hnm->h", (distances, attn_map)).mean()
      dist = dist.item() / N
      if return_map:
          return dist, attn_map
      else:
          return dist
  
imports: import copy
  import logging
  import math
  import torch
  import torch.nn
  from vissl.models.model_helpers import DropPath, to_2tuple, trunc_normal_
  from vissl.models.trunks import register_model_trunk
  from vissl.utils.hydra_config import AttrDict
label_desc: Use image node_id 3 for calling the SelfAttention obj's underlying member method code with example usage: obj.get_attention_map(x, return_map) and returns: dist,  attn_map, dist