GraphNeuralKernel#

class GraphNeuralKernel(width, edge_features, n_layers=2, internal_n_layers=0, internal_layers=None, inner_size=None, internal_func=None, external_func=None, shared_weights=False)[source]#

Bases: Module

Graph Neural Operator kernel model class.

This class implements the Graph Neural Operator kernel network.

See also

Original reference: Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., Anandkumar, A. (2020). Neural Operator: Graph Kernel Network for Partial Differential Equations. DOI: arXiv preprint arXiv:2003.03485

Initialization of the GraphNeuralKernel class.

Parameters:

• width (int) – The width of the kernel.

• edge_features (int) – The number of edge features.

• n_layers (int) – The number of kernel layers. Default is 2.

• internal_n_layers (int) – The number of layers of the neural network inside each kernel layer. Default is 0.

• internal_layers (list[int] | tuple[int]) – The number of neurons for each layer of the neural network inside each kernel layer. Default is None.

• internal_func (torch.nn.Module) – The activation function used inside each kernel layer. If None, it uses the torch.nn.Tanh activation. Default is None.

• external_func (torch.nn.Module) – The activation function applied to the output of the each kernel layer. If None, it uses the torch.nn.Tanh activation. Default is None.

• shared_weights (bool) – If True, the weights of each kernel layer are shared. Default is False.

forward(x, edge_index, edge_attr)[source]#

The forward pass of the Graph Neural Kernel.

Parameters:

• x (torch.Tensor | LabelTensor) – The input tensor.

• edge_index (torch.Tensor) – The edge index.

• edge_attr (torch.Tensor | LabelTensor) – The edge attributes.

Returns:

The output tensor.

Return type:

torch.Tensor