Ready-to-use GNNs

We propose several ready-to-use implementations. One just needs to pass an input structure and an output structure to the constructor.

class TinyRecurrentEquivariantGNN(*args, **kwargs)[source]

Tiny ready-to-use equivariant GNN, with basic recurrent message passing.

  • Normalizer: TDigestNormalizer with 10 breakpoints.

  • Encoder: MLPEncoder with 0 hidden layer and output of size 4.

  • Coupler: RecurrentCoupler with 5 steps, and latent dimension 4.

  • Decoder: MLPEquivariantDecoder with 0 hidden layer.

Parameters:

  • in_structure (GraphStructure) – Structure of the input graph.

  • out_structure (GraphStructure) – Structure of the output graph.

  • seed (int) – Seed for RNG streams.

Return type:

Any

class SmallRecurrentEquivariantGNN(*args, **kwargs)[source]

Small ready-to-use equivariant GNN, with basic recurrent message passing.

  • Normalizer: TDigestNormalizer with 20 breakpoints.

  • Encoder: MLPEncoder with 1 hidden layer of size 16 and output of size 8.

  • Coupler: RecurrentCoupler with 10 steps, hidden layers of size 16 and latent dimension 8.

  • Decoder: MLPEquivariantDecoder with 1 hidden layer of size 16.

Parameters:

  • in_structure (GraphStructure) – Structure of the input graph.

  • out_structure (GraphStructure) – Structure of the output graph.

  • seed (int) – Seed for RNG streams.

Return type:

Any

class MediumRecurrentEquivariantGNN(*args, **kwargs)[source]

Medium ready-to-use equivariant GNN, with basic recurrent message passing.

  • Normalizer: TDigestNormalizer with 50 breakpoints.

  • Encoder: MLPEncoder with 1 hidden layer of size 32 and output of size 16.

  • Coupler: RecurrentCoupler with 20 steps, hidden layers of size 32 and latent dimension 16.

  • Decoder: MLPEquivariantDecoder with 1 hidden layer of size 32.

Parameters:

  • in_structure (GraphStructure) – Structure of the input graph.

  • out_structure (GraphStructure) – Structure of the output graph.

  • seed (int) – Seed for RNG streams.

Return type:

Any

class LargeRecurrentEquivariantGNN(*args, **kwargs)[source]

Large ready-to-use equivariant GNN, with basic recurrent message passing.

  • Normalizer: TDigestNormalizer with 100 breakpoints.

  • Encoder: MLPEncoder with 1 hidden layer of size 64 and output of size 32.

  • Coupler: RecurrentCoupler with 50 steps, hidden layers of size 64 and latent dimension 32.

  • Decoder: MLPEquivariantDecoder with 1 hidden layer of size 64.

Parameters:

  • in_structure (GraphStructure) – Structure of the input graph.

  • out_structure (GraphStructure) – Structure of the output graph.

  • seed (int) – Seed for RNG streams.

Return type:

Any

class ExtraLargeRecurrentEquivariantGNN(*args, **kwargs)[source]

Extra large ready-to-use equivariant GNN, with basic recurrent message passing.

  • Normalizer: TDigestNormalizer with 200 breakpoints.

  • Encoder: MLPEncoder with 2 hidden layers of size 128 and 128 and output of size 64.

  • Coupler: RecurrentCoupler with 200 steps, 2 hidden layers of size 128 and 128 and latent dimension 64.

  • Decoder: MLPEquivariantDecoder with 2 hidden layer of size 128 and 128.

Parameters:

  • in_structure (GraphStructure) – Structure of the input graph.

  • out_structure (GraphStructure) – Structure of the output graph.

  • seed (int) – Seed for RNG streams.

Return type:

Any