dgl.udf.EdgeBatch.edges

EdgeBatch.edges()[source]

返回批次中的边。

返回:: (U, V, EID) – 批次中的边。对于每个 \(i\)，\((U[i], V[i])\) 是一条从 \(U[i]\) 到 \(V[i]\) 的边，其 ID 为 \(EID[i]\)。
返回类型:: (Tensor, Tensor, Tensor)

示例

以下示例使用 PyTorch 后端。

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