code: def relative_matmul(x: Tensor, z: Tensor, transpose: bool) -> Tensor:
      """
      Helper function for relative positions attention.
      https://arxiv.org/pdf/1803.02155.pdf
      x shape [batch_size x heads x q_len x k_len]
      """
      batch_size, heads, length, _ = x.size()
      x_t = x.permute(2, 0, 1, 3)
      x_t_r = x_t.contiguous().view(length, heads * batch_size, -1)
      if transpose:
          z = z.transpose(1, 2)
      x_tz_matmul = torch.matmul(x_t_r, z)
      x_tz_matmul_r = x_tz_matmul.view(length, batch_size, heads, -1)
      x_tz_matmul_r_t = x_tz_matmul_r.permute(1, 2, 0, 3)
      return x_tz_matmul_r_t
  
imports: import torch
  import torch.nn
  from math import log, sqrt
  from torch import Tensor
  from typing import Optional, Tuple
  from torch.nn.functional import scaled_dot_product_attention
  from torch.utils.checkpoint import checkpoint
  from torch.nn.utils import skip_init
  from .alibi_position_bias import AlibiPositionalBias
  from torch.distributed import all_reduce
  from importlib import import_module
label_desc: Use image node_id 4 for calling a global function with example usage: relative_matmul(x, z, transpose) and returns: x_tz_matmul_r_t