code: def gen_relative_position_index(
      q_size: Tuple[int, int],
      k_size: Optional[Tuple[int, int]] = None,
      class_token: bool = False,
  ) -> torch.Tensor:
      # Adapted with significant modifications from Swin / BeiT codebases
      # get pair-wise relative position index for each token inside the window
      assert (
          k_size is None
      ), "Different q & k sizes not currently supported"  # FIXME
  
      coords = torch.stack(
          torch.meshgrid(
              [torch.arange(q_size[0]), torch.arange(q_size[1])]
          )
      ).flatten(
          1
      )  # 2, Wh, Ww
      relative_coords = (
          coords[:, :, None] - coords[:, None, :]
      )  # 2, Wh*Ww, Wh*Ww
      relative_coords = relative_coords.permute(
          1, 2, 0
      )  # Qh*Qw, Kh*Kw, 2
      relative_coords[:, :, 0] += q_size[0] - 1  # shift to start from 0
      relative_coords[:, :, 1] += q_size[1] - 1
      relative_coords[:, :, 0] *= 2 * q_size[1] - 1
      num_relative_distance = (2 * q_size[0] - 1) * (2 * q_size[1] - 1)
  
      # else:
      #     # FIXME different q vs k sizes is a WIP, need to better offset the two grids?
      #     q_coords = torch.stack(
      #         torch.meshgrid([
      #             torch.arange(q_size[0]),
      #             torch.arange(q_size[1])
      #         ])
      #     ).flatten(1)  # 2, Wh, Ww
      #     k_coords = torch.stack(
      #         torch.meshgrid([
      #             torch.arange(k_size[0]),
      #             torch.arange(k_size[1])
      #         ])
      #     ).flatten(1)
      #     relative_coords = q_coords[:, :, None] - k_coords[:, None, :]  # 2, Wh*Ww, Wh*Ww
      #     relative_coords = relative_coords.permute(1, 2, 0)  # Qh*Qw, Kh*Kw, 2
      #     relative_coords[:, :, 0] += max(q_size[0], k_size[0]) - 1  # shift to start from 0
      #     relative_coords[:, :, 1] += max(q_size[1], k_size[1]) - 1
      #     relative_coords[:, :, 0] *= k_size[1] + q_size[1] - 1
      #     relative_position_index = relative_coords.sum(-1)  # Qh*Qw, Kh*Kw
      #     num_relative_distance = (q_size[0] + k_size[0] - 1) * (q_size[1] + k_size[1] - 1) + 3
  
      relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
  
      if class_token:
          # handle cls to token & token 2 cls & cls to cls as per beit for rel pos bias
          # NOTE not intended or tested with MLP log-coords
          relative_position_index = F.pad(
              relative_position_index, [1, 0, 1, 0]
          )
          relative_position_index[0, 0:] = num_relative_distance
          relative_position_index[0:, 0] = num_relative_distance + 1
          relative_position_index[0, 0] = num_relative_distance + 2
  
      return relative_position_index.contiguous()
  
imports: import math
  import os
  from typing import Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from .interpolate import RegularGridInterpolator
  from .mlp import Mlp
  from .weight_init import trunc_normal_
label_desc: Use image node_id 1 for calling a global function with example usage: gen_relative_position_index(q_size, k_size, class_token) and returns: relative_position_index