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)

Я думаю, что разумно думать, что это возможно, но как это можно сделать?