code: def transformer_gat_pgl(
      gw,
      feature,
      hidden_size,
      name,
      num_heads=4,
      attn_drop=0,
      edge_feature=None,
      concat=True,
      is_test=False,
  ):
      """transformer_gat_pgl"""
  
      def send_attention(src_feat, dst_feat, edge_feat):
          if edge_feat is None or not edge_feat:
              output = src_feat["k_h"] * dst_feat["q_h"]
              output = fluid.layers.reduce_sum(output, -1)
              return {
                  "alpha": output,
                  "v": src_feat["v_h"],
              }  # batch x h     batch x h x feat
          else:
              edge_feat = edge_feat["edge"]
              edge_feat = fluid.layers.reshape(
                  edge_feat, [-1, num_heads, hidden_size]
              )
              output = (src_feat["k_h"] + edge_feat) * dst_feat["q_h"]
              output = fluid.layers.reduce_sum(output, -1)
              return {
                  "alpha": output,
                  "v": (src_feat["v_h"] + edge_feat),
              }  # batch x h     batch x h x feat
  
      def reduce_attention(msg):
          alpha = msg[
              "alpha"
          ]  # lod-tensor (batch_size, seq_len, num_heads)
          h = msg["v"]
          alpha = paddle_helper.sequence_softmax(alpha)
          old_h = h
  
          if attn_drop > 1e-15:
              alpha = fluid.layers.dropout(
                  alpha,
                  dropout_prob=attn_drop,
                  is_test=is_test,
                  dropout_implementation="upscale_in_train",
              )
          h = h * alpha
          # h = fluid.layers.lod_reset(h, old_h)
          h = fluid.layers.sequence_pool(h, "sum")
          if concat:
              h = fluid.layers.reshape(h, [-1, num_heads * hidden_size])
          else:
              h = fluid.layers.reduce_mean(h, dim=1)
          return h
  
      q_w_attr = fluid.ParamAttr(
          initializer=fluid.initializer.XavierInitializer()
      )
      q_bias_attr = fluid.ParamAttr(
          initializer=fluid.initializer.ConstantInitializer(0.0)
      )
      q = fluid.layers.fc(
          feature,
          hidden_size * num_heads,
          name=name + "_q_weight",
          param_attr=q_w_attr,
          bias_attr=q_bias_attr,
      )
      q = q / (hidden_size**0.5)
  
      k_w_attr = fluid.ParamAttr(
          initializer=fluid.initializer.XavierInitializer()
      )
      k_bias_attr = fluid.ParamAttr(
          initializer=fluid.initializer.ConstantInitializer(0.0)
      )
      k = fluid.layers.fc(
          feature,
          hidden_size * num_heads,
          name=name + "_k_weight",
          param_attr=k_w_attr,
          bias_attr=k_bias_attr,
      )
  
      v_w_attr = fluid.ParamAttr(
          initializer=fluid.initializer.XavierInitializer()
      )
      v_bias_attr = fluid.ParamAttr(
          initializer=fluid.initializer.ConstantInitializer(0.0)
      )
      v = fluid.layers.fc(
          feature,
          hidden_size * num_heads,
          name=name + "_v_weight",
          param_attr=v_w_attr,
          bias_attr=v_bias_attr,
      )
  
      reshape_q = fluid.layers.reshape(q, [-1, num_heads, hidden_size])
      reshape_k = fluid.layers.reshape(k, [-1, num_heads, hidden_size])
      reshape_v = fluid.layers.reshape(v, [-1, num_heads, hidden_size])
  
      if not isinstance(gw, list):
          msg = gw.send(
              send_attention,
              nfeat_list=[
                  ("q_h", reshape_q),
                  ("k_h", reshape_k),
                  ("v_h", reshape_v),
              ],
              efeat_list=edge_feature,
          )
          output = gw.recv(msg, reduce_attention)
          return output
      else:
          checkpoints = []
          outputs = []
          for batch_no, (batch_gw, batch_edge_feat) in enumerate(
              zip(gw, edge_feature)
          ):
              msg = batch_gw.send(
                  send_attention,
                  nfeat_list=[
                      ("q_h", reshape_q),
                      ("k_h", reshape_k),
                      ("v_h", reshape_v),
                  ],
                  efeat_list=batch_edge_feat,
              )
              output = batch_gw.recv(msg, reduce_attention)
              outputs.append(output)
          outputs = L.sum(outputs)
          return outputs, checkpoints
  
file_path: PGL/ogb_examples/nodeproppred/ogbn-proteins/unimp_cross_feat/module/transformer_gat_pgl.py
imports: import paddle.fluid
  import paddle.fluid.layers
  from pgl import graph_wrapper
  from pgl.utils import paddle_helper
  import math
label_desc: Use image node_id 3 for calling a global function with example usage: transformer_gat_pgl(gw, feature, hidden_size, name, num_heads, attn_drop, edge_feature, concat, is_test) and returns: h, output, outputs,  checkpoints, dict, dict