KHopGraphΒΆ

class dgl.transforms.KHopGraph(k)[source]ΒΆ

Bases: dgl.transforms.module.BaseTransform

Return the graph whose edges connect the \(k\)-hop neighbors of the original graph.

This module only works for homogeneous graphs.

Parameters

k (int) – The number of hops.

Example

>>> import dgl
>>> from dgl import KHopGraph

>>> transform = KHopGraph(2)
>>> g = dgl.graph(([0, 1], [1, 2]))
>>> new_g = transform(g)
>>> print(new_g.edges())
(tensor([0]), tensor([2]))