ChebConv

torchmil.nn.gnns.ChebConv

Bases: Module

__init__(in_channels, out_channels, K=3, compute_lambda_max=True)

Chebyshev spectral graph convolutional operator tailored for sparse adjacency matrices (COO format). Proposed in Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering. This class follows closely the design of PyTorch Geometric's ChebConv, but adapted for COO format adjacency matrices.

Parameters:

• in_channels (int) –

  Number of input channels (features per node).

• out_channels (int) –

  Number of output channels (features per node).

• K (int, default: 3 ) –

  Order of the Chebyshev polynomial approximation. Default is 3.

• compute_lambda_max (bool, default: True ) –

  If True, computes the maximum eigenvalue of the adjacency matrix for normalization. If False and not provided during forward pass, it will be set it to 2.0.

forward(x, adj, lambda_max=None)

Forward pass of the Chebyshev convolutional layer.

Parameters:

• x (Tensor) –

  Input tensor of shape (batch_size, n_nodes, in_channels).

• adj (Tensor) –

  Adjacency matrix of shape (batch_size, n_nodes, n_nodes). It can be a sparse tensor (in COO format) or a dense tensor.

• lambda_max (Tensor, default: None ) –

  Optional maximum eigenvalue of the adjacency matrix for normalization. If not provided, it will be computed.

Returns:

• out ( Tensor ) –

  Output tensor of shape (batch_size, n_nodes, out_channels).