dgl.udf.EdgeBatch.src

property EdgeBatch.src

Return a view of the source node features for the edges in the batch.

Examples

The following example uses PyTorch backend.

>>> import dgl
>>> import torch

>>> # Instantiate a graph and set a node feature 'h'.
>>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
>>> g.ndata['h'] = torch.ones(2, 1)

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

>>> # Copy features from source nodes to edges.
>>> g.apply_edges(edge_udf)
>>> g.edata['src']
tensor([[1.],
        [1.],
        [1.]])

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