code: def forward(
      self,
      key_pos: torch.Tensor,
      query_pos: torch.Tensor,
      key_segment: torch.Tensor,
      query_segment: torch.Tensor,
  ):
      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 keylen == key_segment.size(
              1
          ) and querylen == query_segment.size(1)
  
          key_pos = key_pos.view(batch, -1, keylen)
          query_pos = query_pos.view(batch, querylen, -1)
          key_segment = key_segment.view(batch, -1, keylen)
          query_segment = query_segment.view(batch, querylen, -1)
  
          relative_position_bucket = (
              self._segment_relative_position_bucket(
                  query_segment, key_segment
              )
          )
          relative_position_bucket = (
              relative_position_bucket + self.num_buckets
          )  # 与相对位置编码区间不重叠
  
          # b*q*k
          absolute_position_bucket = self._position_bucket(
              torch.arange(
                  keylen,
                  dtype=torch.int32,
                  device=relative_position_bucket.device,
              )[None, :]
              - torch.arange(
                  querylen,
                  dtype=torch.int32,
                  device=relative_position_bucket.device,
              )[:, None],
              bidirectional=self.bidirectional,
              num_buckets=self.num_buckets,
              max_distance=self.max_distance,
          )
          relative_position_bucket = torch.where(
              (key_segment == query_segment),
              absolute_position_bucket[None, :, :],
              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
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/layers/position_embedding.py
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 SegmentPositionEmbedding obj's underlying member method code with example usage: obj.forward(key_pos, query_pos, key_segment, query_segment) and returns: embeds