Invariant and Equivariant Graph Networks (PyTorch)

A PyTorch implementation of The ICLR 2019 paper “Invariant and Equivariant Graph Networks” by Haggai Maron, Heli Ben-Hamu, Nadav Shamir and Yaron Lipman The official TensorFlow implementation is at


Data should be downloaded from:
Run the following commands in order to unzip the data and put its proper path.

mkdir data
unzip -d data



PyTorch 1.5.0

Additional modules: numpy, pandas, matplotlib

TensorFlow is not neccessary except if you want to run the tests (comparisons) between the PyTorch and TensorFlow versions.

Running the tests

Run the tests comparing between PyTorch and TensorFlow versions’ tensor contractions. All tensor contractions are implemented 1-to-1. The two versions have identical tensor contractions:

cd layers/

Run the (permutation) equivariance tests for the equivariant linear layers implemented in PyTorch (e.g. permute the input tensor and the output tensor must transform covariantly):


Running the experiment

The folder main_scripts contains scripts that run different experiments:

  1. To run 10-fold cross-validation with our hyper parameters run the script. You can choose the datase in 10fold_config.json.
  2. To run hyper-parameter search, run the script with the corresponding config file
  3. To run training and evaluation on one of the data sets run the script

example to run 10-fold cross-validation experiment:

cd main_scripts/
python3 -m main_10fold_experiment --config=../configs/10fold_config.json


PyTorch implementation of tensor contractions and equivariant linear layers is in:


PyTorch implementation of invariant (basic) graph nets:


Related work

Covariant Compositional Networks For Learning Graphs
Predicting molecular properties with covariant compositional networks
The general theory of permutation equivarant neural networks and higher order graph variational encoders


Email: [email protected]


View Github