Decoder¶

class Decoder(*args, **kwargs)[source]¶

Bases: ABC, Module

Interface for all decoders.

A decoder takes as input latent coordinates and an encoded graph context, and produces either a new graph with predictions or a global output vector.

Equivariant Decoder¶

class EquivariantDecoder(*args, **kwargs)[source]¶

Bases: Decoder, ABC

Abstract base class for equivariant decoders that produce outputs per hyper-edges.

Equivariant decoders preserve the structure of the input graph and produce predictions for each hyper-edge in a permutation-equivariant manner.

class MLPEquivariantDecoder(*args, **kwargs)[source]¶

Equivariant decoder that applies class-specific MLPs over hyper-edge features and latent coordinates.

\[\forall c \in \mathcal{C}, \forall e \in \mathcal{E}^c_x, \hat{y}_e = \phi_\theta^c(x_e, h_e),\]

where \(\phi_\theta^c\) is a class specific MLP.

Parameters:

in_graph_structure – Input graph structure.
in_array_size – Size of the input coordinate arrays.
hidden_sizes – Hidden sizes of the MLPs \(\phi_\theta^c\).
activation – Activation of the MLP \(\phi_\theta^c\).
out_structure – Graph structure of the output.
use_bias – Whether to use bias in the MLPs \(\phi_\theta^c\).
kernel_init – Kernel initializer for the MLPs \(\phi_\theta^c\).
bias_init – Bias initializer for the MLPs \(\phi_\theta^c\).
final_activation – Activation of the final layer of the MLPs \(\phi_\theta^c\).
encoded_feature_size – None if the input data has not been encoded, otherwise the size of the encoded features.
dtype – Computation dtype of the decoding (e.g. jnp.bfloat16 for mixed precision): coordinates and features are cast to this dtype before being gathered and fed to the MLPs. MLP parameters are stored in float32 regardless, and the output is cast back to the coordinates dtype. None (default) keeps the computation in the input dtype, i.e. full float32.
seed – Seed for RNG streams for weight initialization.

Return type:

Any

class InvariantDecoder(*args, **kwargs)[source]¶

Bases: Decoder, ABC

Abstract base class for invariant decoders that produce global outputs.

Invariant decoders aggregate information from all addresses in a permutation-invariant manner to produce a single global output vector.

class SumInvariantDecoder(*args, **kwargs)[source]¶

Sum invariant decoder, that sums the information of all addresses.

\[\hat{y} = \phi_\theta \left( \sum_{a \in \mathcal{A}(x)} \psi_\theta(h_a)\right),\]

where \(\phi_\theta\) (outer) and \(\psi_\theta\) (inner) are both trainable MLPs.

Parameters:

Return type:

Any

class MeanInvariantDecoder(*args, **kwargs)[source]¶

Mean invariant decoder, that averages the information of all addresses.

\[\hat{y} = \phi_\theta \left( \frac{1}{\vert \mathcal{A}(x) \vert} \sum_{a \in \mathcal{A}(x)} \psi_\theta(h_a) \right),\]

where \(\phi_\theta\) (outer) and \(\psi_\theta\) (inner) are both trainable MLPs.

Parameters:

Return type:

Any