Multiple built-in functions in message

Hi!
Suppose, I have these UDF:

def message_func(edges: dgl.EdgeBatch) -> Dict:
    return {
        'a': edges.src['a'],
        'b': edges.src['b'],
    }

def reduce_func(nodes: dgl.NodeBatch) -> Dict:
    a_sum = torch.sum(nodes.mailbox['a'], 1)
    b_sum = torch.sum(nodes.mailbox['b'], 1)
    return {'a_sum': a_sum, 'b_sum': b_sum}

I know, that for one feature, I can use built-in functions: copy_u and sum, but can I do something similar for features or more? What if message or reducing functions different?

Hi,
You can do g.update_all([fn.copy_u('a', 'a'), fn.copy_u('b','b')], [fn.sum('a', 'a_sum'), fn.sum('b', 'b_sum')])

2 Likes