API Reference

Module contents

class gfn.GFN(in_features, out_features, bias=True, device=None, dtype=None, nn_backend='scipy')[source]

Bases: Linear

The graph feedforward network (GFN) layer from “GFN: A graph feedforward network for resolution-invariant reduced operator learning in multifidelity applications”.

The layer is an extension of the standard torch.nn.Linear, but with weights and biases are optionally associated to original input \(\mathcal{M}^i_o\) and output \(\mathcal{M}^o_o\) graphs. GFN then defines new weights and biases for new input \(\mathcal{M}^i_n\) and output \(\mathcal{M}^o_n\) graphs, allowing for resolution-invariant learning.

The general GFN equation is given by:

\begin{equation*} \begin{aligned} \tilde{W}_{i_{{\mathcal{M}^o_{n}}}j_{{\mathcal{M}^i_{n}}}} &= \underset{\forall l_{\mathcal{M}^o_{o}} \text{ s.t } l_{\mathcal{M}^o_{o}} {\leftarrow}\!{\backslash}\!{\rightarrow} i_{\mathcal{M}^o_{n}}}{\operatorname{mean}} \sum_{\forall k_{\mathcal{M}^i_{o}} \text{ s.t } k_{\mathcal{M}^i_{o}} {\leftarrow}\!{\backslash}\!{\rightarrow} j_{\mathcal{M}^i_{n}}} \frac{W_{l_{\mathcal{M}^o_{o}}k_{\mathcal{M}^i_{o}}}}{\lvert \{ h_{\mathcal{M}^i_{n}} \text{ s.t. } k_{\mathcal{M}^i_{o}} ~{\leftarrow}\!{\backslash}\!{\rightarrow}~ h_{\mathcal{M}^i_{n}} \}\rvert}, \\ \tilde{b}^d_{i_{\mathcal{M}^o_{n}}} &= \underset{\forall k_{\mathcal{M}^o_{o}} \text{ s.t } k_{\mathcal{M}^o_{o}} {\leftarrow}\!{\backslash}\!{\rightarrow} i_{\mathcal{M}^o_{n}}}{\operatorname{mean}} {b}^d_{k_{\mathcal{M}^o_{o}}}. \end{aligned} \end{equation*}

where:

  • \(\mathcal{M}^i_{o}\) is the original input graph,

  • \(\mathcal{M}^o_{o}\) is the original output graph,

  • \(\mathcal{M}^i_{n}\) is the new input graph,

  • \(\mathcal{M}^o_{n}\) is the new output graph,

  • \(W\) and \(b\) are the weights and biases associated to the original graphs,

  • \(\tilde{W}\) and \(\tilde{b}\) are the new weights and biases associated to the new graphs,

  • \(i_{\mathcal{M}_1} {\leftarrow}\!{\backslash}\!{\rightarrow} j_{\mathcal{M}_2}\) indicates that either node \(i\) in graph \(\mathcal{M}_1\) is the nearest neighbor of node \(j\) in graph \(\mathcal{M}_2\) or vice versa.

GFN also supports mapping from graphs to vectors or vectors to graphs, in which case the above equation can be applied with \(\mathcal{M}^i_{o}=\mathcal{M}^i_{n}\) for a vector input or \(\mathcal{M}^o_{o}=\mathcal{M}^o_{n}\) for a vector output.

Parameters:

  • in_features (int or torch.Tensor) – Either a tensor of shape \((N_{\text{in}}, D_{\text{in}})\) containing the coordinates of each original input node (graph input) or the size \(N_{\text{in}}\) of each input (vector input).

  • out_features (int or torch.Tensor) – Either a tensor of shape \((N_{\text{out}}, D_{\text{out}})\) containing the coordinates of each original output node (graph output) or the size \(N_{\text{out}}\) of each output (vector output).

  • bias (bool, optional) – If set to False, the layer will not learn an additive bias. (default: True)

  • device (torch.device, optional) – The device on which the layer should be allocated. (default: None)

  • dtype (torch.dtype, optional) – The data type for the layer’s parameters. (default: None)

  • nn_backend (str, optional) – The backend to use for nearest neighbor lookup. Can be either "scipy" or "faiss". (default: "scipy")

forward(x, in_graph=None, out_graph=None)[source]

Runs the forward pass of the module.

Parameters:

  • x (torch.Tensor) – The input tensor of shape \((..., N_{\text{in}}^{\prime})\).

  • in_graph (torch.Tensor, optional) – The input graph matching the shape of the input tensor \((N_{\text{in}}^{\prime}, D_{\text{in}})\). If None, treats as vector input i.e. assumes no change from the original input graph. (default: None)

  • out_graph (torch.Tensor, optional) – The output graph of shape \((N_{\text{out}}^{\prime}, D_{\text{out}})\). If None, treats as vector output i.e. assumes no change to the original output graph. (default: None)

Returns:

The output tensor. Matches the shape of the new output graph if provided.

Return type:

torch.Tensor