UDF EdgeBatch.edges() confuses original graph and compute graph when message passing

[lzhmarkk]

🐛 Bug

Edge-wise User-defined Function EdgeBatch.edges() raise error as it confuses original graph and compute graph.

It happens when we use UDF as the message function inside of pull(), push() and send_and_recv() as they all call DGLHeteroGraph._create_compute_graph(). It extracts a subgraph named compute graph from the original one in order to compute message passing.

However, in this function, it sets the edges' id _eid according to the original graph instead of the relabelled ones, making edge index _eid exceeds the total number of edges in the compute graph. Thus, when we call EdgeBatch.edges() in the UDF, trying to find the edges with _eid in the compute graph instead of the original one causes Invalid edge ID error.

To Reproduce

Run the code below. pull(), push() and send_and_recv() all raise this error.

import dgl, torch, dgl.function as fn
g=dgl.graph(([0,1,2,3],[3,2,1,4]))

def msg_fn(edges):
    src, dst, _ = edges.edges()
    return {'m': torch.zeros(len(src), 1)}

g.pull([1], msg_fn, fn.sum('m','h'))
g.push([1], msg_fn, fn.sum('m','h'))
g.send_and_recv([1], msg_fn, fn.sum('m','h'))

Expected behavior

Error dgl._ffi.base.DGLError: Invalid edge ID 1 is raised.

Environment

• DGL Version: 0.9.0
• Backend Library & Version: 1.12.1
• OS: Mac OS 13.0
• How you installed DGL: conda 4.13.0
• Python version: 3.10.4

[Rhett-Ying]

Thanks for reporting this. Do you have any specific reason to do like this? pull/push/send_and_recv are to be deprecated. they can be achieved via in_subgraph() and update_all(). Could you try with it?

[lzhmarkk]

Thanks!

I just want to update the embedding of nodes within the mini-batch, when handling temporal graphs. It seems in_subgraph() and update_all() are sufficient.

Would you mind updating deprecation in the docs later?

[github-actions[bot]]

This issue has been automatically marked as stale due to lack of activity. It will be closed if no further activity occurs. Thank you