dgl.udf.EdgeBatch.data

property EdgeBatch.data

Return a view of the edge features for the edges in the batch.

Examples

The following example uses PyTorch backend.

>>> import dgl
>>> import torch

>>> # Instantiate a graph and set an edge feature 'h'.
>>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
>>> g.edata['h'] = torch.tensor([[1.], [1.], [1.]])

>>> # Define a UDF that retrieves the feature 'h' for all edges.
>>> def edge_udf(edges):
>>>     # edges.data['h'] is a tensor of shape (E, 1),
>>>     # where E is the number of edges in the batch.
>>>     return {'data': edges.data['h']}

>>> # Make a copy of the feature with name 'data'.
>>> g.apply_edges(edge_udf)
>>> g.edata['data']
tensor([[1.],
        [1.],
        [1.]])

>>> # Use edge UDF in message passing, which is equivalent to
>>> # dgl.function.copy_e.
>>> import dgl.function as fn
>>> g.update_all(edge_udf, fn.sum('data', 'h'))
>>> g.ndata['h']
tensor([[1.],
        [2.]])