Pytorch DGL:send（）和receive（）之间的可学习参数？

Question

我想知道是否有可能将参数作为发送消息的一部分进行学习（在减少参数之前）。比如说,

def gcn_message(edges):
    # Can we put a learnable function that we apply to edges.src['h'] here?
    return {'msg' : edges.src['h']}

class GCNLayer(nn.Module):
    def __init__(self, in_feats, out_feats):
        super(GCNLayer, self).__init__()

    def forward(self, g, inputs):
        g.ndata['h'] = inputs
        g.send(g.edges(), gcn_message)
        g.recv(g.nodes(), gcn_reduce)
        h = g.ndata.pop('h')
        return self.linear(h)

我认为有理由认为这是可能的，但如何才能实现呢