dgl.udf.EdgeBatch.edges

EdgeBatch.edges()[source]

Return the edges in the batch

Returns

(U, V, EID) – The edges in the batch. For each \(i\), \((U[i], V[i])\) is an edge from \(U[i]\) to \(V[i]\) with ID \(EID[i]\).

Return type

(Tensor, Tensor, Tensor)

Examples

The following example uses PyTorch backend.

>>> import dgl
>>> import torch
>>> # Instantiate a graph
>>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
>>> # Define a UDF that retrieves and concatenates the end nodes of the edges
>>> def edge_udf(edges):
>>>     src, dst, _ = edges.edges()
>>>     return {'uv': torch.stack([src, dst], dim=1).float()}
>>> # Create a feature 'uv' with the end nodes of the edges
>>> g.apply_edges(edge_udf)
>>> g.edata['uv']
tensor([[0., 1.],
        [1., 1.],
        [1., 0.]])
>>> # Use edge UDF in message passing
>>> import dgl.function as fn
>>> g.update_all(edge_udf, fn.sum('uv', 'h'))
>>> g.ndata['h']
tensor([[1., 0.],
        [1., 2.]])