code: def forward(
      self,
      query_pos: torch.Tensor,  # (batch, len_q)
      key_pos: torch.Tensor,  # (batch, len_k)
      rel_buckets: torch.Tensor,  # (batch, len_q, len_k)
  ):
      with torch.no_grad():
          batch = key_pos.size(0)
          keylen = key_pos.size(1)
          querylen = query_pos.size(1)
  
          assert key_pos.size(0) == query_pos.size(0)
          assert (
              rel_buckets.size(0) == batch
              and rel_buckets.size(1) == querylen
              and rel_buckets.size(2) == keylen
          )
  
          relative_position_bucket = (
              rel_buckets - 1 + self.num_buckets
          )  # 与相对位置编码区间不重叠
  
          # b*q*k
          inner_segment_bucket = self._position_bucket(
              key_pos[..., None, :] - query_pos[..., :, None],
              num_buckets=self.num_buckets,
              max_distance=self.max_distance,
          )
          relative_position_bucket = torch.where(
              rel_buckets == 0,
              inner_segment_bucket,
              relative_position_bucket,
          )
          # (batch, len_q, len_k)
  
      # (batch, len_q, len_k, num_heads)
      embeds = F.embedding(
          relative_position_bucket, self.relative_attention_bias
      )
      # (batch, num_heads, len_q, len_k)
      embeds = embeds.permute(0, 3, 1, 2).contiguous()
      return embeds
  
imports: import math
  from typing import Union
  import torch
  import bmtrain
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the BucketPositionBias obj's underlying member method code with example usage: obj.forward(query_pos, key_pos, rel_buckets) and returns: embeds