SparseALIGNNConv

SparseALIGNNConv#

class nfflr.nn.SparseALIGNNConv(in_features: int, out_features: int, norm: Literal['layernorm', 'batchnorm'] = 'layernorm', skip_last_norm: bool = False)[source]#

ALIGNN with sparser line graph

Methods

forward(g, lg, x, y, z[, y_mask])

Node and Edge updates for ALIGNN layer.

forward(g: DGLGraph, lg: DGLGraph, x: Tensor, y: Tensor, z: Tensor, y_mask: Tensor | None = None)[source]#

Node and Edge updates for ALIGNN layer.

x: node input features y: edge input features z: edge pair input features