code: def select_edges(src, dst, input_mask, num_nodes, max_seqlen):
      src = fluid.layers.elementwise_max(src, num_nodes * 0)
      dst = fluid.layers.elementwise_max(dst, num_nodes * 0)
      src = fluid.layers.elementwise_min(src, num_nodes - 1)
      dst = fluid.layers.elementwise_min(dst, num_nodes - 1)
  
      conditions = []
      conditions.append(L.gather(input_mask, src) > 0.5)
      conditions.append(L.gather(input_mask, dst) > 0.5)
      block_src = src / max_seqlen
      block_dst = dst / max_seqlen
      conditions.append(block_src == block_dst)
      mask = None
      for cond in conditions:
          if mask is None:
              mask = cond
          else:
              mask = L.logical_and(mask, cond)
  
      dst = masked_select(dst, mask)
      src = masked_select(src, mask)
      return src, dst
  
file_path: PGL/legacy/examples/xformer/sparse_scaled_dot_product_attention.py
imports: import paddle.fluid
  import paddle.fluid.layers
  import paddle.fluid.layers
  from pgl.utils import paddle_helper
  import pgl
label_desc: Use image node_id 2 for calling a global function with example usage: select_edges(src, dst, input_mask, num_nodes, max_seqlen) and returns: src,  dst