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.]])