code: def sparse_scaled_dot_product_attention(
      q, k, v, input_mask, dropout_rate, n_head, d_key, d_value
  ):
      def send_q_k_spmm(src_feat, dst_feat, edge_feat):
          # q [ num_edges, n_head * dim]
          # k [ num_edges, n_head * dim]
          # v [ num_edges, n_head * dim]
          _q = dst_feat["q"]
          _k = src_feat["k"]
          _v = src_feat["v"]
          _q = L.reshape(_q, [-1, n_head, _q.shape[-1] // n_head])
          _k = L.reshape(_k, [-1, n_head, _k.shape[-1] // n_head])
          score = L.reduce_sum(_q * _k, -1)  # [num_edge, n_head]
          return {"score": score, "value": _v}
  
      def recv_score_v_spmm(msg):
          score = msg["score"]
          score = paddle_helper.sequence_softmax(score)
          score = layers.dropout(
              score,
              dropout_prob=dropout_rate,
              dropout_implementation="upscale_in_train",
              is_test=False,
          )
  
          score = L.reshape(score, [-1, n_head, 1])
          _v = msg["value"]
          _new_v = L.reshape(_v, [-1, n_head, _v.shape[-1] // n_head])
  
          _new_v = _new_v * score
  
          _new_v = L.reshape(_new_v, [-1, _v.shape[-1]])
          _new_v = L.lod_reset(_new_v, _v)
          return L.sequence_pool(_new_v, "sum")
  
      graph_wrapper = BigBirdWrapper(input_mask)
      old_v = v
  
      q = L.reshape(q, [-1, d_key * n_head])
      k = L.reshape(k, [-1, d_key * n_head])
      v = L.reshape(v, [-1, d_value * n_head])
  
      q = L.scale(q, scale=d_key**-0.5)
      msg = graph_wrapper.send(
          send_q_k_spmm, nfeat_list=[("k", k), ("v", v), ("q", q)]
      )
      out = graph_wrapper.recv(msg, recv_score_v_spmm)
      out = L.reshape(out, [-1, L.shape(old_v)[1], d_value * n_head])
      return out, out
  
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 5 for calling a global function with example usage: sparse_scaled_dot_product_attention(q, k, v, input_mask, dropout_rate, n_head, d_key, d_value) and returns: out,  out, dict, L