torch_geometric.nn.conv.WLConv

class WLConv[source]

Bases: Module

The Weisfeiler Lehman (WL) operator from the “A Reduction of a Graph to a Canonical Form and an Algebra Arising During this Reduction” paper.

WLConv iteratively refines node colorings according to:

\[\mathbf{x}^{\prime}_i = \textrm{hash} \left( \mathbf{x}_i, \{ \mathbf{x}_j \colon j \in \mathcal{N}(i) \} \right)\]
Shapes:

  • input: node coloring \((|\mathcal{V}|, F_{in})\) (one-hot encodings) or \((|\mathcal{V}|)\) (integer-based), edge indices \((2, |\mathcal{E}|)\)

  • output: node coloring \((|\mathcal{V}|)\) (integer-based)

forward(x: Tensor, edge_index: Union[Tensor, SparseTensor]) → Tensor[source]

Runs the forward pass of the module.

reset_parameters()[source]

Resets all learnable parameters of the module.

histogram(x: Tensor, batch: Optional[Tensor] = None, norm: bool = False) → Tensor[source]

Given a node coloring x, computes the color histograms of the respective graphs (separated by batch).